RePo: Language Models with Context Re-Positioning

AI summary

Overview

• Research area: Machine learning / large language model architecture, specifically position encoding and in-context learning.
• Technical level: Advanced. The paper assumes familiarity with Transformer attention, rotary position encoding (RoPE), and continual pre-training, though its central idea can be grasped without that background.
• One-sentence scope: The paper introduces RePo, a lightweight differentiable module that lets a language model assign its own continuous, non-linear position values to tokens instead of relying on fixed linear or constant indices, and shows gains on noisy-context, structured-data, and long-context tasks after continual pre-training of OLMo-2 1B and 7B.

What This Paper Is About

Most LLMs assign positions to tokens in a rigid way: either consecutive integer indices from 0 to L−1 (as in RoPE) or a single constant index for every token (as in NoPE). That fixed structure forces the attention layers to figure out the organization of the input on their own, spending capacity that could otherwise go toward reasoning. The paper's goal is to give the model a learnable module that re-positions tokens according to their relevance, so attention can focus on the content rather than reconstructing the input's structure.

Key Contributions

1. A context re-positioning mechanism (RePo). A light-weight, differentiable module f_phi maps each token's hidden state to a real-valued position. It first extracts a position representation from the hidden state using a SwiGLU sub-layer, then projects it to a scalar position per attention head. Because modern position encodings such as RoPE and ALiBi are continuous and differentiable, these assigned positions can be trained end-to-end with the LLM.

2. A design that reuses existing pre-trained models. RePo is inserted before the position encoding step and applied only from the ceil(L/3)-th layer upward (the 5th layer for the 1B model and the 10th for the 7B model), leaving lower layers unchanged. Positions are used only inside the attention score computation, so the auto-regressive order and KV-cache behavior are preserved. The paper states this adds a 0.9% increase in parameters.

3. A contamination-aware evaluation setup. The experiments continually pre-train OLMo-2 1B and 7B, chosen because the model is fully open-sourced (training data, weights, and code), to avoid biased or misleading results from data contamination.

4. Analyses of where the gains come from. The paper measures attention mass allocated to needle, query, and rest tokens; characterizes the range and pattern of assigned positions; checks general-benchmark performance; and reports FLOPs and inference-time overhead.

Main Findings

• Noisy context (RULER NIAH, 1B): RePo reaches an average of 91.3, versus RoPE at 85.9 (+5.4), NoPE at 80.2 (−5.7), R2N1 at 90.5 (+4.6), and N2R1 at 88.3 (+2.4). At 7B, RePo scores 97.9 versus RoPE at 97.3 (+0.6), with perfect 100.0 on the multi-value needle type.

• Structured data (HybridQA table, exact match): On the 1B model RePo scores 26.70 versus RoPE 24.43 (+2.27), NoPE 23.52 (−0.91), R2N1 25.11 (+0.68), and N2R1 23.86 (−0.57). At 7B, RePo reaches 37.61 versus RoPE 33.52 (+4.09).

• Longer context (LongBench, F1 for QA and few-shot, Rouge-L for summarization): On 1B, RePo averages 27.86 versus RoPE 20.93 (+6.93), NoPE 6.98 (−13.95), R2N1 22.95 (+2.02), and N2R1 11.29 (−9.64). At 7B, RePo averages 32.03 versus RoPE 25.65 (+6.38).

• General short-context tasks: RePo is competitive rather than better. On the 1B model it averages 53.73 versus RoPE 53.70 (+0.03); at 7B it averages 70.30 versus RoPE 70.92 (−0.62). NoPE drops to 47.74 (−5.96) on 1B, R2N1 is 53.88 (+0.18), and N2R1 is 50.02 (−3.68).

• Attention shifts toward distant relevant tokens: On NIAH within the 4K training context, attention mass (scale 10⁻²) to needle tokens is 2.013 for RePo versus 1.754 for linear assignment (e.g., RoPE) and 1.572 for constant assignment (e.g., NoPE). Attention to nearby query tokens is lower for RePo (1.046) than for linear (1.123) or constant (1.135) assignment. Attention to the rest of the context is similar across methods (0.015 for RePo, 0.014 for both others).

• Positions live in a denser, non-linear space: Across attention heads, RePo assigns larger positional distances on longer contexts, but the largest distance is still much smaller than the raw context length (2K or 4K). The paper interprets this as reducing the out-of-distribution problem that RoPE faces when extrapolating low-frequency dimensions.

• Learned positions are mostly neither constant nor monotonic: Splitting positions into chunks of Δ=16 tokens with ε=0.2, the Mono pattern appears in 4% of chunks, Constant in 22%, and Hybrid dominates. The authors describe this as a hybrid of prior strategies such as NoPE's constant a ± ε and RoPE-style monotonic sequences within a span.

• Low overhead: Training FLOPs for OLMo-2 1B on 50B tokens are 3.84 × 10²⁰ for RoPE and 3.87 × 10²⁰ for RePo; decoding time per token using vLLM is 0.0176 seconds for RoPE and 0.0182 seconds for RePo.

Methodology in Plain English

The authors start from OLMo-2 checkpoints that have completed stage-1 pre-training on 4 trillion tokens, and continually pre-train them on 50B tokens of stage-2 data with a 4096-token context, using 4 H100 GPUs for 50B tokens and keeping the training configuration and codebase identical to the released OLMo setup.

For the RePo variant, instead of giving each token a fixed index, a small network reads the token's hidden state and outputs a position value. That network has two parts: an extraction step that pulls position information out of the hidden state, and a projection step that turns it into a single number. The extraction step is shared across a layer, while the projection is learned separately for each attention head. These numbers then feed into the same RoPE function the base model already uses, so the change is a drop-in replacement rather than a new encoder.

Four comparisons are used: plain RoPE (linear assignment from 0 to 4095), NoPE (no position encoding, equivalent to constant assignment), R2N1 (two RoPE layers followed by one NoPE layer, repeated), and N2R1 (the reverse). Evaluation uses the allenai/olmes codebase across three categories: noisy context (RULER NIAH variants, within the 4K training length), structured data (HybridQA, also within 4K), and longer context (RULER and LongBench subsets of 4K to 16K tokens, with YaRN applied to RoPE layers for extrapolation across all methods).

Why This Matters

• Research impact: The work reframes position assignment as something a model can learn rather than something a designer fixes in advance, and it presents an alternative to the COPE approach, which the paper notes relies on attention logits, requires [B, L, L] tensors, and is incompatible with RoPE and flash attention. RePo is described as compatible with most differentiable position-encoding methods and usable on standard pre-trained LLMs without training from scratch.

• Real-world applications:

  • Long-document understanding, where relevant passages sit far from the question.
  • Retrieval-augmented generation, where retrieved passages are diluted by irrelevant context.
  • Agentic systems that accumulate long, mixed-purpose interaction histories.
  • Table and structured-data reasoning, where linearizing a table into text loses structural cues.

• Industry relevance: Because RePo adds 0.9% parameters, leaves auto-regressive order and KV-cache handling unchanged, and has near-identical FLOPs and per-token decoding time to the vanilla model, it is a relatively low-cost architectural adaptation that can be layered onto existing open-weight models through continual pre-training rather than pre-training from scratch. The paper's use of fully open-sourced OLMo-2 also gives practitioners a reproducible baseline.

Future Directions

• Understanding why Hybrid patterns dominate: The paper reports that most assigned positions are neither constant nor monotonic, but does not explain why the model prefers that mixture. A mechanistic account of what these hybrid patterns encode would be a natural next step.
• Re-positioning beyond attention scores: The authors note that queries and keys could in principle be sorted by their assigned positions within an attention head, but they avoid this because it would require recomputing the KV cache at every time step. Finding an efficient way to exploit this is an open engineering problem.
• Choosing the RePo insertion depth: RePo is applied from the ceil(L/3)-th layer up based on the assumption that lower layers capture surface-level features. The paper mentions preliminary trials where applying RePo to very few layers made it behave like RoPE, and applying it to the first 1/3 of layers made training unstable. A principled way to select which layers get RePo remains unresolved.
• Generalization to other position encodings and scales: The paper states RePo is not restricted to RoPE and could extend to other differentiable encodings such as ALiBi, and it evaluates only at 1B and 7B. Testing other encodings, larger scales, and whether the learned positions transfer across model families are all open.

Target Audience

This paper is most useful to machine learning researchers and engineers who work on Transformer architecture, position encoding, or long-context modeling, and who want a concrete, low-overhead alternative to fixed linear or constant position assignment. It is also relevant to practitioners doing continual pre-training of open-weight models who need robustness on noisy, structured, or long inputs without sacrificing general short-context performance. Readers without a background in attention mechanisms and position encoding will find the analyses (attention mass, position patterns, FLOPs) difficult to interpret, though the motivation and headline results are accessible.

Authors’ abstract

In-context learning is fundamental to modern Large Language Models (LLMs); however, prevailing architectures impose a rigid and fixed contextual structure by assigning linear or constant positional indices. The rigid position information poses the full burden of organizing the input structure to attention layers, thus reducing the amount of attention that could be allocated for more critical information. To address this, we propose RePo, a novel mechanism that alleviates the burden for attention layers via context re-positioning. Unlike conventional approaches, RePo utilizes a differentiable module, $f_φ$, to assign token positions that capture contextual dependencies, rather than replying on pre-defined order. By continually pre-training on the OLMo-2 1B \& 7B models, we demonstrate that RePo consistently enhances performance on tasks involving noisy contexts, structured data, and longer context length, while maintaining competitive performance on general short-context tasks. Analysis reveals that RePo successfully allocates more attention mass to distant but relevant information, assigns positions in a dense and non-linear space, and captures the intrinsic structure of the input context. Our code is at https://github.com/SakanaAI/repo.