Remember to Forget: Gated Adaptive Positional Encoding

The gist The GAPE (Gated Adaptive Positional Encoding) introduces a drop-in augmentation for Rotary Positional Encoding that decouples distance-based suppression from token importance through querydependent and keydependent gates, improving long-context robustness up to 8× the training length.

How it works

GAPE augments the native rotary geometry of RoPE with a data-dependent content-aware mask over the attention logits, introducing two adaptive routing signals: a query gate that allows each query token to modulate how strongly its context is contracted, and a key landmark that controls how strongly each key is protected from distance-based penalization This mechanism selectively suppresses irrelevant background context while preserving distant useful information, without modifying the underlying rotary geometry The structural bias is realized exactly as a rank-at-most-two correction to the pre-softmax logits through the same matrix multiplication used for attention

The GAPE mask separates four factors with distinct roles: the relative distance (i − j)/T defines the positional envelope, gi controls query-dependent contraction, (1 − lj) modulates how strongly that contraction applies to key j, and Γh sets the head-wise strength The landmark gate lj continuously protects selected keys from distance-based suppression, with lj → 1 approaching full protection The query gate gi controls the effective positional horizon, while lj controls how strongly each key is exposed to the structural penalty

Key Mechanisms and Properties

The paper details several theoretical properties of GAPE, including its equivalence to other formulations and its behavior under specific conditions

[i] We introduce GAPE, a drop-in augmentation for positional encodings that contracts weakly protected distant context while preserving selected keys, with a FlashAttention-compatible implementation

[ii] We show that GAPE provides a mechanism for suppressing spurious OOD long-range alignments while preserving distant relevant tokens, and validate the resulting robustness in controlled Needle-in-a-Haystack retrieval up to 8× the training context length

[iii] We show that GAPE improves long-context robustness in language modeling and can be retrofitted into a pretrained 7B model without degrading standard capabilities, while improving performance at the longest evaluated context

Empirical Validation

The framework is validated through controlled retrieval and language modeling experiments

[i] GAPE improves long-context robustness across controlled retrieval and language-modeling experiments, extrapolating up to 8× the training length

[ii] GAPE improves retrieval robustness under OOD context extension by adapting context contraction while preserving relevant distant keys, consistent with the mitigation mechanism analyzed above

[iii] Across rotary encodings, augmenting RoPE and p-RoPE with GAPE improves robustness beyond the training horizon

The results show that GAPE yields lower perplexity in the OOD regime, especially at the largest evaluated context lengths

Mechanistic Insights

The analysis of the routing variables reveals how GAPE actively uses its structural pathway

[i] The landmark gate lj protects keys from distance-based suppression, and the query gate gi controls the strength of context contraction

[ii] GAPE learns non-uniform head-wise masks: several heads develop strong biases, while others remain weakly biased

The query gate gi adapts to the retrieval regime: close targets induce larger values and therefore stronger contraction, whereas far targets keep gi near zero, preserving a broader effective horizon

Conclusion

GAPE separates context contraction from key preservation: gi controls how strongly the current query contracts distant context, while lj continuously modulates how strongly each key is exposed to that contraction The head-specific amplitude Γh calibrates the overall strength of this structural pathway GAPE retains selected distant information while suppressing weakly protected background context, and its factorized form is compatible with fused SDPA and FlashAttention. The results support the complementary roles of positional extrapolation and adaptive structural routing.

--- Page 33 ---

Preprint

Remember to Forget:

Gated Adaptive Positional Encoding

Supplementary Material

CONTENTS

A Gated Adaptive Positional Encoding: Proofs and Further Explanations 15

A.1 Landmark Protection.............) 15

A.2 Modulation of Recency Bias 15

A.3 Softmax Normalization Growth 17

A.4 Entropy Contraction 18

A.5 Hardware-Native Factorisation 33

A.6 Translation Invariance 20

A.7 Mitigating Phase Hallucinations 21

A.8 GAPE Theoretical Properties in Needle-In-A-Haystack Retrieval 23

B Extended Experiments 25

B.1 nanoGPT: Training Setup and Datasets........) 25

B.2 OLMo-3 7B Retrofit: Setup and Detailed Results 25

B.3 Ablation Study of the GAPE Mask Components 27

B.4 Needle-In-A-Haystack Synthetic Experiment 28

B.5 Frequency-Channel Norm Distribution 30

B.6 FULL ROUTING DYNAMICS OF GAPE AND HEAD SPECIALISATION 32

--- Page 34 ---

Preprint

Remember to Forget:

Gated Adaptive Positional Encoding

FLOPs. For a single Transformer layer, both RoPE and GAPE share the same costs across the QKV projection, causal attention, output projection, and MLP FLOPsROPE = FLOPsGAPE = 2BHT2d + 24BT C2 + O(BT C)

KV-cache Memory at Inference. During autoregressive decoding, each layer stores past keys and values in a KV cache For GAPE, the landmark value and positional information are already encoded in the modified key tensor, so the modified key kt and the value vt are cached Thus, both RoPE and GAPE cache tensors of identical shape K, V ∈ R B×H×Tctx×d

Decode FLOPs per layer = 24C 2 + 4CTctx + O(C) where the O(C) term contains the positional-encoding computation

--- Page 35 ---

Preprint

Remember to Forget:

Gated Adaptive Positional Encoding

Attention Maps Entropy Analysis. Having verified GAPE’s gating dynamics in a controlled synthetic setting in Section 4, we now examine whether the same behavior emerges on real longcontext inputs GAPE produces lower-entropy attention maps than RoPE and p-RoPE across layers, suggesting that the learned mask suppresses diffuse background context

--- Page 36 ---

Preprint

Remember to Forget:

Gated Adaptive Positional Encoding

Head Specialization in Attention Maps.

1. Bold header: GAPE Augmentation for Long-Context Retrieval Robustness

This improvement allows AI systems to maintain retrieval accuracy when sequences are extended beyond their training horizon by introducing a content-aware mask that contracts weakly protected distant context while preserving selected keys. This directly addresses the OOD failure mode of spurious long-range alignments, enabling reliable performance in controlled retrieval up to 8× the training length.

2. Bold header: Adaptive Context Contraction via Query Gate

The system can dynamically adjust its focus based on query relevance by employing a query gate that allows each query token to modulate how strongly its context is contracted. This means the model can learn to strengthen recency when the answer is local and to suppress it when successful retrieval requires preserving long-range evidence, effectively setting an effective context length tailored to the specific task.

3. Bold header: Landmark Protection for Distant Evidence

The architecture can selectively retain critical distant information through a key landmark that controls how strongly each key is protected from distance-based penalization. This mechanism ensures that selected keys can remain accessible through landmark protection, preventing the loss of semantically relevant distant tokens that might otherwise be erased by rigid distance penalties.

4. Bold header: Improved Language Modeling Stability

By employing GAPE, language models exhibit improved performance in extrapolation regimes, as demonstrated by lower OOD perplexity compared to standard RoPE or p-RoPE. This allows the model to generate text with greater stability and accuracy when processing genuinely long documents that exceed the training context window.

5. Bold header: Retrofit Capability for Pretrained Models

GAPE can be retrofitted into a pretrained 7B model without degrading standard capabilities, suggesting that existing, high-performing models can be enhanced for long-context generalization with minimal retraining effort, making adaptation more practical.

Rotary Positional Encoding (RoPE) is widely used in modern large language models. However, when sequences are extended beyond the range seen during training, rotary phases can enter out-of-distribution regimes, leading to spurious long-range alignments, diffuse attention, and degraded retrieval. Existing remedies only partially address these failures, as they often trade local positional resolution for long-context stability. We propose GAPE (Gated Adaptive Positional Encoding), a drop-in augmentation for positional encodings that introduces a content-aware bias directly into the attention logits while preserving the rotary geometry. GAPE decouples distance-based suppression from token importance through a query-dependent gate that contracts irrelevant context and a key-dependent gate that preserves salient distant tokens. We show that weakly protected distant context is exponentially attenuated as a function of the query gate, while selected keys can remain accessible through landmark protection. We further show that GAPE can be implemented within standard scaled dot-product attention. Empirically, GAPE improves long-context robustness across controlled retrieval and language-modeling experiments, extrapolating up to 8x the training length. We further retrofit GAPE into a pretrained 7B model, maintaining performance on standard benchmarks and improving performance at the longest evaluated context. These results support adaptive context suppression as a complement to positional representation for long-context generalization.