Decomposing Query-Key Feature Interactions Using Contrastive Covariances

Andrew Lee, Yonatan Belinkov, Fernanda Viégas, Martin Wattenberg
Proceedings of the 43rd International Conference on Machine Learning, PMLR 306:63430-63449, 2026.

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 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.

Cite this Paper

BibTeX
@InProceedings{pmlr-v306-lee26b, title = {Decomposing Query-Key Feature Interactions Using Contrastive Covariances}, author = {Lee, Andrew and Belinkov, Yonatan and Vi\'{e}gas, Fernanda and Wattenberg, Martin}, booktitle = {Proceedings of the 43rd International Conference on Machine Learning}, pages = {63430--63449}, year = {2026}, editor = {Zhang, Tong and Dudik, Miroslav and Jaggi, Martin and Agarwal, Alekh and Li, Sharon and Schuurmans, Dale and Zhu, Jerry and Berkenkamp, Felix and Dong, Hanze and Bietti, Alberto}, volume = {306}, series = {Proceedings of Machine Learning Research}, month = {06--11 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v306/main/assets/lee26b/lee26b.pdf}, url = {https://proceedings.mlr.press/v306/lee26b.html}, 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 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.} }
Endnote
%0 Conference Paper %T Decomposing Query-Key Feature Interactions Using Contrastive Covariances %A Andrew Lee %A Yonatan Belinkov %A Fernanda Viégas %A Martin Wattenberg %B Proceedings of the 43rd International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2026 %E Tong Zhang %E Miroslav Dudik %E Martin Jaggi %E Alekh Agarwal %E Sharon Li %E Dale Schuurmans %E Jerry Zhu %E Felix Berkenkamp %E Hanze Dong %E Alberto Bietti %F pmlr-v306-lee26b %I PMLR %P 63430--63449 %U https://proceedings.mlr.press/v306/lee26b.html %V 306 %X 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.
APA
Lee, A., Belinkov, Y., Viégas, F. & Wattenberg, M.. (2026). Decomposing Query-Key Feature Interactions Using Contrastive Covariances. Proceedings of the 43rd International Conference on Machine Learning, in Proceedings of Machine Learning Research 306:63430-63449 Available from https://proceedings.mlr.press/v306/lee26b.html.

Related Material