Riemannian Attention Mechanisms for Transformers: A Theoretical Framework and Architecture Design
Not on the map. Why?
Abstract (arXiv)
All Transformer-based large language models compute attention via the Euclidean inner product, an architectural choice that Dong et al. (2021) proved causes representational rank to decay doubly exponentially with depth in pure self-attention stacks. We develop a theoretical framework that targets this structural limitation at the mathematical level by replacing the flat Euclidean metric with learned per-token Riemannian metrics. Our contributions are threefold. (1) We prove that Riemannian attention scores with heterogeneous per-token metrics are non-Gram---they cannot be factorized as QK^T with factorization dimension O(d). We are explicit that this is a structural observation, not a proof of rank preservation. (2) We establish that low-rank metric factors render all geometric operations tractable: geodesic distance in O(d*r) per token and metric inversion in O(d*r^2) via the Woodbury identity---both far below the O(d^3) cost of a general matrix---making Riemannian attention feasible at billion-parameter scale with negligible overhead. (3) We present the Fiber Bundle Transformer, a complete architecture specification in which each token position carries its own Riemannian metric, attention is geodesic distance computation, feed-forward updates use metric-preconditioned steps, and the connection carries explicit curvature and torsion proxies. We derive formal predictions about correctly implemented geometric architectures and identify the central open problem: proving or disproving that heterogeneous Riemannian metrics prevent the rank collapse that row-stochastic attention matrices otherwise cause. This paper presents theoretical analysis and architectural design; empirical validation is the subject of future work.
Key facts. No citing papers in the map yet (data through 7 Oct 2026). Builds on 18 of its 38 references that are in the map. First posted on arXiv 2 Aug 2026.
Builds on
18 of 38 references.
References · 18
- 2017Attention Is All You Needcited by 59,932
- 2019Parameter-Efficient Transfer Learning for NLPcited by 4,195
- 2022FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awarenesscited by 3,611
- 2020Reformer: The Efficient Transformercited by 1,658
- 2021Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolationcited by 1,085
- 2017Poincar'e Embeddings for Learning Hierarchical Representationscited by 674
- 2021Attention is Not All You Need: Pure Attention Loses Rank Doubly Exponentially with Depthcited by 411
- 2020ReZero is All You Need: Fast Convergence at Large Depthcited by 260
- 2018Hyperbolic Attention Networkscited by 122
- 2021Do Transformer Modifications Transfer Across Implementations and Applications?cited by 87
- 2025mHC: Manifold-Constrained Hyper-Connectionscited by 76
- 2023Reversible Vision Transformerscited by 33
6 more
- 2025HELM: Hyperbolic Large Language Models via Mixture-of-Curvature Expertscited by 22
- 2025RiemannFormer: A Framework for Attention in Curved Spacescited by 8
- 2021Invertible Attentioncited by 5
- 2025The Curved Spacetime of Transformer Architecturescited by 5
- 2025ManifoldFormer: Geometric Deep Learning for Neural Dynamics on Riemannian Manifoldscited by 4
- 2025CAT: Curvature-Adaptive Transformers for Geometry-Aware Learningcited by 1
Built on by
Not cited yet.