RePo: Language Models with Context Re-Positioning

This paper introduces RePo, a novel mechanism that leverages a differentiable module to dynamically re-position tokens based on contextual dependencies rather than fixed linear indices, thereby reducing extraneous cognitive load and enhancing LLM performance on tasks involving noisy contexts, structured data, and long-range dependencies.

The Gist

Imagine you are trying to solve a complex puzzle, but instead of having the pieces laid out in a neat, logical order, they are dumped on the table in a chaotic pile. To find the piece you need, you have to scan the whole mess, ignoring the junk, and hope your brain doesn't get tired from the sheer effort of sorting it all out.

This is essentially what happens inside modern Large Language Models (LLMs) when they read long or messy texts.

Here is a simple breakdown of the paper "REPO: Language Models with Context Re-Positioning" using everyday analogies.

1. The Problem: The "Linear" Trap

Currently, most AI models read text like a person reading a book page by page, from left to right. They assign every word a number: 1, 2, 3, 4... all the way to the end.

• The Issue: This "linear" order is rigid. It doesn't care if word #100 is actually the most important clue for word #500.
• The Cognitive Load: The paper argues this is like forcing a human to solve a math problem while someone is shouting random numbers in their ear. The brain (or the AI) wastes energy just trying to figure out where things are, rather than what they mean. This wasted energy is called "extraneous cognitive load."

2. The Solution: REPO (The "Smart Librarian")

The authors propose a new system called REPO (Context Re-Positioning).

Imagine the AI has a Smart Librarian inside its brain.

• Old Way (Standard AI): The librarian puts every book on a shelf based strictly on its arrival time. Book #1 goes on shelf 1, Book #2 on shelf 2. If you need a book from the back, you have to walk all the way down the aisle.
• REPO Way: The Smart Librarian reads the books first. If it sees that Book #1 and Book #500 are talking about the same topic, it magically moves them to sit right next to each other on the shelf, regardless of when they arrived. It ignores the arrival order and organizes the books based on how they relate to each other.

3. How It Works (The "Magic Module")

The paper introduces a small, lightweight computer program (a "differentiable module") that acts as this librarian.

• Instead of saying, "You are word #500," it says, "You are word #500, but in terms of importance to the current question, you are actually sitting right next to word #10."
• It creates a flexible, non-linear map of the text. It can group related ideas together and push irrelevant "noise" (like ads or random sentences) to the side, even if they were physically written in the middle of the text.

4. The Results: Why It Matters

The researchers tested this on the OLMo open-source models. Here is what happened:

• The "Needle in a Haystack" Test: Imagine hiding a needle (the answer) in a giant haystack (a long document). Standard AIs often get lost in the hay. REPO, however, seems to "smell" the needle and zooms right to it, ignoring the rest of the hay.
• Structured Data: When reading tables or charts turned into text, REPO understands the structure better because it can group related rows together, rather than reading them strictly line-by-line.
• Longer Contexts: As the text gets longer (from 4,000 words to 16,000 words), standard AIs start to forget things. REPO stays sharp because it isn't wasting energy on the "distance" between words; it's focused on the "relationship" between them.

5. The Best Part: It's Efficient

You might think, "If the AI is reorganizing the whole text, isn't that slow?"
Surprisingly, no. The "Smart Librarian" is very lightweight. It adds almost no extra cost to the computer's processing power. It's like adding a sticky note to a book rather than rewriting the whole library.

Summary

REPO is like giving an AI a pair of glasses that lets it see the true connections between words, rather than just their order on the page. By letting the AI decide where to "place" information based on relevance, it frees up its brainpower to do deep reasoning, solve harder problems, and handle longer, messier documents without getting confused.

It's a shift from "Read it in the order it was written" to "Read it in the order that makes the most sense."

Technical Summary

Here is a detailed technical summary of the paper "REPO: Language Models with Context Re-Positioning".

1. Problem Statement

Current Large Language Models (LLMs) rely on In-Context Learning (ICL), where information within a limited context window is processed to solve tasks. However, prevailing architectures impose a rigid, fixed contextual structure by assigning tokens linear, consecutive integer indices (e.g., $0$to$L-1$) or constant indices.

The authors argue that this static organization creates extraneous cognitive load, a concept borrowed from Cognitive Load Theory (CLT). In human cognition, working memory capacity is finite; if too much capacity is consumed by processing the organization of information (extraneous load), less capacity remains for deep reasoning and attention allocation (germane load).

• The Issue: Rigid linear or constant position structures force the model to process irrelevant or noisy information with the same structural weight as critical dependencies.
• The Consequence: This leads to performance degradation in tasks requiring:
  • Noisy Contexts: Finding specific information ("needles") in large amounts of irrelevant data.
  • Structured Data: Processing tables or structured inputs where linearization loses semantic structure.
  • Long Contexts: Maintaining attention over long distances (e.g., 16K+ tokens) where linear position bias fails.

2. Methodology: REPO (Context Re-Positioning)

The authors propose REPO, a novel mechanism that allows the model to dynamically re-organize token positions based on their semantic relevance rather than their input order.

Core Mechanism

REPO introduces a lightweight, differentiable module, $f_\phi$, which assigns a continuous, real-valued position $z_i$ to each token $x_i$ based on its hidden state $h_i$.

• Position Representation: A lightweight SwiGLU sub-layer extracts a position representation $r_i$ from the token's hidden state $h_i$.
  $$r_i = \text{Swish}(h_i W_g) \odot (h_i W_c)$$
• Position Assignment: A linear transformation maps this representation to a new position value $z_i$.
  $$z_i = r_i W_z$$
• Integration: These learned positions $z_i$ replace the standard integer indices in the position encoding function (e.g., RoPE). The attention score between token $i$ and $j$ becomes:
  $$A^{REPO}_{i,j} = q_i^\top g_\theta(z_j - z_i) k_j$$
  Where $g_\theta$ is the standard rotary position encoding (RoPE) function, but applied to the learned distance $z_j - z_i$ rather than the fixed index distance $j - i$.

Training Strategy

• Continual Pre-training: The method is applied via continual pre-training on the open-source OLMo-2 models (1B and 7B parameters) to avoid data contamination issues.
• Layer Placement: To balance efficiency and effectiveness, REPO is applied starting from the 1/3rd layer of the model. Lower layers are kept with standard linear positions as they primarily capture surface-level features (syntax, POS tagging) that benefit less from reorganization.
• Efficiency: The module is lightweight, adding only 0.9% to the parameter count. It does not alter the auto-regressive generation order (KV cache remains unchanged); it only modifies the attention calculation logic.

3. Key Contributions

1. Theoretical Framework: The first work to explicitly frame position encoding limitations through Cognitive Load Theory, arguing that dynamic re-positioning reduces extraneous load and frees up working memory for reasoning.
2. Novel Architecture: The design of a differentiable, data-driven position assignment module ($f_\phi$) that learns to assign positions in a dense, non-linear continuous space.
3. Dynamic Adaptability: Unlike static methods (RoPE) or constant methods (NoPE), REPO learns to dynamically switch between patterns:
   • Constant: Assigning similar positions to related tokens (like NoPE).
   • Monotonic: Maintaining order for sequential dependencies (like RoPE).
   • Hybrid: A mix of both, capturing the intrinsic structure of the input (e.g., segmenting few-shot examples).
4. Open Source: Release of code and model weights based on fully open-source OLMo-2 models.

4. Experimental Results

The authors evaluated REPO on OLMo-2 1B and 7B models across three critical dimensions:

A. Noisy Context (Needle-in-a-Haystack)

• Benchmark: RULER (NIAH tasks).
• Result: REPO significantly outperformed RoPE.
  • OLMo-2 1B: +5.4 points average improvement over RoPE.
  • OLMo-2 7B: +0.6 points improvement, achieving perfect accuracy on multi-value tasks.
• Analysis: REPO allocates more attention mass to distant "needle" tokens and less to nearby "query" tokens, effectively breaking the locality bias inherent in linear position encodings.

B. Structured Data

• Benchmark: HybridQA (Table reasoning).
• Result: REPO improved Exact Match (EM) scores significantly.
  • OLMo-2 1B: +2.27 points over RoPE.
  • OLMo-2 7B: +4.09 points over RoPE.
• Insight: REPO better preserves latent structural cues when linearizing tables into text.

C. Long Context & Extrapolation

• Benchmark: RULER and LongBench (4K to 16K tokens).
• Result: REPO showed superior generalization to unseen context lengths.
  • OLMo-2 1B: +6.93 points on LongBench average.
  • OLMo-2 7B: +6.38 points improvement.
• Mechanism: The learned positions exist in a denser, non-linear space, allowing the model to generalize better to lengths beyond training (e.g., 4K $\to$ 16K) compared to standard RoPE.

D. General Tasks

• Result: REPO maintained competitive performance on standard short-context benchmarks (ARC, MMLU-Pro, etc.), demonstrating that the re-positioning mechanism does not degrade general capabilities.

5. Significance and Impact

• Architectural Shift: REPO challenges the "de facto standard" of fixed linear position encoding, suggesting that context-aware position assignment is a more efficient use of model capacity.
• Robustness: It offers a solution to the "needle-in-a-haystack" problem and long-context degradation without requiring massive architectural overhauls or expensive retraining from scratch.
• Interpretability: Analysis reveals that REPO learns to mimic human-like chunking (grouping related info) and structural segmentation, validating the Cognitive Load Theory motivation.
• Future Direction: This work opens a new avenue for flexible context management, where models actively restructure their input to optimize reasoning, potentially benefiting Retrieval-Augmented Generation (RAG) and Agentic systems.

In summary, REPO demonstrates that by allowing LLMs to "re-order" their internal view of the context based on relevance, we can significantly reduce cognitive load, leading to superior performance in complex, noisy, and long-context scenarios while maintaining efficiency.