RiemannFormer: A Framework for Attention in Curved Spaces

“Spacetime tells matter how to move; matter tells spacetime how to curve.”

— John Archibald Wheeler

Official implementation of the paper "RiemannFormer: A Framework for Attention in Curved Spaces".

We introduce RiemannFormer, a paradigm shift that reimagines the Transformer through the lens of differential geometry. By rejecting the implicit flat Euclidean assumption of standard Self-Attention, we posit that token embeddings reside on a Curved Riemannian Manifold.

Our framework introduces PaTESO (Parallel Transport of Embeddings via Subspace Orientation), a geometrically principled position encoding derived from the isomorphism $\mathfrak{so}(4) \cong \mathfrak{su}(2)_L \oplus \mathfrak{su}(2)_R$. This allows us to decouple Flat Translation from Intrinsic Curvature, enabling a dynamic geometry where "Matter determines the Metric".

🚀 News & Updates

• [2025-06-10] 🏆 SOTA Efficiency: Our metric-gated architecture outperforms standard ViT and RoPE variants on CIFAR benchmarks with high parameter efficiency.

🏗️ Methodology

0. New Paradigm

We constructed the foundational parallel transport operator $\mathcal{P}_{n \to m} = \mathbf{T}_m \mathbf{T}_n^{-1}$. While this formulation elegantly aligns tangent spaces, it implicitly assumes a trivial (flat) connection on the manifold; that is, parallel transport is path-independent and lacks Holonomy effects.

However, adapting the standard Transformer for higher-dimensional data — such as images or voxels — requires the ability to model curvature and capture complex topological structures. To address this and incorporate richer semantic topology, we introduce an interaction operator $\mathcal{F}_{nm}$ within the common reference space $\mathbb{R}^D$. This generalizes our transport protocol into a "sandwich" structure:

$$ \mathcal{P}_{n \to m} = \mathbf{T}_m \cdot \mathcal{F}_{nm} \cdot \mathbf{T}_n^{-1} $$

In this formulation, the intermediate term $\mathcal{F}{nm}$ encodes the manifold's intrinsic geometry. To preserve the norm of feature vectors and focus on directional geometric relationships, we specifically instantiate $\mathcal{F}{nm}$ as a rotation operation within the structure group.

As illustrated by the blue arc in the figure, this transformation is generated by an element of the Lie algebra, $\mathbf{\Omega}_{nm}$:

$$ \mathcal{F}_{nm} = \exp(\mathbf{\Omega}_{nm}) $$

Consequently, $\mathbf{\Omega}_{nm}$ serves as the crucial geometric quantity that captures the manifold's intrinsic curvature and holonomy effects. This formulation yields our enhanced parallel transport operator:

$$ \mathcal{P}_{n \to m} = \mathbf{T}_m \cdot \exp(\mathbf{\Omega}_{nm}) \cdot \mathbf{T}_n^{-1} $$

This approach not only breaks the assumption of spatial flatness but also enables interactions between tokens to explicitly perceive the intrinsic geometric structure of the manifold.

1. The Attention Manifold

We model the interaction between tokens $n$ and $m$ as Parallel Transport $\mathcal{P}_{n \to m}$ along a geodesic. The transport operator is factorized into a flat component and a curvature component:

$$ \mathcal{P}_{n \to m} = \underbrace{\exp\left( \Delta \mathbf{p}_{nm} \cdot \mathbf{J}_L \right)}_{\text{Flat Translation}} \cdot \underbrace{\exp\left( \lambda \cdot \Omega_{nm} \cdot \mathbf{J}_R \right)}_{\text{Symplectic Curvature}} $$

2. Matter Determines Geometry (Metric Gating)

Unlike standard attention where the metric is fixed ($\mathbf{I}$), we learn a Dynamic Riemannian Metric $\mathbf{g}_{ij} = s_i s_j$. The scaling factor $s_i$ is derived directly from the semantic content (Matter):

$$ s_i = \sigma(\text{Proj}(\mathbf{x}_i)) $$

This allows the model to spatially expand salient regions (foreground) and compress noise (background), effectively acting as a "Gravitational Lens" for information flow.

💻 Usage

📂 Project Structure

🖊️ Citation

If you find this work helpful or inspiring, please cite us:

@article{2025riemannformer,
  title={RiemannFormer: A Framework for Attention in Curved Spaces},
  author={Zhongping Ji},
  journal={arXiv preprint arXiv:2506.07405},
  year={2025}
}

🙏 Acknowledgement

We stand on the shoulders of giants. This work draws inspiration from Riemannian geometry，General Relativity, Gauge Theory.