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

2026-08-02

AI Summary

This work addresses the double-exponential decay of representation rank with depth in Transformers, caused by Euclidean inner-product attention. To mitigate this issue, the authors propose replacing the Euclidean metric with token-wise learnable Riemannian metrics, thereby constructing a geodesic-distance-based attention mechanism and introducing a complete Fiber Bundle Transformer architecture. This is the first approach to incorporate heterogeneous Riemannian metrics into attention, revealing their non-Gram structure. Efficient computation is achieved through low-rank decomposition combined with the Woodbury identity, while curvature and torsion are explicitly modeled within the architecture. The resulting framework scales to billion-parameter models with controlled computational overhead, offering a novel pathway to alleviate representation rank collapse.

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.

Problem

• rank collapse
• self-attention
• Riemannian geometry
• Transformer architecture
• representational degeneracy

Innovation

• Riemannian attention
• rank collapse
• fiber bundle
• geodesic distance
• low-rank metric