IC-013Entity recognition directions regulate attention to entity tokens in attribute extraction heads in Gemma and Llama models.

Javier Ferrando, Oscar Balcells Obeso, Senthooran Rajamanoharan, Neel Nanda

SourceDo I Know This Entity? Knowledge Awareness and Hallucinations in Language Models

The study shows that attention scores from downstream heads (e.g., l18h5 and l20h3 in Gemma 2 2b) to the last token of an entity are higher for known entities than for unknown ones. Steering the residual stream at the entity's last token with the discovered known entity latent increases attention to the entity, while steering with the unknown entity latent decreases it. This causal relationship is observed in Gemma 2 2b, Gemma 2 9b, and Llama 3.1 8b.

Evidence
interventional
Key metric
Change in attention scores from attribute extraction heads to entity tokens. Statistical tests comparing steering with entity recognition latents versus random SAE latents show significant differences in the expected direction in most cases. For Gemma 2 2b, known latent steering increases attention (10/10 tests), unknown latent steering decreases attention (9/10 tests). For Gemma 2 9b, known latent steering increases attention (10/10 tests), unknown latent steering shows mixed results. For Llama 3.1 8b, known latent steering increases attention (7/10 tests), unknown latent steering decreases attention (10/10 tests).
Caveat
The effect of steering with the unknown entity latent on attention in Gemma 2 9b was less consistent than for other models. Attention score changes are measured on specific attribute extraction heads identified via activation patching.
Model
Gemma 2 2B, 9B, Llama 3.1 8B
Concepts
Linear representation
Datasets
Wikidata / WikidataRecent [source]
Methods
Activation patching / Activation replacement / Cross-model activation patching (CMap) [primary], Activation steering / Mean steering / PCA steering [primary]
Related work
Fact finding: Attempting to reverse-engineer factual recall on the neuron level [builds-on], Attention satisfies: a constraint-satisfaction lens on factual errors of language models [context]
Extraction
automatic-extraction