Decomposing Query-Key Feature Interactions Using Contrastive Covariances

Abstract: Despite the central role of attention heads in Transformers, we lack tools to understand why a model attends to a particular token. To address this, we study the query-key (QK) space -- the bilinear joint embedding space between queries and keys. We present a contrastive covariance method to decompose the QK space into low-rank, human-interpretable components. It is when features in keys and queries align in these low-rank subspaces that high attention scores are produced. We first study our method both analytically and empirically in a simplified setting. We then apply our method to LLMs to identify human-interpretable QK subspaces for categorical semantic features and binding features. Finally, we demonstrate how attention scores can be attributed to our identified features.

• The paper introduces a method using contrastive covariances to decompose query-key interactions into interpretable, low-rank subspaces.
• It validates the approach through synthetic tasks and applications to LLMs, confirming accurate recovery of feature subspaces and effective causal interventions.
• The study demonstrates how targeted manipulation of attention scores in Transformers can be achieved by isolating semantic and binding features.

Decomposing Query-Key Feature Interactions Using Contrastive Covariances

Introduction and Methodology

The paper "Decomposing Query-Key Feature Interactions Using Contrastive Covariances" (2602.04752) addresses interpretability in attention mechanisms of Transformers by introducing a principled method for decomposing the bilinear query-key (QK) interaction space into interpretable, low-rank components. The central innovation is the use of contrastive covariances to isolate specific feature-aligned subspaces in the joint space of queries and keys, facilitating direct analysis of what information drives attention scores between tokens.

Traditional attention analysis is limited by the scalar nature of attention scores, which obscure the multi-dimensional structure of query-key interactions. The proposed method computes contrastive covariances between matched and mismatched feature conditions, thereby isolating subspaces responsible for a specific feature by holding confounding factors constant. Singular Value Decomposition (SVD) applied to these contrastive matrices then identifies both the rank and basis of the underlying feature subspaces in QK space.

Figure 1: Schema illustrating contrastive covariance; positive and negative conditions on QK pairs isolate features in QK space.

The methodology is first validated in analytically controlled settings with known latent variables and is then applied to actual LLMs, including Llama-3.1-8B-Instruct and Qwen 3-4B-Instruct, focusing on filter heads (category selectors) and entity binding heads to expose interpretable structure in the QK space.

Empirical Validation in Synthetic Settings

To precisely characterize the contrastive covariance approach, a synthetic "payload retrieval" task is constructed with discrete and continuous latent variables ($z_1, z_2$ for binary and $s_1, s_2$ for Gaussian). The embedding process composes query and key tokens with independently controlled latent features, which guarantees ground-truth knowledge of feature location, rank, and interaction.

Empirically, the technique consistently recovers the correct rank and subspace for the latent variables, as long as these fit within the attention head's dimensionality ($r_1 + r_2 < d_{head}$). Superposition is observed when the model's representational capacity is saturated.

Figure 2: Heatmap showing recovered ranks for each latent variable, validating the method’s ability to identify true feature rank outside the superposition regime.

SVD projections onto the identified subspaces show that the intrinsic structure of the latent variable (e.g., 3D hypercube for binary, hypersphere for Gaussian) is preserved after recovery.

Figure 3: Projected keys and queries using PCA in the $z_1$ subspace reveal the geometric structure of the latent variable, with precise query-key alignment.

Causal interventions within these subspaces offer further validation. By swapping latent-variable coordinates between key tokens, the authors show that virtually all attention can be shifted from an original to a target token, provided the true latent subspace is used for the intervention. Interventions on random subspaces yield much lower attention transfer, supporting the causal relevance of the identified features.

Figure 4: Attention transfer via causal intervention in QK subspaces, drastically more effective for true feature subspaces compared to random subspaces.

Analysis of the bilinear interaction matrix $G$ ($q^\top k = z_q^\top G z_k$) reveals further structure: for sufficient dimension, features decompose into tensor products with strong block-diagonalization; under capacity constraints, features enter superposition with substantial off-diagonal interaction.

Figure 5: Visualization of QK bilinear interaction matrix $G$, showing independent feature blocks with sufficient dimension and superposition under compression.

Application to LLMs: Discovering Semantic and Binding Features

Categorical and Semantic Structure in Filter Heads

The approach is leveraged to dissect filter heads, which attend to semantically defined categories (e.g., "fruits", "animals", "vehicles", "drinks", "countries"). Using controlled prompts, the contrastive covariance method exposes that categorical QK subspaces are low-rank (rank 1 per category) and capture clear alignment between queries (queries for a specific category) and relevant keys (items of that category).

PCA of the joint QK subspace demonstrates that keys and queries from the same category align distinctly. Expanded category sets yield interpretable clustering along semantic axes.

Causal interventions in these semantic QK subspaces show that attention can be systematically shifted between category tokens, although not perfectly, indicating the presence of additional, possibly higher-order, features.

Figure 6: Causal interventions on category QK subspaces shift attention between semantic groups more effectively than random subspaces.

Entity Binding and Mechanism Analysis

A parallel analysis is conducted for entity binding tasks, motivated by recent studies showing multiple binding mechanisms in LLMs (order-ID, lexical references). The method recovers both "order-ID" (group order) and "lexical" (token identity) features as interpretable, multi-dimensional subspaces (typically rank 2-3 for order-ID, 9-10 for lexical in nine-entity tasks).

Projection (PCA/UMAP) into these subspaces visualizes clusters corresponding to feature identities and confirms query-key alignment.

Figure 7: Embedded projections of order-ID and lexical subspaces displaying distinct cluster structure and key-query alignment.

Causal interventions again indicate that manipulating the relevant subspaces reliably redirects attention, either via order-ID or lexical mechanism, and strongest when both are intervened.

Figure 8: Causal intervention in binding QK subspaces shows that intervening in order-ID or lexical subspaces shifts attention; both together maximize transfer.

Logit Decomposition and Feature Attribution

The paper extends analysis by decomposing attention logits into contributions from each of the identified low-rank feature subspaces. Given the linearity of logits in query space, projectors onto order-ID and lexical subspaces yield not only a summary attribution but also reveal the residual logit component not accounted for by these features.

Visualization of logit attributions demonstrates that most (but not all) attention scores can be explained by the low-rank subspaces, with some unexplained logit mass, attributable to latent features not isolated by the chosen contrastive conditions.

Figure 9: Token-level attribution of logits to order-ID and lexical subspaces, showing relative feature importance and residuals not explained by decomposed features.

Limitations and Implications

A key limitation of the approach is its dependence on human-defined positive/negative feature conditions, i.e., the selection of supervised contrastive axes. The method does not address unsupervised discovery of features, especially complex multi-dimensional or latent features beyond human intuition. Furthermore, in capacity-limited regimes, superposition leads to ambiguous decompositions.

Despite this, the work demonstrates that low-rank, interpretable feature subspaces in QK space govern much of the attention mechanism in both synthetic and real LLM scenarios, offering a procedural path to causal, component-level transparency.

Practically, this enables targeted manipulation of attention, direct attribution of attention scores to cognitively meaningful features, and systematic identification of the representational bottlenecks (when superposition arises). Theoretically, it provides evidence for the structured, modular nature of attention-based computations in large models, furthering a mechanistic interpretability agenda.

Future research directions include unsupervised subspace discovery (e.g., using methods like sparse coding, ICA, or other contrastive learning variants in QK space), multi-dimensional and hierarchical feature analysis, and extension to interactions beyond query-key (e.g., value or output heads).

Conclusion

This work introduces and validates a unified principled approach—contrastive covariance decomposition—for isolating human-interpretable, low-rank feature subspaces in the query-key interaction space of Transformer attention heads. Both analytic theory and extensive experiments on controlled and large-scale settings reveal that attention scores are, to a significant degree, driven by the alignment of features in these learned subspaces. The approach enables interventions that causally redirect attention, quantifies feature attributions in logit space, and provides insights into the superposition phenomenon under dimensional bottlenecks. This methodological advance supports more transparent, modular interpretations of Transformer-based models and opens new avenues for mechanistic understanding and targeted model editing.