Riemannian Attention Mechanisms for Transformers: A Theoretical Framework and Architecture Design
cs.LG, cs.AI
Submitted: 2026-08-02
Updated: 2026-08-26
Comments: 23 pages, theoretical paper
License: http://creativecommons.org/licenses/by-nc-sa/4.0/
The gist: All Transformer-based large language models compute attention via the Euclidean inner product, an architectural choice that Dong et al.
Terminology
Abstract
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.
Sources
- Position: Beyond Euclidean -- Foundation Models Should Embrace Non-Euclidean Geometries
- HELM: Hyperbolic Large Language Models via Mixture-of-Curvature Experts
- mHC: Manifold-Constrained Hyper-Connections
- Do Transformer Modifications Transfer Across Implementations and Applications?
- ReZero is All You Need: Fast Convergence at Large Depth
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: Local-Consistency Optimisation for Bayesian Neural Networks
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks