# Riemannian Attention Mechanisms for Transformers: A Theoretical Framework and Architecture Design

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 18 of its 38 references that are in the map.
- First posted on arXiv 2 Aug 2026.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Parameter-Efficient Transfer Learning for NLP](https://mlsanity.com/paper/1902.00751/) (2019)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Reformer: The Efficient Transformer](https://mlsanity.com/paper/2001.04451/) (2020)
- [Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation](https://mlsanity.com/paper/2108.12409/) (2021)
- [Poincar'e Embeddings for Learning Hierarchical Representations](https://mlsanity.com/paper/1705.08039/) (2017)
- [Attention is Not All You Need: Pure Attention Loses Rank Doubly Exponentially with Depth](https://mlsanity.com/paper/2103.03404/) (2021)
- [ReZero is All You Need: Fast Convergence at Large Depth](https://mlsanity.com/paper/2003.04887/) (2020)
- [Hyperbolic Attention Networks](https://mlsanity.com/paper/1805.09786/) (2018)
- [mHC: Manifold-Constrained Hyper-Connections](https://mlsanity.com/paper/2512.24880/) (2025)
- [Do Transformer Modifications Transfer Across Implementations and Applications?](https://mlsanity.com/paper/2102.11972/) (2021)
- [Reversible Vision Transformers](https://mlsanity.com/paper/2302.04869/) (2023)
- [HELM: Hyperbolic Large Language Models via Mixture-of-Curvature Experts](https://mlsanity.com/paper/2505.24722/) (2025)
- [RiemannFormer: A Framework for Attention in Curved Spaces](https://mlsanity.com/paper/2506.07405/) (2025)
- [Invertible Attention](https://mlsanity.com/paper/2106.09003/) (2021)
- [The Curved Spacetime of Transformer Architectures](https://mlsanity.com/paper/2511.03060/) (2025)
- [ManifoldFormer: Geometric Deep Learning for Neural Dynamics on Riemannian Manifolds](https://mlsanity.com/paper/2511.16828/) (2025)
- [CAT: Curvature-Adaptive Transformers for Geometry-Aware Learning](https://mlsanity.com/paper/2510.01634/) (2025)

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