On the Geometry of Positional Encodings in Transformers

1. Problem Statement

Transformer models rely on self-attention mechanisms, which are inherently permutation-equivariant (i.e., they treat input sequences as sets rather than ordered sequences). To handle sequential data, positional encodings (PE) are added to token embeddings. However, the design of these encodings has historically been driven by empirical trial-and-error (e.g., sinusoidal, learned, RoPE, ALiBi) rather than a rigorous mathematical theory.

The paper addresses three fundamental gaps:

1. Necessity: Is positional information strictly required, or can the model learn order implicitly?
2. Learning Dynamics: What geometric structure does a learned positional encoding acquire at convergence?
3. Optimality: What constitutes an "optimal" positional encoding, and how can it be constructed theoretically without training?

2. Methodology and Theoretical Framework

The paper develops a mathematical theory based on Information Geometry and Multidimensional Scaling (MDS).

A. Statistical Geometry of Positions

The authors define the "statistical distance" between two positions $i$ and $j$ in a corpus using the Hellinger distance ($d_H$) between their marginal token distributions ($\mu_i$ and $\mu_j$).
$$d_H(\mu_i, \mu_j) = \left( \sum_{v \in V} (\sqrt{\mu_i(v)} - \sqrt{\mu_j(v)})^2 \right)^{1/2}$$
This metric is chosen because it is a true metric (unlike KL divergence), bounded, and intrinsic to the probability simplex.

B. The Stress Criterion

To evaluate how well a positional encoding $P$ (where $p_i$ is the vector for position $i$) represents the corpus, the paper introduces the Stress metric:
$$\text{stress}(P) = \frac{\sum_{i<j} (\|p_i - p_j\| - d_H(\mu_i, \mu_j))^2}{\sum_{i<j} d_H(\mu_i, \mu_j)^2}$$

• Zero stress implies a perfect isometric embedding of the corpus's positional geometry into the vector space.
• High stress indicates a geometric mismatch.
• Note: The paper clarifies that low stress measures geometric faithfulness, not necessarily downstream task accuracy, though the two are hypothesized to be related.

C. Theoretical Tools

• Neural Tangent Kernel (NTK): Used to analyze the gradient flow of positional encodings during training, specifically to prove the "Monotonicity Conjecture."
• Classical MDS: Used to construct the information-optimal encoding by finding the best low-dimensional flat approximation of the curved statistical manifold defined by the Hellinger distances.

3. Key Contributions and Theorems

The paper presents four main theoretical results:

1. The Necessity of Positional Information (Theorem 1)

• Claim: A Transformer without any explicit positional signal computes a function that is permutation-equivariant.
• Implication: Such a model cannot distinguish between a sequence and its permutations. Therefore, it cannot solve any task sensitive to word order. Positional signals are strictly necessary for order-sensitive tasks.

2. The Positional Separation Theorem (Theorem 4)

• Claim: Under mild conditions (non-stationarity of token distributions, order sensitivity of the task, and non-degeneracy of the attention matrix), every global minimizer of the loss function assigns distinct vector representations to distinct positions ($p_i^* \neq p_j^*$ for $i \neq j$).
• Significance: Training guarantees that the model will not collapse different positions into the same embedding, provided the task requires order sensitivity.

3. The Monotonicity Conjecture (Conjecture 5) & Proof in NTK Regime

• Claim: If the statistical distance between positions increases with their sequence distance (i.e., $|i-j| \le |i-k| \implies d_H(\mu_i, \mu_j) \le d_H(\mu_i, \mu_k)$), then the learned embeddings should preserve this order: $\|p_i^* - p_j^*\| \le \|p_i^* - p_k^*\|$.
• Proof: The paper provides a rigorous proof of this conjecture within the Neural Tangent Kernel (NTK) regime (where the attention matrix $M$ is nearly stationary). It establishes that the gradient flow converges to a monotone fixed point.

4. The Information-Optimal Encoding (Proposition 6 & Algorithm 1)

• Construction: Since the statistical manifold is curved, an exact isometric embedding into flat Euclidean space is generally impossible. The paper proposes Classical MDS on the Hellinger distance matrix as the best approximation.
• Algorithm:
  1. Compute Hellinger distances between all positional distributions.
  2. Construct the doubly-centered Gram matrix $B$.
  3. Perform eigendecomposition to find the optimal embedding $P_{MDS}$.
• Result: This $P_{MDS}$ minimizes the Stress criterion.
• Justification for Sinusoidal PE: The paper proves that if positional statistics vary smoothly (approximate stationarity), the MDS solution approximates a sinusoidal function. This provides the first theoretical justification for the original Vaswani et al. (2017) sinusoidal encoding.

5. Minimal Parametrization (Remark 8)

• The optimal encoding has an effective rank $r = \text{rank}(B) \le n-1$.
• Instead of a full $n \times d$ matrix ($nd$ parameters), the encoding can be represented as $P = AB^\top$ with $A \in \mathbb{R}^{n \times r}$ and $B \in \mathbb{R}^{d \times r}$, requiring only $r(n+d)$ parameters.
• This offers significant parameter savings when $r \ll d$.

4. Experimental Results

The theory was validated on a synthetic corpus and two real-world datasets (SST-2 and IMDB) using BERT-base.

• Stress Comparison:

  • MDS Encoding: Achieved near-zero stress (perfect isometry) on both datasets, confirming the theoretical construction.
  • ALiBi: Surprisingly achieved very low stress (0.563 on SST-2), significantly outperforming Sinusoidal and RoPE. This suggests that for short sentences (SST-2), the positional geometry is approximately rank-1 (shift-equivariant), which aligns with ALiBi's linear bias structure.
  • Sinusoidal & RoPE: Had identical, high stress values. Their performance is determined by their shared frequency schedule, not the absolute/relative distinction.
  • Random: Highest stress.

• Dimensionality Analysis:

  • The stress of MDS drops to zero once the embedding dimension $d$ exceeds the intrinsic rank $r$ of the corpus (e.g., $r=64$ for SST-2, $r=254$ for IMDB).
  • Sinusoidal and RoPE stress grows exponentially with $d$ due to frequency misalignment.

• Separation & Monotonicity:

  • Experiments confirmed the Positional Separation Theorem: minimum pairwise distance between embeddings remained strictly positive throughout training.
  • Monotonicity Violation: The MDS encoding had a violation rate of 22% (vs. 50% for random), supporting the Monotonicity Conjecture. Pre-trained models showed similar behavior (25.4%), while scratch-trained models were near random (48.2%) due to insufficient convergence time.

• Layer-wise Stress:

  • In pre-trained BERT, Layer 3 showed a massive spike in stress when projecting the sinusoidal PE through the attention matrix. This correlates with literature suggesting Layer 3 specializes in syntactic and positional processing, aggressively reorganizing the signal.

5. Significance and Impact

1. Theoretical Foundation: Moves positional encoding design from empirical heuristics to a principled geometric framework based on Information Geometry.
2. Justification of Existing Methods: Explains why sinusoidal encodings work (they approximate MDS for smooth corpora) and why ALiBi works (it approximates the rank-1 MDS solution for shift-equivariant corpora).
3. Efficiency: Demonstrates that positional encodings can be compressed significantly (low-rank parametrization) without losing geometric fidelity, potentially reducing model size.
4. Diagnostic Tool: The Stress metric provides a corpus-specific, training-free diagnostic to evaluate how well a chosen encoding matches the statistical structure of a specific dataset.
5. Open Problems: The paper identifies the extension of the Monotonicity proof beyond the NTK regime and the precise link between geometric stress and downstream task accuracy as key areas for future research.

In summary, this paper establishes that positional encodings are not arbitrary additives but are geometrically constrained by the statistical structure of the training corpus. The optimal encoding is the one that best preserves the Hellinger distances between positional distributions, achievable via MDS and approximable by low-rank structures.