Invertible Query-Key Coupling Composes with Attention Mechanisms
Not on the map. Why?
Abstract (arXiv)
Scaled dot-product attention forms its queries and keys as independent linear projections, so the two never interact before the dot product that scores them. We study coupled query-key dynamics, a pre-scoring transformation that evolves each token's query and key jointly through a shared invertible coupling before standard scoring. We realize it as an alternating affine map in the style of real non-volume-preserving flows: the coupling is the identity at initialization, adds a small fraction of parameters per head, and leaves the softmax and surrounding architecture unchanged. We place coupling on top of existing attention methods rather than replacing them, and ask whether that composition helps. On WikiText-103, adding coupling to Differential Attention improves on it at both 150M and 455M parameters. At 455M the gain is significant at sequence length 512 (p=0.003, six seeds), survives a Bonferroni correction and replicates on a held-out test split; it also holds across rotary-embedding training lengths 512, 1024, and 2048. The same additive direction appears when coupling is added to query-key normalization (significant at 150M) and Multi-Token Attention (directional). Matched controls attribute the gain to the joint pre-scoring coupling rather than to added capacity, and show that removing the invertibility guarantee preserves the 455M gain yet is far worse than the base method at 150M, so invertibility is what makes the coupling reliable across scales. On its own, coupling lowers perplexity at 60M and 150M (one-sided Welch tests, p<0.05) but the gain narrows with scale and is not significant at 455M. We relate the construction to the expressivity of coupling flows, use an associative-recall study to map where coupling helps and where it degrades sharp retrieval, and conclude that coupling is most useful in composition with a scoring-stage method rather than as a standalone change.
Key facts. No citing papers in the map yet (data through 7 Oct 2026). Builds on 23 of its 27 references that are in the map. First posted on arXiv 2 Apr 2026.
Builds on
23 of 27 references.
References · 23
- 2017Attention Is All You Needcited by 59,932
- 2021RoFormer: Enhanced Transformer with Rotary Position Embeddingcited by 4,581
- 2014Neural Machine Translation by Jointly Learning to Align and Translatecited by 4,530
- 2016Pointer Sentinel Mixture Modelscited by 2,992
- 2016Density estimation using Real NVPcited by 2,029
- 2023GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpointscited by 1,311
- 2015Effective Approaches to Attention-based Neural Machine Translationcited by 1,024
- 2019Hamiltonian Neural Networkscited by 702
- 2020Hopfield Networks is All You Needcited by 461
- 2021Attention is Not All You Need: Pure Attention Loses Rank Doubly Exponentially with Depthcited by 411
- 2020Query-Key Normalization for Transformerscited by 266
- 2025Gated Attention for Large Language Models: Non-linearity, Sparsity, and Attention-Sink-Freecited by 266
11 more
- 2020Synthesizer: Rethinking Self-Attention in Transformer Modelscited by 204
- 2023Zoology: Measuring and Improving Recall in Efficient Language Modelscited by 188
- 2023A mathematical perspective on Transformerscited by 160
- 2023Stabilizing Transformer Training by Preventing Attention Entropy Collapsecited by 154
- 2020Coupling-based Invertible Neural Networks Are Universal Diffeomorphism Approximatorscited by 85
- 2024nGPT: Normalized Transformer with Representation Learning on the Hyperspherecited by 68
- 2020Talking-Heads Attentioncited by 57
- 2025Multi-Token Attentioncited by 10
- 2024Elliptical Attentioncited by 7
- 2024Expanding Expressivity in Transformer Models with M"obiusAttentioncited by 4
- 2026Beyond Linearity in Attention Projections: The Case for Nonlinear Queriescited by 2
Built on by
Not cited yet.