Technical Summary: PJ-RoPE: A Fourier–Jet–Affine Position Space for Relative Attention

Problem Statement

Relative position representations are central to long-context attention in Transformers, yet existing methods are often treated as competing designs rather than components of a unified structure.

• RoPE encodes relative lag via Fourier characters (rotations).
• ALiBi adds a monotone affine recency term directly to attention logits.
• Jordan-RoPE replaces semisimple rotary frequencies with finite Jordan blocks, generating derivative-like "finite-jet" features.

The paper argues that these methods share a common algebraic structure rooted in constant-coefficient difference modules. Specifically, simple roots correspond to Fourier/RoPE characters, repeated nonzero roots generate Jordan/Fourier jets, and repeated unit roots yield ALiBi-like affine recency. The core problem addressed is the lack of a unified framework that treats these as selectable sectors within a single learnable space, and the absence of diagnostics to determine which sectors specific tasks actually utilize. Furthermore, high-order jet features (involving powers of relative distance) pose stability challenges at long contexts, causing transformed features and logits to grow dramatically.

Methodology

The authors propose PJ-RoPE (Poincaré-type RoPE), a framework that unifies these primitives into a learnable Fourier–Jet–Affine relative-position space.

1. Theoretical Formulation

The framework is grounded in the algebra of constant-coefficient difference modules. It defines a position space $P_{PJ}$ as a direct sum of three sectors:

• Homogeneous Fourier–Jet Sector ($H_{FJ}$): Contains RoPE (order-0) and Jordan-RoPE (order-$r > 0$) features. Algebraically, these are solutions to equations with repeated nonzero roots, yielding terms like $(d/L)^r e^{-cd/L} e^{i\omega d}$.
• Affine Recency Sector ($A_{rec}$): Contains constant and linear terms (e.g., $-sd/L$), corresponding to the repeated unit root. This supplies the "translation-like" direction analogous to the affine extension of the Lorentz group (hence "Poincaré-type").
• Light-Cone (LC) Sector ($H_{LC}$): A compactified coordinate chart for high-order behavior to address stability. It replaces raw distance $d$ with a rapidity-like phase $\phi_L(d) = L \text{asinh}(d/L)$ and a saturating amplitude $\beta_L(d) = d/\sqrt{d^2 + L^2}$.

The framework distinguishes between two implementation regimes:

• PJ-bias: A scalar additive attention kernel containing Fourier characters, damped jets, affine terms, and LC variants. It is learnable but not an exact feature transform.
• PJ-rotary: An exact feature-transform implementation applying relative actions to query and key features, used to verify closure properties of RoPE and Jordan-RoPE.

2. Adaptive Diagnostics

To determine which regions of the space tasks select, the authors introduce adaptive sector gates ($g_{FJ}, g_{aff}, g_{LC}$) and three diagnostics:

• Effective Mass: Measures the allocation of gated parameter magnitude across sectors and jet orders.
• Functional Energy: Measures the realized scale of kernel components over an evaluation window, accounting for damping and phase cancellation.
• Leave-One-Order-Out (LOO): Assesses whether removing a specific jet order degrades the fitted function, separating parameter selection from actual functional contribution.

3. Light-Cone Stabilization

To prevent the exponential growth of high-order jets at long contexts, the LC branch applies a coordinate transformation. This bounds the magnitude of transformed features and logits but compresses the phase resolution at long ranges, explicitly framing the trade-off between stability and resolution.

Key Contributions

1. Position-Space Formulation: The paper formulates relative attention as a learnable Fourier–Jet–Affine space, separating scalar bias kernels from exact rotary feature transforms.
2. Sector Containment: It identifies RoPE, Jordan-RoPE (finite jets), ALiBi (affine recency), and LC coordinates as distinct sectors within a single mathematical structure.
3. Adaptive Diagnostics: The introduction of sector gates, effective mass, functional energy, and LOO ablations provides a method to measure task-level sector selection empirically.
4. Evidence Chain: The authors evaluate sector recovery, trainable usage, and natural-task allocation across controlled kernels, synthetic tasks, byte-level language models, and symbolic music-token streams.

Experimental Results

The experiments utilize small causal Transformers (approx. 0.15M–0.20M parameters) trained at short contexts and evaluated at long contexts (up to 32,768 tokens).

• Fixed-Kernel Recovery: Scalar PJ-bias successfully recovers pure phase, finite-jet, affine, and LC targets, confirming the space spans the designed primitive families.
• Adaptive Sector Recovery: When trained with "teacher" signals (e.g., pure affine or specific jet orders), the learned gates and diagnostics correctly shift mass toward the corresponding sectors.
• Synthetic Sequence Tasks: Trainable attention models can utilize these sectors when the positional branch carries the signal. Affine teachers activate the recency branch; signed jet teachers activate FJ variants.
• Natural Language (Byte-Level): In language tasks (Tiny Shakespeare, WikiText-2, War and Peace), the models predominantly select affine/recency behavior. High-order FJ/LC mass appears only as a minor diagnostic correction. NTK-style RoPE scaling combined with affine recency yields the best long-context performance.
• Music-Token Streams: In symbolic music tasks (MAESTRO, MusicNet), the allocation differs. LC-affine variants remain strong, but there is a consistent, measurable allocation to high-order corrections (small but nonzero mass). This aligns with the need to model motif returns and rhythm envelopes.
• LC Stability vs. Resolution: LC variants successfully bound Q/K norms and cache/logit scales at 32,768 tokens. However, this comes at a cost: the "far-range phase span" is significantly reduced, leading to lower retrieval resolution in hard-negative probes compared to raw coordinates.

Significance and Claims

The paper positions PJ-RoPE not as a single new architecture, but as a position-space diagnostic framework. Its primary significance lies in reframing relative position mechanisms as a unified, learnable space where different tasks select different regions:

• Task-Dependent Selection: The empirical evidence suggests that language tasks favor affine/recency structures, while structured symbolic streams (like music) utilize a mix of affine and high-order jet corrections.
• Stability-Resolution Tradeoff: The work explicitly quantifies the trade-off introduced by Light-Cone coordinates: they stabilize long-context extrapolation by bounding growth but compress phase resolution.
• Separation of Concerns: The paper distinguishes between "representation correctness" (exact PJ-rotary closure) and "long-context stability" (scalar LC path), noting that exact feature transforms can degrade sharply at long ranges without stabilization.

The authors conclude that stable long-context use depends on a task-dependent allocation between affine recency, Fourier–jet structure, and LC compactification. They remain modest regarding generalization, noting that these findings are based on small models and specific tokenization schemes, and that larger-scale validation is required before these allocations can be treated as deployment rules.