Remember to Forget: Gated Adaptive Positional Encoding

2026-05-11

AI Summary

This work addresses the degradation of attention focus and retrieval performance in long-sequence extrapolation caused by phase mismatch in Rotary Position Embedding (RoPE). To mitigate this issue, the authors propose GAPE, a plug-and-play positional encoding enhancement that employs a query- and key-dependent gating mechanism. GAPE dynamically suppresses irrelevant context while preserving critical long-range information, all without compromising the geometric structure of RoPE. The method innovatively introduces a content-aware bias into attention logits, decoupling distance-induced decay from token importance to ensure accessibility of key tokens and enable controlled attention quality degradation. Experimental results demonstrate that GAPE significantly improves attention concentration and extrapolation robustness on both synthetic retrieval tasks and long-context benchmarks, outperforming existing RoPE-based baselines.

Abstract

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 prove that protected tokens remain accessible, while the attention mass assigned to unprotected distant tokens decays as a function of the query gate. We further show that GAPE can be implemented within standard scaled dot-product attention. We validate these properties empirically, finding that GAPE consistently yields sharper attention and improved long-context robustness over rotary baselines across both synthetic retrieval and long-context benchmarks.

Problem

• Rotary Positional Encoding
• long-context robustness
• attention dispersion
• out-of-distribution
• positional encoding

Innovation

• Gated Adaptive Positional Encoding
• Rotary Positional Encoding
• long-context robustness
• content-aware gating
• attention sparsity