PaperlayerView exact revision on arXiv

Source-linked AI summary

RePo: Language Models with Context Re-Positioning

Huayang Li, Tianyu Zhao, Deng Cai, Richard Sproat

arXiv:2512.14391v3cs.LGcs.AIcs.CL

TL;DR

Fixed positional indices leave attention layers to organize context, with the introduction identifying this as a burden for contextual reasoning. RePo learns context-dependent token positions, and continual pre-training experiments report gains across context-demanding benchmarks while maintaining comparable general-benchmark performance.

  • Problem

    Fixed positional schemes impose a rigid context organization, leaving attention layers to recover structure and limiting performance on tasks requiring long-range or fine-grained dependencies.

    Source: §1 · p. 1, §3.1 · p. 3

  • Method

    RePo uses a differentiable module to assign tokens continuous, non-constrained positions based on contextual relevance rather than predefined order.

    Source: §1 · pp. 1–2

  • Results

    On OLMo-2 1B, REPO outperforms ROPE by at least 5.4 points on noisy-context, 2.2 on structured-data, and 6.9 on long-context benchmarks, with comparable performance on general benchmarks.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Analysis finds that RePo allocates more attention to distant relevant tokens and learns positions that capture intrinsic context structure.

    Source: §1 · p. 2

    The method introduces no new data sources or supervision and operates within existing model architectures.

    Source: Impact Statement · p. 9

Abstract

from arXiv · show

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_\phi$, 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.

1. Introduction

Fixed positional schemes leave attention layers to organize linearized context, a burden that can impair tasks requiring long-range or fine-grained dependencies. RePo instead learns token positions from contextual relevance, and experiments report gains on context-demanding tasks alongside analyses of attention allocation and learned position patterns.

1. Introduction

  • Fixed position assignments leave attention responsible for recovering input structure, which can impair contextual reasoning on tasks with long-range or fine-grained dependencies.The introduction frames this as a competition for limited attention capacity between understanding structure and processing the input.Source: §1 · p. 1
  • REPO uses a differentiable module fϕ to assign tokens continuous, non-constrained positions based on their hidden states and contextual relevance.The module can be learned independently for each attention head, and continuous position encodings permit end-to-end optimization.Source: §1 · pp. 1–2
  • On OLMo-2 1B, REPO outperforms ROPE by at least 5.4 points on noisy-context, 2.2 on structured-data, and 6.9 on long-context benchmarks.It also achieves comparable performance on general benchmarks, and similar gains are reported at the 7B scale.Source: §1 · p. 2
  • REPO allocates more attention to distant relevant “needle” tokens, learns dense non-linear positions, and captures context structure such as few-shot-example segmentation.Its learned positional patterns also resemble a hybrid of constant and monotonic assignment strategies.Source: §1 · p. 2

2. Background

This section introduces attention scores and position encoding, then uses rotary position encoding as an example of how token positions affect query–key interactions. It contrasts the common linear assignment with NOPE’s constant-position interpretation.

2. Background

  • Position encoding maps token positions into embeddings or biases before attention, and ROPE applies a differentiable rotation to queries and keys using integer token indices.In ROPE, gθ(j − i) captures the relative positional relationship between tokens at distance j − i.Source: §2 · pp. 2–3
  • Self-attention projects token hidden states into query, key, and value representations, with an attention score computed between each token pair.The attention matrix A contains scores for all token pairs; the background simplifies to a single head, noting extension to multiple heads.Source: §2 · p. 2
  • Linear position assignment is standard in related work, while NOPE’s omission of position assignment is equivalent to assigning every token the same constant position.Source: §2 · p. 3

3. Methods

REPO replaces fixed linear context positions with learned, relevance-sensitive positions in a continuous, nonlinear space. Its module derives position representations from token hidden states and assigns positions used by positional encoding.

3.1. Overview

  • REPO assigns token positions based on contextual relevance, aiming to reduce the organizational burden fixed linear positions place on attention.The assigned positions are continuous and nonlinear, and the module operates before position encoding.Source: §3.1 · p. 3

3.2. Context Re-positioning

  • The module has two components: extracting position representations from token hidden states and assigning real-valued positions from those representations.Source: §3.2 · p. 3
  • A lightweight SwiGLU sub-layer extracts a lower-dimensional position representation r_i from each token’s hidden state h_i.The representation dimension d_p is smaller than the hidden-state dimension d.Source: §3.2 · p. 3
  • A linear transformation maps each position representation to a token position, offering comparable performance with lower latency than other tested strategies.Source: §3.2 · p. 3
  • Position representations are applied independently per layer, while position assignment is independent per attention head and can use differentiable positional encodings beyond RoPE.Source: §3.2 · p. 4
  • REPO begins at layer ceil(L/3), retaining standard positional encoding in lower layers because they primarily capture local, surface-level features.Source: §3.2 · p. 4
  • Assigned positions affect positional encoding without changing autoregressive token order, avoiding the KV-cache recomputation that sorting by assigned positions would require.Source: §3.2 · p. 4

4. Experiments

The experiments continually pre-train OLMo-2 1B and 7B models and test REPO on noisy, structured, and longer-context tasks; REPO improves across all three dimensions.

4.1. Setup

  • Experiments start from OLMo-2 1B and 7B stage-1 checkpoints, continue pre-training on 50B tokens with a 4K context, and evaluate longer contexts up to 16K.YaRN is applied to RoPE layers for context extrapolation.Source: §4.1 · pp. 4–5

4.2. Results

  • REPO yields performance gains across all three evaluation dimensions: noisy context, structured data, and longer context.Source: §4.2 · p. 5
  • On OLMo-2 1B, REPO reaches 91.3 average accuracy, 5.4 points above ROPE; on 7B, its average score rises 0.6 points over ROPE.The 7B gain corresponds to a reported 22% reduction in remaining errors and perfect accuracy on the multi-value task.Source: §4.2 · pp. 5–6
  • On HybridQA, REPO exceeds ROPE exact-match accuracy by 2.27 points at 1B and 4.09 points at 7B.Source: §4.2 · p. 6
  • REPO beats all baselines at 4K, with wider gaps at unseen 8K and 16K; LongBench average gains are 6.93 points at 1B and 6.38 at 7B over ROPE.LongBench includes QA and few-shot tasks scored with F1 and summarization tasks scored with Rouge-L.Source: Figure 2 · p. 5, §4.2 · p. 6, Table 3 · p. 6

5. Analyses

The analyses examine where REPO’s gains come from and what position patterns it learns. They report greater attention to distant relevant tokens, adaptive position patterns, comparable general-task performance, and low inference overhead.

5.1. Attention Mass on Relevant Tokens

  • REPO assigns substantially more attention mass to distant needle tokens than RoPE and NoPE, while RoPE favors nearby queries and NoPE varies less across context parts.The authors relate these attention patterns to REPO’s noisy-context gains and their Cognitive Load Theory motivation.Source: §5.1 · p. 7

5.2. Position Patterns Learned by REPO

  • Across 2K- and 4K-token contexts, REPO’s positional distances grow with context length but remain below raw context length and have a non-uniform distribution.The authors hypothesize that a denser, non-linear learned space can better use the training context length for long-context performance.Source: §5.2 · p. 7, Figure 3 · p. 7
  • Hybrid is the dominant assigned-position pattern, while Mono occurs in 4% and Constant in 22% of chunks; positions can also capture input structure such as few-shot example segmentation.The pattern statistics use RULER; the authors report consistent observations on other benchmarks.Source: §5.2 · pp. 7–8, Figure 4 · p. 8

5.3. Performance on General Tasks

  • REPO achieves performance comparable to RoPE on extensive general benchmarks despite the position-assignment change between pre-training and continued training.The authors interpret this as evidence that REPO learned from general data generalizes to diverse data, including short, precise questions needing little context reorganization.Source: §5.3 · p. 8, Table 5 · p. 8

5.4. Efficiency

  • REPO adds 0.9% to parameters, while its RULER inference time within training context length is comparable to the vanilla model.Table 6 compares FLOPs and inference time; the cited passage reports the parameter increase and inference-time comparison.Source: Table 6 · p. 8, §5.4 · p. 8

6. Related Work

Prior position-encoding work includes fixed and data-dependent schemes, while REPO predicts positions before encoding. The paper distinguishes its lightweight, broadly compatible context reorganization from both attention-logit gating and model-agnostic context editing.

6. Related Work

  • REPO predicts token positions before the position-encoding process, unlike methods that map position indices into attention biases or embeddings.The related-work discussion describes RoPE as a de facto choice in modern LLMs and says REPO is orthogonal to enhanced RoPE methods.Source: §6 · p. 8
  • Unlike COPE’s attention-logit gate, which predicts non-decreasing positions, REPO repositions tokens within context using a lightweight module compatible with most positional encodings and standard pre-trained LLMs.The paper notes COPE’s [B, L, L] computation, incompatibility with RoPE and flash attention, and resulting scalability limits.Source: §6 · pp. 8–9
  • REPO changes Transformer context organization, an approach the paper describes as orthogonal to model-agnostic document reordering and agentic context evolution.The cited examples include manually moving high-attention documents to the end of context and ACE’s iterative context evolution.Source: §6 · p. 9

7. Conclusion

REPO replaces rigid linear position encoding with a lightweight differentiable module that learns positions reflecting context structure and dependencies. In continual pre-training of OLMo-2, it outperforms strong baselines on contextual-dependency tasks, maintains general short-context performance, and allocates more attention to distant relevant tokens.

7. Conclusion

  • REPO replaces rigid linear position encoding with a lightweight differentiable module that learns positions capturing context structure and dependencies.Continual pre-training of OLMo-2 models yielded gains on noisy-context, structured-data, and long-context extrapolation benchmarks while maintaining general short-context performance.Source: §7 · p. 9
  • REPO mitigates locality bias by allocating more attention to distant, relevant needle tokens and learns adaptive position patterns in a denser, more non-linear space.The conclusion presents active context repositioning as a direction for flexible context management through LLM architecture.Source: §7 · p. 9

Impact Statement

The impact statement describes REPO as operating within existing model architectures without new data sources or supervision, while noting that improved contextual reasoning may amplify beneficial and harmful uses. The appendix contrasts REPO’s adaptive assignment with fixed linear and constant positions and records evaluation and context-extension settings.

Impact Statement

  • REPO uses existing model architectures without new data sources or supervision; improved contextual reasoning may amplify both beneficial and harmful LLM uses.The statement says the work does not inherently introduce ethical risks beyond those already associated with LLMs.Source: Impact Statement · p. 9

A. What’s the Real Difference between Conventional PEs, NoPE, and RePo?

B.1. Noisy Context

  • The noisy-context appendix lists four RULER NIAH evaluation task IDs.The listed tasks are ruler niah s 1 4096::std, ruler niah mk 1 4096::std, ruler niah mv 4096::std, and ruler niah mq 4096::std.Source: §B.1 · p. 13

B.2. Extrapolation

  • The extrapolation appendix specifies YaRN extensions to 8K and 16K tokens, using factors 2.0 and 4.0 respectively from 4096 original maximum position embeddings.The 16K setting is used for all LongBench experiments.Source: §B.2 · p. 13

B.3. General Tasks

The evaluation passages list Olmes tasks for general benchmarks and describe long-context evaluation on OLMo-2 7B using RULER with YaRN applied to all RoPE layers. The context-extrapolation results are presented in Figure 6.

B.3. General Tasks

  • Olmes evaluation tasks include ARC Challenge, ARC Easy, BoolQ, and CoQA.The task IDs are arc challenge:rc::large, arc easy:rc::olmes, boolq:rc::large, and coqa::large.Source: §B.3 · p. 13
  • The Olmes evaluation also includes DROP, HellaSwag, MMLU Pro, and TriviaQA.The task IDs are drop::large, hellaswag:rc::large, mmlu pro:cot::none, and triviaqa::olmes.Source: §B.3 · p. 13

B.4. Additional Results on Long Context

  • Long-context evaluation on OLMo-2 7B uses RULER data and applies YaRN to all RoPE layers to extend context; context-extrapolation results appear in Figure 6.The supplied passages do not report the results shown in the figure.Source: Figure 6 · p. 14, §B.4 · p. 14

B.5. Ablation Study

The starting-layer ablation finds overall performance robust to where RePo begins, with a small trade-off between general benchmarks and NIAH as the starting layer increases.

B.5. Ablation Study

  • Overall performance is robust to RePo’s starting layer, with results consistent across other evaluation benchmarks.The ablation evaluates OLMo-2 1B on NIAH and MMLUPro.Source: §B.5 · p. 14, Figure 5 · p. 14
  • Increasing the starting layer slightly improves general-benchmark performance while negatively affecting NIAH performance.Figure 5 reports the sensitivity experiment for starting layers l = 3, 5, and 7.Source: §B.5 · p. 14, Figure 5 · p. 14
  • Here, l marks the layer where RePo begins; lower layers retain vanilla RoPE.For example, l = 5 means RePo starts at the fifth LLM layer.Source: §B.5 · p. 14

C. Preliminary Experiments

In text reversal, RePo is tested on sequences whose relevant input tokens become farther away as generation proceeds. It outperforms NoPE and RoPE on longer evaluation sequences, while learned positions exhibit task-specific patterns without an imposed inductive bias.

C. Preliminary Experiments

  • The reversal task tests re-positioning when each generated token’s corresponding input token grows linearly farther away; locality bias does not apply.The input sequence is generated in reverse order.Source: C. Preliminary Experiments · p. 14

C.1. Setup

  • Training uses sequence lengths [2, 20] and evaluation uses [2, 30], with shared subwords lacking grammatical or semantic structure.The data and splits also follow Kazemnejad et al. (2023).Source: §C.1 · p. 14
  • The experiment trains a four-layer GPT-2 with NoPE, RoPE, and REPO8; RePo is a position-prediction module used with RoPE encoding in this work.The fϕ parameters are shared across attention heads within each layer.Source: §C.1 · p. 14

C.2. Findings

  • For input lengths L > 20, RePo outperforms NoPE and RoPE; all models are nearly perfect for L ≤ 20.Figure 8 reports accuracy across input lengths.Source: §C.2 · p. 16, Figure 8 · p. 16
  • Predicted positions differ from fixed indices: reversal-token pairs share positions in layers 0, 2, and 3, alongside hybrid prompt and mirror-symmetry patterns.The patterns emerged without task-specific inductive bias; Figure 7 plots input order against predicted position.Source: Figure 7 · p. 15, §C.2 · p. 16

D. Case Study

The case study visualizes RePo-assigned positions on MMLUPro few-shot examples. The positions vary across layers and heads, roughly align with example segmentation, and include negative values and punctuation-related outliers.

D. Case Study

  • Assigned positions roughly align with semantic segmentation in few-shot examples, and patterns differ across layers and attention heads.Figure 9 visualizes positions on MMLUPro few-shot examples.Source: D. Case Study · p. 16, Figure 9 · p. 17
  • Some assigned positions are negative, corresponding to reverse-direction rotations under RoPE; outliers correspond to non-informative punctuation and function words.Examples include “.” and “such.”Source: D. Case Study · p. 16
  • Figure 9 distinguishes prompt symbols in orange from few-shot questions in blue and answers in red.The figure visualizes positions assigned by RePo.Source: Figure 9 · p. 17
Loading 2512.14391v3…