Singular Vectors of Attention Heads Align with Features

AI summary

Overview

Research area: Mechanistic interpretability of transformer language models, specifically the geometry of attention head weights and how features are represented inside models.

Technical level: Advanced. The paper uses singular value decomposition, matrix spectra, Gram matrices, and transformer internals, and it states and proves theorems about trained attention weights.

Scope: The paper investigates why and when the singular vectors of an attention head's QK matrix align with the features that the head uses, using a controllable toy model, formal theorems, and checks on GPT-2 and Pythia.

What This Paper Is About

Several prior studies assumed or observed that the singular vectors of an attention head tend to line up with the features that head uses, but no work had explained why this happens or under what conditions it should be expected. This paper asks the motivating question directly: when features are represented linearly, under what conditions will the singular vectors of attention head matrices align with those features? The authors also ask how such alignment can be recognized operationally in real models where feature representations cannot be directly observed.

Key Contributions

1. Demonstrates singular vector–feature (SVF) alignment robustly in a setting where features can be directly observed. Using an extension of the toy autoencoder model from Elhage et al. (2022) with an added attention head, the authors show that when a head attends to specific feature pairs, the head's singular vectors align with those features, and they confirm this over variations in the relative loss weight λ, the number of features N, context length m, head dimension H, and random seeds.

2. Provides theory explaining when alignment must occur. Theorem 1 shows that when the head's target logits are 1 if feature x₁ is present in the query token and y₁ in the key token (0 otherwise), the trained Ω⋆ is rank-1 with left and right singular vectors u₁ ∝ Σ_X⁻¹x₁ and v₁ ∝ Σ_Y⁻¹y₁. Corollary 1 shows that with isotropic features the singular vectors align exactly with the features. Theorem 2 shows approximate alignment when features deviate from isotropy, given sufficiently bounded interference between features, and Theorem 3 shows that with a penalty for interference the training objective is minimized when features are orthogonal.

3. Identifies sparse attention decomposition (SAD) as a testable prediction of SVF alignment and confirms it in both the toy model and real models, including evidence that it is not merely a consequence of attention matrices being low-rank or ill-conditioned.

4. Validates the framework in real models across two architectures. Evidence is reported in Pythia-160M and GPT-2, which the authors note use different positional schemes (Pythia uses rotary positional encoding, GPT-2 does not), mitigating the risk of architectural dependency.

Main Findings

• SVF alignment arises when a head attends to specific feature pairs: In the toy model with N=20 features in D=10 dimensions and head dimension H=10, adding the head causes feature w₀ to align with left singular vector u₀ and feature w₁ to align with right singular vector v₀, with the spectrum showing only a single large singular value. The simplified isolated autoencoder without attention instead arranges features into 10 antipodal pairs, one per dimension, in an isotropic arrangement.

• The top singular vectors align with the most important feature pair for computing attention: With 100 features in hidden dimension 50 and target logits set to T_{i,i+20} = 26 − i for 0 ≤ i < 20 (zero otherwise), 20 feature pairs are attended with logit values declining linearly in i, and the singular values of the head reflect the values of those target logits.

• Features orthogonalize to minimize interference, enabling multi-pair alignment: Features of interest become orthogonal to the other features — shown for 20 features in which w₀ and w₁ are of interest, and for 100 features in dimension 50 with 40 features of interest (20 pairs). The authors hypothesize that singular vectors align with features to minimize attention loss while features orthogonalize to minimize reconstruction loss.

• Features and singular vectors evolve together during training, in order of importance: In a run with 20 features, 10 hidden dimensions, and four pairs (w_i, w_{i+4}) for i = 0, 1, 2, 3, over 10,000 training steps, feature w₀ changes little while u₀ moves into its final position around step 1500; u₁ shifts until about step 4000, after which feature w₁ abruptly shifts toward a position orthogonal to w₀. Alignment occurs for the highest-logit features first.

• Anisotropy does not necessarily destroy alignment: Using SAE dictionary elements as a proxy for features in GPT-2, the authors report that ‖E_X‖₂ ranges between 10 and 55 depending on layer. In a toy experiment with frozen features and controlled anisotropy over that observed range, SVF alignment remains good with mean cosine similarity above 0.75 even when anisotropy is close to the maximum seen in GPT-2. Theorem 2 is described as a lower bound that can be quite loose in practice, and anti-whitening singular vectors using Σ_X is reported to significantly improve SVF alignment.

• SAD emerges during training in the toy model: Decomposition of relative attention becomes sparse later in training, and sparsity is sparser when fewer features of interest are present. Measured with the sparsity metric S(v) (from Rolls and Tovee, 1995), which is 1 when inputs are minimally sparse and 1/n when maximally sparse, decomposition is sparse when features of interest are present and much less sparse when they are absent.

• SAD emerges in Pythia-160M: Using a prompt from the Indirect Object Identification (IOI) task and capturing relative attention at the heads and token pairs identified in Tigges et al. (2024), relative attention decomposition becomes significantly sparse over training, with most of the contribution to relative attention coming from a small set of singular vectors. Not all attention heads show such strong sparsity emergence, which the authors suggest may be due to a larger number of features playing a role in some heads.

• Sparsity is not explained by a few large singular values: Figure 8 shows the largest contributors to relative attention can come from among the smallest singular values, and recomputing S(v) after randomly rotating the U and V singular vector matrices — which changes no spectrum or rank — eliminates the decline in S(v). Additional comparisons to spectrum-preserving rotations of the SVD bases are reported.

• Few singular vectors are needed to reconstruct relative attention in real models: The N_recon(j) metric counts the minimum number of singular vectors needed to reconstruct relative attention on key token s_j. In Pythia, typically only 1–4 singular vectors are needed, and this value declines during training. For GPT-2 on the IOI task, averaged over 128 IOI prompts varying names, prompt templates (including name order), and objects, N_recon is small in most cases, with typical values varying across attention heads and reflecting different functional roles.

• Alignment extends to positional encoding and to the overcomplete regime: Appendix D extends the toy model to include rotary positional encoding (RoPE) and shows SVF alignment can still occur, both for features with position-independent target logits and for features with position-dependent logits. Appendix E probes the regime where there are more features of interest than head dimensions, noting that some features move into superposition within the head's representation space, but SVF alignment still holds for the majority of singular vectors.

Methodology in Plain English

The authors build a small, fully controllable model in which both the features and the attention weights can be inspected. It extends the toy autoencoder of Elhage et al. (2022), which represents inputs as sums of feature vectors, by adding a single attention head that scores a query token r against key tokens using ℓ = rᵀW_QᵀW_K s, with Ω = W_QᵀW_K. Because the model is trained end to end, the feature matrix W and the head weights Ω are learned together, so the researchers can compute cosine similarities between features and singular vectors directly rather than guessing.

Attention behavior is specified through target logits parameterized by a matrix T, where T_ij states how much should be added to the target logit to the extent features w_i and w_j are present in the two tokens. The head is trained with cross-entropy against the softmax of these target logits, and the total loss combines reconstruction loss with attention loss weighted by λ. This design lets the authors dial in exactly which feature pairs a head should attend to.

To understand training dynamics, the authors track how features and singular vectors move over training. To test whether alignment happens in real models where features cannot be observed, they derive the implication called sparse attention decomposition: if features align with singular vectors (the "alignment hypothesis") and unimportant features are roughly orthogonal to important ones, then a head's logit decomposed in the SVD basis — and more appropriately relative attention, ℓ_j minus the mean of the other logits — should show a sparse representation, with large contributions only from singular vectors matching features actually present. They measure this with the S(v) sparsity metric and the N_recon(j) reconstruction count, and rule out low-rank explanations by randomly rotating the singular vector bases while preserving the spectrum.

Why This Matters

Impact on research. Methods for finding features in models all have drawbacks: probing is primarily correlational, and sparse autoencoders are expensive to train and have documented issues, possibly because they are built only from activations and ignore model weights. Exploiting SVF alignment uses both activations and weights, and reduces feature identification to decomposing activations in the SVD basis — something that can be done per prompt in a single forward pass. By supplying theoretical justification and conditions for alignment, this paper puts that strategy on firmer footing and makes singular vectors usable as a comprehensible, tractable search space of candidate features.

Potential applications (these follow from the method's properties; the paper does not itself report deployed applications):

• Auditing model behavior and supporting safety-oriented interpretability work, a use the authors cite as motivation.
• Localizing model components for debugging, by narrowing attention to a small set of directions worth investigating.
• Causal analysis of specific heads, since candidate features found this way can be tested for causal impact on task performance.
• Per-prompt, single-forward-pass feature decomposition that could be integrated into interpretability tooling more cheaply than training SAEs.

Industry relevance. The approach is described as more straightforward, efficient, and scalable than SAEs or linear probes because it requires no auxiliary training and can run per prompt. That makes it attractive to teams that need interpretability checks on deployed language models without large compute budgets, and the validation across both Pythia (with RoPE) and GPT-2 (without) suggests the finding is not tied to a single architecture.

Future Directions

• The overcomplete regime. The authors explicitly frame the case where a head has more features of interest than dimensions as a regime needing follow-on study; in their Appendix E probe, some features fall into superposition, and alignment holds only for the majority of singular vectors rather than all of them.

• Explaining heads where sparsity does not strongly emerge. Not all attention heads in Pythia show strong sparsity emergence, which the authors suggest may be because more features play a role in some heads — an open question about how the framework behaves as the effective number of features grows.

• Tighter theory under anisotropy. The authors note that Theorem 2 is a lower bound that can be quite loose in practice, leaving room for sharper bounds, and they show empirically that anti-whitening with Σ_X improves alignment.

• Extending positional and multi-head settings. RoPE is treated in an appendix as an extension of the toy model, and the analysis focuses on a single head's QK matrix, leaving the interaction of multiple heads and further positional schemes as natural next steps.

Target Audience

This paper is most valuable to mechanistic interpretability researchers and graduate students working on feature identification in language models, particularly those using or evaluating SVD-based methods and sparse autoencoders. It also suits machine learning theorists interested in provable statements about what training objectives induce in attention weights, and safety-focused practitioners who need a cheap, weight-aware method for locating candidate features. Readers should be comfortable with linear algebra, singular value decomposition, and transformer internals; the full paper content provided does not report any experiments beyond the toy model, GPT-2, and Pythia.

Authors’ abstract

Identifying feature representations in language models is a central task in mechanistic interpretability. Several recent studies have made the observation that feature representations can be inferred in some cases from singular vectors of attention matrices. However, sound justification for this phenomenon is lacking. In this paper we address that question, asking: why and when do singular vectors align with features? First, we demonstrate that singular vectors robustly align with features in a model where features can be directly observed. We then show theoretically that such alignment is expected under a range of conditions. We close by asking how, operationally, alignment may be recognized in real models where feature representations are not directly observable. We identify sparse attention decomposition as a testable prediction of alignment, and show evidence that it emerges in real models in a manner consistent with predictions. Together these results suggest that alignment of singular vectors with features can be a sound and theoretically justified basis for feature identification in language models.