Extending the Context of Pretrained LLMs by Dropping their Positional Embeddings
Extending the Context of Pretrained LLMs by Dropping their Positional Embeddings
tl;dr
We introduce DroPE, a method for extending context windows without long-context fine-tuning. By removing positional embeddings and running a short recalibration, we achieve seamless context extension that maintains base performance while far outperforming RoPE scaling methods. Proven effective across scales up to 7B parameters and trillion-token datasets.
DroPE uses positional embeddings as a training-time scaffold: pretrain with RoPE, then drop positional embeddings and briefly recalibrate at the original context length to recover RoPE-level perplexity while improving length generalization.
Introduction
Many valuable real-world tasks are long: reviewing a change that touches many files, continuing a month-old chat, or answering questions about a 200-page contract or an hours-long transcript. In these settings, the useful details often lie far into the provided context, and the model must keep track of names, variables, assumptions, instructions, or feedback across long stretches of text. Today’s strongest models do offer large context windows, but their accuracy and recall capabilities considerably suffer once we go past typical sequence lengths seen in training.
“So, why don’t we just train on longer sequences?”
Unfortunately, training language models on long sequences is not easy. Part of the difficulty lies in the data: truly long, clean, and relevant contexts are rare and expensive to curate. The other part is compute: attention compares every token to every other token, which means training costs grow quadratically with sequence length — making long-context training brutally expensive.
Therefore, what we actually want is length generalization: models that were trained on, say, 4k tokens should still reason over 16K or 32K at test time without retraining. Today’s transformer LMs often don’t. Push them past their training context, and they don’t merely get a bit worse; they break: when evaluated out-of-the-box on longer sequences, they fail to produce coherent completions, and even with different scaling tricks, they stop using information that sits far away in the prompt. Retrieval-heavy tasks, “needle-in-a-haystack” evaluations, and multi-hop reasoning across distant passages all fall apart.
Why does this happen? Transformers need a way to know where tokens are. Raw attention is permutation-invariant (treating its input like a bag of words: “man eats fish” and “fish eats man” look the same if you ignore order), so we need to inject positional information directly into the representations. The standard choice is positional embeddings (PEs), and the dominant one in modern LMs is RoPE (rotary positional embeddings). RoPE is fantastic for training—it bakes in a strong, learnable sense of order. But when we stretch sequences beyond what the model saw during training, RoPE becomes the main culprit: those position-dependent rotations drift out of distribution, and the very mechanism that made learning fast ends up warping long-range attention.
This post digs into that tension, why RoPE helps in-distribution yet undermines extrapolation, and shows a simple path to keep the training benefits without paying the long-context extrapolation price. We introduce DroPE—a simple method for extending a pretrained language model’s usable context without long-context fine-tuning:
Just drop the model’s positional embeddings after pretraining and run a short recalibration.
The result is a seamless zero-shot context extension that preserves in-context performance and far outperforms RoPE scaling methods and specialized long-context architectures on downstream tasks. We extensively show our method can be easily integrated across small to large parameter and data scales, with results on models with up to 7B parameters and trillions of pretraining tokens.
Why do we need PEs in the first place?
The defining feature of transformers is abandoning architectural inductive biases such as convolutions and recurrences in favor of the highly general self-attention layer. The attention mechanism does not directly encode relative distances between queries and keys. Therefore, raw attention is invariant to prefix permutations: for any permutation ( \sigma \in S_m ), if ( y_1, \ldots, y_n = Attn(x_1, \ldots, x_n) ), then
[ Attn(x_{\sigma^{-1}(1)}, \ldots, x_{\sigma^{-1}(n)}) = y_{\sigma^{-1}(1)}, \ldots, y_{\sigma^{-1}(n)}. ]
Therefore, for sequence modeling tasks such as language modeling, we need to directly inject positional information about the tokens through positional embeddings (PE) and causal masking. While the original motivation for causal masking was not to provide positional information, but instead to have efficient parallelizable training, it turns out that a consistent
While the original Attention Is All You Need paper suggested using absolute positional information (i.e., the token position), injected once for the input tokens, more recently, the community has settled on using Rotary Positional Embeddings, which encodes relative token positions by rotating key and query vectors on every attention head.
Visualizing rotary position embeddings
RoPE rotates each pair of embedding dimensions by an angle that grows with the relative token position. The rate at which the angle grows is determined by the frequency of the pair, with higher-frequency for lower indices and lower-frequency for higher indices. Drag the sliders to see the rotation at different relative token distances and observe how many full turns each pair has accumulated.
Transformers train faster with PEs
While several works have demonstrated the viability of language modeling without positional embeddings (using only the causal masking for positional information), transformers without PE, commonly referred to as NoPE transformers, consistently underperform their PE counterparts. Put differently, under a fixed data and compute budget, a RoPE transformer LM will achieve better results than a NoPE transformer LM.
Attention non-uniformity develops faster with RoPE. In essence, our analysis shows that RoPE plays a crucial role in breaking “attention uniformity” in transformer LMs, providing models with an important inductive bias that allows the model to efficiently learn important positional-aware features in its parameters. First, we define a non-uniformity measure for attention heads that captures the alignment of the head with a predefined positional pattern
[ A^c = \frac{1}{T} \sum_{1 \leq j \leq i \leq T}\alpha_{ij} c_{ij}. ]
Here, ( c ) encodes the positional pattern we are interested in. We then empirically examine ( A^c )‘s gradients at initialization for transformers with and without PEs. High gradient norm means that positional bias can develop fast, right off the bat, while low gradient norm means that, regardless of data, bias takes time to develop, and the heads remain uniform for a long time. As shown in the Figure below, across all layers, gradient norms are higher for RoPE transformers, meaning that attention heads can become diagonal or off-diagonal much faster. Since we know these types of heads are critical for language modeling, this explains the pretraining gap.
Why is long-context hard?
Even as we scale hardware and optimize kernels, self-attention remains a quadratic bottleneck: every token compares to every other token. Training at very long sequence lengths is therefore not just slower, it’s disproportionately memory- and compute-hungry. As a consequence of these prohibitive costs coupled with the scarcity of high-quality long context data, most models are pretrained at limited context sizes, with an expectation that they’ll still behave appropriately when we stretch them at test time. This “zero-shot context extension” has become a central challenge in developing frontier models.
While long-context performance is impacted by many factors, the first issue to address is that performance doesn’t simply taper off beyond the training window—the models often simply break. Push a standard transformer past its pretraining length, and its perplexity shoots through the roof, and the generated text stops being coherent. Why does this happen? The main culprit is positional embeddings.
Positional Embeddings pose a challenge
RoPE (and any PE scheme) has a failure mode: when test sequences are longer than what the model saw during training, the induced rotations (phases) move out of distribution. This means attention heads see attention scores never seen in training, and performance drops. Popular “RoPE-scaling” tricks (PI, NTK-aware scaling, YaRN) try to fix this by compressing low frequencies to keep phases in range. That preserves perplexity but quietly shifts semantic heads—the ones that match content across large distances—so the model behaves as if the context were effectively cropped to the original length. In practice, you get near-constant perplexity with poor long-range retrieval—exactly what long-context tasks need most.
Embeddings as a train-time scaffold
Taken together, the observations from the previous section imply that PEs are a key component for effective LM training, but are also a fundamental barrier to long-context generalization. This raises a natural question:
“Is it possible to harness the inductive bias from positional embeddings exclusively during pretraining?”
It turns out the answer is yes! We propose a new method for extending the context of LMs by Dropping their Positional Embeddings after pretraining (DroPE). Following a short calibration phase on the original pretraining data, performed at the original context length, DroPE models perfectly reproduce the “in-context” performance of the base model. This simple procedure unlocks strong zero-shot context generalization to unseen sequence lengths, beyond highly-tuned RoPE extensions and alternative architectures.
Integrating DroPE in mid-training with no additional cost
To demonstrate DroPE’s performance as a zero-cost addition to pretraining, in our first set of experiments, we train from scratch different LMs with half a billion parameters on 16B tokens. We repeat this recipe for RoPE and NoPE transformers, as well as our DroPE variant. We implement DroPE by taking the 14B tokens RoPE transformer checkpoint, removing positional embeddings from every layer, and resuming training for the final 2B tokens. Despite only recalibrating at the very end of training, at no extra cost, DroPE matches the final in-context validation perplexity of RoPE trained on the full 16B tokens, and improves over the NoPE baseline trained without positional embedding all the way.