cs.LGOct 21, 2026
SaveGeometric knowledge spaces, Attention pairing and Architectures
Abstract
In this paper, we propose a geometric approach to deriving attention mechanisms and exploring potential new transformer architectures. We begin by formalizing the notions of knowledge worlds and attention pairings, then demonstrate how this framework naturally leads to the architectures of standard attention, causal attention, linear attention, tensor attention, and recurrent neural networks. Finally, we introduce several new geometric spaces that may inspire novel neural network architectures.
Figures & tables
| Architecture | knowledge worlds | Intelligence | Update and readout |
|---|---|---|---|
| Attention (ordinary or causal) | , , , . | , , . | . , . |
| Linear attention | , , , . | , , , . | . . Normalized evaluation: . |
| Tensor attention | , , , . | , . | . , . Concatenate heads and apply . |
| Bilinear RNN | , , , . | , . The evaluation is an increment to memory. | , , . |
| Vanilla RNN | , , , . | , . The evaluation is the next hidden state. | , . The hidden state is replaced by its evaluation. |
Table 1. knowledge worlds, intelligence pairings, and the resulting neural architectures. The RNN formulas describe the stabilized stages , with initial stages as defined in the preceding constructions.
More in this topic
Attention Mechanisms
Attention mechanisms let neural networks weight information from other tokens or regions according to relevance, enabling context-dependent processing.
See all 517 papersAttention normalizes learned scores into a Markov operator. In the continuum, the local sector becomes diffusion on a learned Riemannian manifold, but attention natively bonds distant points, creating nonlocal structure such that no tensor geometry can describe. The Bakry-Emery calculus, built directly from the transition operator, extends geometry to these potentially stratified spaces, giving a coordinate-free account of what vision and language transformers learn.
Query Expansion and Key Specialization in Transformer Attention Geometry
The projection of queries and keys are central to the attention mechanism in Transformer architectures. While they are mathematically symmetric, they play different roles in attention mechanisms. The question of whether there is an effect from their functional distinction on their geometric development in training remains unanswered. We investigate the problem through the training of small GPT-like Transformers on character-level WikiText-103 for three different depths (4, 6, and 8 layers), three types of initialization for queries and keys, and four random seeds, resulting in 36 runs and 54 trajectories of average layers across seeds. We track the effective dimensionality of those layers using participation ratios and discover that effective dimension of queries expand while keys shrink, and that is positive in all trajectories studied. In connection to attention, the shrinking of keys leads to a narrower spectrum of and more peaked attention weights. In order to determine if this connection is causal or coincidental, we directly control the spectrum of keys during training across five seeds: restricting it to make it shrink sharpens the attention with high directional confidence, while keeping it constant to the level of initial dispersion makes attention softer. Additional token-level checkpoint analyses show that the monotonic paired-contrast trend is not universal across pretrained families, but survives as an early-training regime that later decays over a full pretraining run, and the link between interaction-rank geometry and attention entropy remains visible in several models.
Riemannian Attention Mechanisms for Transformers: A Theoretical Framework and Architecture Design
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(dr) per token and metric inversion in O(dr^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.