IC-484Re-scaling the first 2-3 principal components of Llama-3.1-8B token representations causally shifts next-token predictions toward the target graph position

Core Francisco Park, Andrew Lee, Ekdeep Singh Lubana, Yongyi Yang, Maya Okawa, Kento Nishi, Martin Wattenberg, Hidenori Tanaka

SourceICLR: In-Context Learning of Representations

The authors intervene on Llama-3.1-8B's activations by re-scaling the first 2 or 3 principal components of a token's representation to match another token's mean representation, effectively 'moving' the token to a different position in the graph. Across 1,000 randomly generated contexts and all n-1 target positions, this intervention significantly improves hit@1, hit@3, and accumulated probability mass on valid neighbors compared to null or random-vector interventions, across ring, grid, and hexagonal graphs. The effect is strongest with 3 PCs and weakest for the hexagonal graph.

Evidence
interventional
Key metric
Ring hit@1: 0.61 (2 PCs), 0.77 (3 PCs), 0.20 (null), 0.17 (random). Grid hit@1: 0.57 (2 PCs), 0.68 (3 PCs), 0.17 (null), 0.16 (random). Hex hit@1: 0.30 (2 PCs), 0.42 (3 PCs), 0.07 (null), 0.06 (random).
Caveat
The authors state that the principal components 'have some causal effect on the model's output predictions, but does not provide a full explanation' and that 'our simple re-scaling approach does not perfectly capture a causal relationship between principal components and model predictions.'
Model
Llama 3.1 8B
Concepts
Linear representation
Methods
Principal component analysis [primary]
Related findings
IC-481, IC-482, IC-483
Extraction
automatic-extraction