Decomposing Query-Key Feature Interactions Using Contrastive Covariances

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 large language models 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.

Language models work by attending to different parts of the input context to predict a next token. Despite the central role of attention, we do not know how to explain why the model attends to the tokens it chooses to attend to. To study this, we study the "QK-Space" of attention modules, which captures how token representations at different positions interact with one another to determine whether to attend to a token or not. We propose a new method that relies on counterfactual pairs, one in which a high attention score is produced and another in which low attention score is produced, which can then be used to identify interpretable features in the attention module's QK-space.