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

Standard Transformer architectures rely on the Euclidean inner product for attention, which has been shown to cause representational rank to decay doubly exponentially with depth. This paper investigates whether replacing this flat Euclidean geometry with learned, per-token Riemannian metrics can mitigate this structural limitation. The author proposes a framework where each token position carries its own metric, transforming attention into a geodesic distance computation. To ensure computational feasibility, the author utilizes low-rank metric factors, allowing for efficient distance calculations and metric inversions using the Woodbury identity.

A key theoretical contribution is the proof that Riemannian attention scores are non-Gram, meaning they cannot be factorized as the product of query and key matrices (QK^T) with O(d) dimensionality. This structural change is a direct consequence of metric heterogeneity. Despite this, the author demonstrates that by restricting the Riemannian metrics to a low-rank form (I + UU^T), the geometric operations remain tractable with O(d*r) complexity, making the approach viable for billion-parameter models with negligible overhead.

The paper introduces the Fiber Bundle Transformer, an architectural design that incorporates these geometric principles. In this model, token positions are treated as fibers over a base space, and the architecture includes specific components for metric learning (MetricNet), geodesic attention, and metric-preconditioned feed-forward updates. The author derives formal predictions, suggesting that curvature heterogeneity should emerge during optimization and identifying metric collapse to identity as a primary failure mode that requires architectural countermeasures.