# Directions That Don't Drift: Stiefel Manifold Routing for Transformer Attention

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 14 of its 23 references that are in the map.
- First posted on arXiv 16 Sep 2026.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Progress measures for grokking via mechanistic interpretability](https://mlsanity.com/paper/2301.05217/) (2023)
- [Are Sixteen Heads Really Better than One?](https://mlsanity.com/paper/1905.10650/) (2019)
- [Scaling Vision Transformers to 22 Billion Parameters](https://mlsanity.com/paper/2302.05442/) (2023)
- [Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets](https://mlsanity.com/paper/2201.02177/) (2022)
- [Unitary Evolution Recurrent Neural Networks](https://mlsanity.com/paper/1511.06464/) (2015)
- [Query-Key Normalization for Transformers](https://mlsanity.com/paper/2010.04245/) (2020)
- [Riemannian Adaptive Optimization Methods](https://mlsanity.com/paper/1810.00760/) (2018)
- [Stabilizing Transformer Training by Preventing Attention Entropy Collapse](https://mlsanity.com/paper/2303.06296/) (2023)
- [Omnigrok: Grokking Beyond Algorithmic Data](https://mlsanity.com/paper/2210.01117/) (2022)
- [Full-Capacity Unitary Recurrent Neural Networks](https://mlsanity.com/paper/1611.00035/) (2016)
- [The Slingshot Mechanism: An Empirical Study of Adaptive Optimizers and the Grokking Phenomenon](https://mlsanity.com/paper/2206.04817/) (2022)
- [Orthogonal Weight Normalization: Solution to Optimization over Multiple Dependent Stiefel Manifolds in Deep Neural Networks](https://mlsanity.com/paper/1709.06079/) (2017)
- [Efficient Riemannian Optimization on the Stiefel Manifold via the Cayley Transform](https://mlsanity.com/paper/2002.01113/) (2020)

## Abstract (arXiv)

The query and key projections WQ,WK in attention are almost always trained by Euclidean optimizers with no geometric constraint. We constrain them to the Stiefel manifold and optimize with a Riemannian Adam carrying one scalar second moment per frame---the form of citetbecigneul2019, here extended to the compact, non-Hadamard St(d,r) with a tangent projector, step-norm cap, and polar retraction. Four propositions prove steepest descent in the embedded metric, gradient-scale independence, well-conditioning, and exact O(d)-equivariance. A fifth records that weight decay has identically zero Riemannian gradient on St(d,r) (W=WIr lies in the normal space), so decay cannot act on the constrained frames. On a CIFAR-10 patch benchmark at n=10k this rule gains +6.79,pp over AdamW across 12 paired starts (t=38.33, 12/12); earlier fixed-step Riemannian SGD gains +1.97,pp, of which +1.69,pp comes from frozen orthonormal initialization alone. The corrected Adam's lead grows with data: +1.9,pp at n=1k to +6.7,pp at n=50k. A 12-seed ablation credits all gain to the scale-free step (+4.63,pp, 12/12), nothing to the projector or equivariance; a targeted varepsilon-sweep causally confirms the mechanism (-2.6,pp at varepsilon=0.1, p<0.001). Two five-seed grokking studies confirm the constrained arm does not grok better than the baseline (p=0.019, A2 wins): the weight-decay exemption has no grokking consequence. A single-seed pilot exploiting this localization achieves the first stable grokking under slingshot conditions---Stiefel + targeted circuit regularization keeps routing-frame isometry error 106× lower than the unconstrained ablation through every collapse.