Token-Aware Phase Attention (TAPA)

• Token-Aware Phase Attention (TAPA) is a novel positional encoding technique that replaces fixed phase shifts with a token-specific, learnable phase function to maintain stable long-range interactions.
• It mitigates the distance-dependent biases of standard RoPE by employing an amplitude/phase split and a quadratic phase function, ensuring that distant tokens retain meaningful influence.
• Experimental evaluations on the LLaMA3 7B architecture show that TAPA maintains low perplexity at extended context lengths (32K–64K tokens) while outperforming other methods.

Token-Aware Phase Attention (TAPA) is a positional encoding technique for transformer models that introduces a token-dependent, learnable phase modulation into the attention mechanism. Unlike standard Rotary Positional Embeddings (RoPE), which exhibit intrinsic distance-dependent biases and require post-hoc adjustment or hyperparameter retuning to handle long contexts, TAPA preserves long-range token interactions and allows transformers to extrapolate context length efficiently while maintaining stability and low perplexity at scales far beyond typical RoPE limits (Yu et al., 16 Sep 2025).

1. Rotary Positional Embedding (RoPE) and Its Limitations

RoPE encodes token position through complex rotations applied to each pair of query and key subcomponents in a transformer head of dimension $D$. For tokens at positions $m$ and $n$, RoPE computes the attention as

$\text{Attn}_\text{RoPE}(q^{(m)},k^{(n)}) = \frac{1}{\sqrt D}\Re\!\Bigl[ \sum_{d=0}^{D/2-1} q_{[2d:2d+1]}^\C\,(k_{[2d:2d+1]}^\C)^*\,e^{\,i(m-n)\theta_d} \Bigr]$

where $\theta_d$ is a base frequency and $q_{[2d:2d+1]}^\C$ denotes the complex embedding of the $d$th subcomponent.

Under i.i.d. assumptions on $q$ and $k$, the expectation of the RoPE attention kernel exhibits an explicit distance-dependent bias: $\E[\text{Attn}_\text{RoPE}] = \frac{1}{\sqrt D}\sum_{d}(\mu_0\cos 2\pi\lambda\theta_d + \nu_0\sin 2\pi\lambda\theta_d),\quad \lambda = m-n.$ Two key theorems characterize the resulting pathologies:

• As $|m-n|$ increases, the expected attention oscillates and can cluster arbitrarily within $[-\sqrt{\mu_0^2+\nu_0^2}, +\sqrt{\mu_0^2+\nu_0^2}]$, introducing instability at long context lengths (Theorem 2.1).
• RoPE intrinsically favors nearer tokens regardless of content, with a non-trivial bias between “medium” and “large” distances (Theorem 2.2).
• Increasing RoPE's base frequency can reduce—but not eliminate—this bias; however, practical usage still suffers from context “collapse” or blow-up at great lengths (Theorem 2.3) (Yu et al., 16 Sep 2025).

2. Token-Aware Phase Attention: Definition and Construction

TAPA replaces RoPE’s fixed, position-only phase shift with a learnable, token-dependent phase function, fundamentally altering the structure of positional encoding. The general TAPA formulation for tokens $q, k$ at positions $m, n$ is

$\boxed{ \text{Attn}_{\phi, \M, \alpha}(q, k) = q^\top\,\M\,k \cdot \cos\!\bigl(2\pi\,|m-n|^\alpha\,\phi(q,k)\bigr), }$

where:

• $\M \in \R^{D\times D}$ is a learnable amplitude matrix,
• $\alpha > 0$ controls distance scaling by a power law,
• $\phi(q,k)$ is a smooth, learnable function $\R^D \times \R^D \to \R$.

A practical instantiation of $\phi$ is a quadratic form: $$\phi(q,k) = q^\top\,\N\,k \qquad \N \in \R^{D\times D},$$ chosen for its stationary-phase property—yielding controlled, oscillatory cancellation over long distances.

To eliminate excess parameters and maintain computational efficiency, TAPA adopts an amplitude/phase split: $$q = (q_A, q_P),\; k = (k_A, k_P),\quad \dim(q_A) = \theta D,\;\dim(q_P) = (1-\theta)D,$$ with the specific attention kernel

$$\boxed{ \text{Attn}_{\theta, \alpha}(q, k) = \frac{q_A^\top k_A}{\sqrt{\theta D}} \cdot \cos\Bigl(2\pi\,|m-n|^\alpha\,\frac{q_P^\top k_P}{\sqrt{(1-\theta)D}}\Bigr) }$$

where $\M$ and $\N$ are block diagonal (zero off subspace), and $\theta$ is a hyperparameter controlling the split.

3. Theoretical Properties and Guarantees

TAPA’s theoretical analysis demonstrates that it overcomes RoPE’s critical limitations:

• Vanishing Bias: Under Schwartz-class density on $(q_A, k_A, q_P, k_P)$, the expected TAPA attention decays to zero polynomially in distance:

$\left|\E_{q,k} \text{Attn}_{\theta, \alpha}(q^{(m)}, k^{(n)})\right| < C(\rho, D) |m-n|^{-\alpha(1-\theta)D}$

This result (Theorem 3.3) is based on Fourier analysis of oscillatory integrals, confirming that TAPA removes the non-trivial, non-decaying bias even as $|m-n|\to\infty$.

• Non-Degeneracy: The variance of the attention kernel remains asymptotically non-zero for arbitrarily long-range pairs:

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

(Theorem 3.4), ensuring that distant tokens retain influence on the resulting attention output.

In contrast to RoPE, where content similarity is eventually overwhelmed by the positional bias, TAPA guarantees both the elimination of bias and preservation of signal variance over long distances.

4. Experimental Results and Comparative Evaluation

Experiments in (Yu et al., 16 Sep 2025) utilize the LLaMA3 7B architecture pretrained on the Pile (420B tokens), with subsequent long-context fine-tuning on PG19 chunks (32K tokens). The implementation uses PyTorch SDPA-style attention kernels without FlashAttention optimizations; RoPE and TAPA are compared directly under matched settings.

The principal findings are:

• At 1K–16K context, all methods exhibit similar perplexity (perp declines $\sim$13.0→11.8).
• At 32K context, TAPA achieves a perplexity of 11.74 (9.4% lower than RoPE/PI, 3.5% better than YaRN).
• At 49K–64K, RoPE/PI “collapse” to $\sim10^3$–$10^4$ perplexity, YaRN to 300–2000, while TAPA remains stable (11.67 at 49K, 11.75 at 64K).
• In zero-shot evaluation (no long-context fine-tuning), TAPA maintains perplexity 17.96 at 16K and 122.7 at 32K—over $100\times$ lower than the next best method.
• Ablation studies (Table 3) show that replacing the quadratic phase with a linear phase yields instability and poor long-range behavior, corroborating the theoretical necessity of the stationary-phase construction.

Context Length	RoPE/PI	YaRN	TAPA
1K–16K	~13.0–11.8 (stable)	stable	stable
32K	collapse	collapse	11.74
49K–64K	$10^3$–$10^4$ (collapse)	300–2000	11.67–11.75

TAPA imposes a minor computational overhead: each attention head performs an additional $QK^\top$ multiplication (for the phase), resulting in a $\sim2\times$ slowdown in unoptimized PyTorch compared to fused FlashAttention. A custom kernel implementation is suggested to close this gap.

5. Applications, Limitations, and Comparisons

TAPA is directly applicable to:

• Long-range language modeling tasks (e.g., summarization, retrieval-augmented generation, code),
• Transformers requiring contexts significantly beyond 8K tokens,
• Potentially multimodal architectures with very long memory.

Current limitations include:

• Evaluation is restricted to the LLaMA3 7B architecture and perplexity as the metric,
• Additional computational and memory cost without a specialized kernel,
• Restriction to quadratic (stationary) phase functions, while other phase families may be unexplored.

Compared to established extrapolation methods (Position-Interpolation, YaRN, RoPE with frequency adjustments), TAPA eliminates the need for post-hoc rescaling, hyperparameter tuning, or architectural modifications before/after pretraining.

6. Prospects and Future Directions

Future avenues proposed include:

• Efficient FlashAttention-style kernels or Triton implementations for phase attention,
• Scaling pretraining to 32K–1M context lengths,
• Systematic exploration of higher-order or adaptive phase families for phase function $\phi$,
• Benchmarking on diverse downstream long-context tasks (e.g., QA, reasoning, code completion),
• Extension to bidirectional or non-causal transformer settings.

A plausible implication is that TAPA, with efficient implementation and further phase function generalization, could form the basis of long-memory Transformer architectures suitable for both unimodal and multimodal extended context applications (Yu et al., 16 Sep 2025).