Token-Aware Phase Attention (TAPA)

• TAPA is a positional encoding method that employs a learnable phase function to modulate cosine attention scores based on token content and relative distance.
• It replaces fixed rotary positional embeddings with a dynamic, content-aware modulation that eliminates intrinsic distance bias while preserving non-degenerate token interactions.
• Empirical evidence shows TAPA robustly extrapolates to long contexts, reducing perplexity by up to 327× compared to RoPE and its extensions.

Token-Aware Phase Attention (TAPA) is a positional encoding methodology that integrates a learnable, content-dependent phase function into the attention mechanism of transformer architectures. Designed as a principled alternative to Rotary Positional Embedding (RoPE) and its extensions, TAPA addresses intrinsic distance-dependent biases present in RoPE, enabling stable long-context modeling without the need for post-hoc modifications, hyperparameter rescaling, or architectural changes. TAPA achieves robust extrapolation to unseen context lengths, preserves non-degenerate token interactions across large distances, and demonstrates significantly improved perplexity—especially on long-context tasks—relative to RoPE-based methods (Yu et al., 16 Sep 2025).

1. Mathematical Formulation of TAPA

TAPA modifies the standard attention score by introducing a phase-modulated cosine term, whose magnitude depends both on the relative position and the token content. The general attention formulation is

$$\mathrm{Attn}_{\phi,\mathcal{M},\alpha}(q,k) = q^\top \mathcal{M} \, k \times \cos\left(2\pi\,|m-n|^{\alpha}\, \phi(q,k)\right),$$

where $q,k \in \mathbb{R}^D$ are the query and key vectors at positions $m,n$, $\mathcal{M} \in \mathbb{R}^{D\times D}$ is the learnable amplitude matrix, $\phi(q,k)$ is the learnable phase function, and $\alpha > 0$ scales the phase with respect to relative distance $|m-n|$.

The principal parameterization employs a quadratic phase:

$$\phi(q,k) = q^\top \mathcal{N} k, \qquad \mathcal{N} \in \mathbb{R}^{D\times D} \; \text{learnable},$$

which models all pairwise second-order interactions between query and key dimensions. Each attention head's representation is divided into amplitude and phase segments: $q = (q_A, q_P), k = (k_A, k_P)$, with segmentation controlled by $\theta$:

$$q_A, k_A \in \mathbb{R}^{\theta D},\qquad q_P, k_P \in \mathbb{R}^{(1-\theta) D}.$$

The diagonalized, practical attention equation is then

$$\mathrm{Attn}_{\theta,\alpha}(q,k) = \frac{q_A^\top k_A}{\sqrt{\theta D}} \; \cos \left(2\pi\,|m-n|^\alpha\, \frac{q_P^\top k_P}{\sqrt{(1-\theta) D}} \right),$$

introducing only the hyperparameters $\theta, \alpha$, and requiring no architectural changes beyond substituting this computation for standard dot-product attention.

2. Integration into the Attention Mechanism

RoPE modifies the baseline dot-product attention $\frac{q^\top k}{\sqrt{D}}$ via a fixed, position-dependent, complex-valued rotation, realized as:

$$\mathrm{Attn}_\text{RoPE}(q^{(m)},k^{(n)}) = \frac{1}{\sqrt{D}} \sum_{d=0}^{D/2-1} \big(A^{(m,n)}_d \cos(2\pi(m-n)\theta_d) + B^{(m,n)}_d \sin(2\pi(m-n)\theta_d)\big),$$

where $A^{(m,n)}_d$ and $B^{(m,n)}_d$ are bilinear forms of query and key components. TAPA, in contrast, supplants this fixed trigonometric modulation with a dynamically computed, content-aware cosine factor as shown above, enabling learnable position-content interaction effects without the limitations imposed by RoPE’s fixed basis rotation.

3. Analysis of Distance-Dependent Bias and Variance

RoPE introduces a provable, intrinsic bias in the expected attention score with respect to token distance (Theorems 2.1, 2.2):

$$\mathbb{E}_{q,k}[\mathrm{Attn}_\text{RoPE}(q^{(m)},k^{(n)})] = \frac{1}{D} \sum_{d=0}^{D/2-1} \big(\mu_0 \cos 2\pi \lambda \theta_d + \nu_0 \sin 2\pi \lambda \theta_d\big), \qquad \lambda = m-n,$$

which allows arbitrary baseline shifts over distance and systematically preferences nearby tokens. Consequently, attention scores on long contexts become unstable, and distant token interactions degenerate.

TAPA overcomes this through its phase-cosine modulation. Under mild assumptions (joint density of $q_P, k_P$ Schwartz-class), Theorem 3.1 demonstrates that the expected attention decays rapidly to zero with growing distance:

$$\big|\mathbb{E}_{q,k}\, \mathrm{Attn}_{\theta,\alpha}(q^{(m)},k^{(n)})\big| = O(|m-n|^{-\alpha(1-\theta) D}),$$

thereby eliminating constant distance bias. Theorem 3.2 establishes that the attention variance remains bounded away from zero at large distances,

$$\liminf_{|m-n|\to\infty} \mathrm{Var}[\mathrm{Attn}_{\theta,\alpha}(q^{(m)},k^{(n)})] \geq \frac{\sigma_0^2}{2},$$

ensuring that signal is preserved and attention does not degenerate pointwise in the long-range regime.

4. Extrapolation and Fine-Tuning for Long-Context

TAPA’s mathematical structure, in which the relative distance $|m-n|$ enters as a power in the cosine modulation, allows immediate extension to context lengths far beyond those seen in pretraining. There is no requirement for input position “shifting,” frequency rescaling, or interpolation.

In large-scale experiments, pretraining proceeds identically to a standard transformer (with RoPE replaced by TAPA) at 8k context for 420B tokens. For extension to 32k, the model is fine-tuned on 32k segments for 500 steps (0.25% of pretraining tokens), retaining all hyperparameters: $\alpha = 0.5, \theta = 0.1$, with no additional schedule or rescaling. Evaluation employs sliding windows up to 64k length using FlashAttention or standard attention kernels, requiring no further adjustment.

5. Comparison to RoPE Extensions and Hyperparameter Stability

Conventional methods to extend RoPE’s context, such as base-frequency retuning, positional interpolation (PI), or non-uniform dimension-wise scaling (YaRN), all require ad hoc post-hoc modifications after pretraining and introduce additional tunable parameters. For instance, increasing RoPE’s base frequency necessitates fine-tuning the base to extremely large values at adaptation time; PI linearly rescales long positions into the original pretraining range via an interpolation factor; YaRN applies non-uniform scaling factors to frequency bases across RoPE dimensions.

TAPA dispenses entirely with these procedural modifications. Its hyperparameters, $\theta$ and $\alpha$, are fixed at initialization and need not be revisited. Optimization proceeds with the same scheduler and learning rate as the LLaMA3 architecture, and no new architectural components are added. This allows direct and robust adaptation to long context lengths without manual intervention.

6. Empirical Evaluation and Performance

Comprehensive experiments show that TAPA not only matches RoPE, PI, and YaRN at short and intermediate context lengths but also delivers substantial gains as context length grows. The following table summarizes test perplexities (PG19, LLaMA3-7B) for sliding windows from 1k to 64k after fine-tuning on 32k contexts:

Context Length	RoPE (b=5e5)	PI	YaRN	TAPA
1024	12.97	12.99	13.05	13.04
8192	11.97	11.98	12.03	12.07
16384	11.79	11.80	11.85	11.83
32768	12.96	12.97	12.16	11.74
49152	938.23	939.17	322.14	11.67
65536	2280.16	2282.44	1962.55	11.75

At and beyond 32k, TAPA perplexity is 9.4% lower than RoPE/PI and 3.5% lower than YaRN. The RoPE-family methods exhibit catastrophic perplexity increases (hundreds to thousands) as context length exceeds training regimes, whereas TAPA remains stable.

In zero-shot settings with no fine-tuning, all RoPE-family baselines degrade irrecoverably past 8k tokens (e.g., ~6000 at 16k, ~16000 at 32k), while TAPA degrades gracefully to 17.96 at 16k and 122.71 at 32k—representing 327× to 133× improvement over other methods.

Ablation studies confirm that the quadratic phase parameterization substantially outperforms a linear variant, with 1.3–1.6 points lower perplexity at short/mid ranges and dramatically greater stability at extreme length scales.

7. Theoretical and Practical Implications

TAPA provides a robust mechanism for eliminating the constant distance bias inherent in RoPE, while preserving high-variance, non-degenerate token interactions at all length scales. Its ability to extrapolate to previously unseen context lengths without manual intervention is operationally significant for long-context applications. The approach exhibits compatibility with current transformer training pipelines, introduces negligible computational or optimization overhead, and requires only light fine-tuning for substantial context extension.

The evidence supports TAPA as a theoretically sound and empirically effective mechanism for long-context attention, establishing a new paradigm for context-length extrapolation and stability in transformer models (Yu et al., 16 Sep 2025).