# Invertible Query-Key Coupling Composes with Attention Mechanisms

arXiv: https://arxiv.org/abs/2604.01683 · HTML: https://mlsanity.com/paper/2604.01683/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 23 of its 27 references that are in the map.
- First posted on arXiv 2 Apr 2026.

## Builds on (most cited first, 23 of 23)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [Neural Machine Translation by Jointly Learning to Align and Translate](https://mlsanity.com/paper/1409.0473/) (2014)
- [Pointer Sentinel Mixture Models](https://mlsanity.com/paper/1609.07843/) (2016)
- [Density estimation using Real NVP](https://mlsanity.com/paper/1605.08803/) (2016)
- [GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints](https://mlsanity.com/paper/2305.13245/) (2023)
- [Effective Approaches to Attention-based Neural Machine Translation](https://mlsanity.com/paper/1508.04025/) (2015)
- [Hamiltonian Neural Networks](https://mlsanity.com/paper/1906.01563/) (2019)
- [Hopfield Networks is All You Need](https://mlsanity.com/paper/2008.02217/) (2020)
- [Attention is Not All You Need: Pure Attention Loses Rank Doubly Exponentially with Depth](https://mlsanity.com/paper/2103.03404/) (2021)
- [Gated Attention for Large Language Models: Non-linearity, Sparsity, and Attention-Sink-Free](https://mlsanity.com/paper/2505.06708/) (2025)
- [Query-Key Normalization for Transformers](https://mlsanity.com/paper/2010.04245/) (2020)
- [Synthesizer: Rethinking Self-Attention in Transformer Models](https://mlsanity.com/paper/2005.00743/) (2020)
- [Zoology: Measuring and Improving Recall in Efficient Language Models](https://mlsanity.com/paper/2312.04927/) (2023)
- [A mathematical perspective on Transformers](https://mlsanity.com/paper/2312.10794/) (2023)
- [Stabilizing Transformer Training by Preventing Attention Entropy Collapse](https://mlsanity.com/paper/2303.06296/) (2023)
- [Coupling-based Invertible Neural Networks Are Universal Diffeomorphism Approximators](https://mlsanity.com/paper/2006.11469/) (2020)
- [nGPT: Normalized Transformer with Representation Learning on the Hypersphere](https://mlsanity.com/paper/2410.01131/) (2024)
- [Talking-Heads Attention](https://mlsanity.com/paper/2003.02436/) (2020)
- [Multi-Token Attention](https://mlsanity.com/paper/2504.00927/) (2025)
- [Elliptical Attention](https://mlsanity.com/paper/2406.13770/) (2024)
- [Expanding Expressivity in Transformer Models with M"obiusAttention](https://mlsanity.com/paper/2409.12175/) (2024)
- [Beyond Linearity in Attention Projections: The Case for Nonlinear Queries](https://mlsanity.com/paper/2603.13381/) (2026)

## Built on by (most cited first, 1 of 1)

- [Phase Space Attention:A Hairer Lift Circumvents the Single-Layer Induction Obstruction](https://mlsanity.com/paper/2609.32319/) (2026)

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