Randomized YaRN: Teaching LLMs to Reason Over Lengths They've Never Seen

Paper: 2606.23687 Authors: Manas Mehta, Fangcong Yin, Greg Durrett Categories: cs.CL

The Gap

Existing work on length generalization has a ceiling problem. Standard fine-tuning with rotary positional encoding (RoPE) works well on training-length sequences, but fails catastrophically on longer ones because the model never sees out-of-distribution (OOD) position values. YaRN (a RoPE extension) improves extrapolation by simply scaling position frequencies, but still: if you only train on 8K contexts, the model’s attention patterns have never needed to handle positions beyond that range. Post-training methods like positional interpolation compress positions back into the training range, losing high-frequency (local) detail. The community has been stuck between “train on long data (expensive)” and “accept degrading performance on long sequences (bad)”.

This paper asks: What if we could expose a model to long-range position signals while still training on cheap short data? Their answer: randomize the position assignments within each training batch, so a 4K-length input might get positions sampled from a 64K range, forcing the model to learn robust attention patterns that work for arbitrary sequence lengths.

[Problem: LLMs fail on long contexts]                  [Assumption: OOD positions are the bottleneck]
         |                                                          |
         v                                                          v
[Mainstream: train on short, fail on long]         [YaRN: scale positions linearly, still fails beyond ~2x]
         |                                                          |
         +-------------------[Gap: no method teaches OOD position robustness without long data]---+
                                                     |
                                                     v
                                          [Method: Randomized YaRN]
                                          - assign random positions from large range to short inputs
                                          - use length curriculum to gradually increase the range
                                                     |
                                                     v
                                          [Evidence: 2 benchmarks, 16K-128K, outperforms FT & YaRN]
                                                     |
                                                     v
                                          [Conclusion: training with randomized positional distribution is a recipe for length generalization]

The Increment

One sentence: Before this paper, length generalization required either expensive long-sequence training or accepting degraded performance at >2x extrapolation; after this paper, you can train on 8K data and get non-trivial reasoning performance at 128K by exposing the model to randomized OOD positions during training.

Core Mechanism

The method has three components that work together: a position sampling scheme, a length curriculum, and a YaRON base position encoding.

First, during training, each input sequence of actual length L is assigned a “virtual length” V sampled from a uniform range [L, M], where M is the maximum virtual length for that training stage. The token positions are then scaled by (V / L) and fed into YaRON’s position encoding. This means a 4K-token input might get positions that span 0 to 32K, forcing the attention mechanism to handle distances it has never seen.

Second, a length curriculum gradually increases M across training steps. Early in training, M might be set to 8K; later, it increases to 64K or 128K. This prevents the model from being overwhelmed by extreme OOD positions before it has learned basic attention patterns.

Third, the YaRON base frequency scaling is kept constant throughout, providing the backbone for extrapolation. The randomization essentially acts as a data augmentation that forces the model to learn position-agnostic attention patterns—patterns that depend only on relative distances, not on absolute position magnitudes.

[Training Data: short sequences, L=2K to 8K]
         |
         v
[Position Sampler: sample V ~ Uniform(L, M)]
   - M starts small (e.g., 8K), grows via curriculum to 128K
   - scale: positions *= V / L
         |
         v
[YaRON Encoding: apply base-scaled RoPE to scaled positions]
   - base frequency scaling fixed (e.g., 1/10000^(2i/d))
         |
         v
[Transformer: processes position-scaled sequence]
   - attention patterns forced to handle OOD distances
   - gradients update weights towards length-robust representations
         |
         v
[Loss: standard cross-entropy on short-context labels]
         |
         v
[Gradually increase M: length curriculum]
   - schedule: step-wise or cosine decay of "how far out to sample"

Think of it like training a pilot in a flight simulator before letting them fly a real plane. A traditional pilot learns to fly only within the training airport’s airspace (say, a 10 km radius). When sent to a different airport, they panic because the landmarks are all different and way farther apart. YaRN is like giving them a map that stretches the old landmarks to cover the new airport—helpful, but the relative positions are still confusing. The randomized training in this paper is like running the simulator with randomly generated landscapes: mountains, oceans, deserts, all at different scales. By forcing the pilot (model) to land safely in randomly scaled virtual environments (OOD position ranges), they learn to focus on *relative navigation cues (relative attention patterns) rather than memorized landmark positions (absolute position biases). The length curriculum is the training schedule: start with landscapes that are only slightly larger than the familiar airport, then gradually increase to alien worlds.

Key Concepts

• Length Generalization: A model’s ability to perform reasoning on sequences longer than those seen during training. This is not about memorization; it’s about the underlying attention and position mechanisms being robust enough to handle unseen distances. Example: a model trained on 8K-long Wikipedia articles should be able to answer a question that requires looking at token #10,000 and token #12,000 (a 2K relative distance it has seen, but an absolute position >8K it hasn’t). Most models fail here because they overfit to absolute position ranges, using them as “breadcrumbs” rather than learning general relative distance computations.

• YaRON (Yet Another RoPE NTK-aware scaling): A variant of Rotary Position Encoding (RoPE) that adjusts the frequency of position basis functions. Standard RoPE assigns each dimension a base frequency (like 1/10000^(2i/d)), and positions are computed as rotations. YaRON modifies these frequencies to “stretch” the encoding space, so that a position value of, say, 64000 maps to the same rotation angle that 8000 would have in standard RoPE. This lets the model reinterpret unseen large positions as “compressed” versions of known small positions. Think of it like switching a ruler from centimeters to kilometers: the physical distance might be huge, but the number of ticks the model sees stays manageable. The key insight in this paper is that YaRON provides a good *basis for extrapolation, but still needs the training exposure to OOD positions that randomized sampling provides.

• Length Curriculum: A training schedule where the maximum virtual length M is gradually increased. This is analogous to how you learn a new language: start with single words (short context), then short sentences (medium), then paragraphs (long). The curriculum prevents “catastrophic forgetting” of short-range patterns while slowly expanding the model’s comfort zone. In this paper, M starts at 8K (the max training length) and increases to 128K over a few hundred training steps. The rate of increase is a hyperparameter—too fast, and the model gets confused; too slow, and training time grows.

Framework Shift

Before (mainstream approach):                    After (this paper):
[Train on short data]                            [Train on short data, but...]
   |                                                  |
   v                                                  v
[Standard RoPE/YaRON on actual positions]        [Randomized YaRON on scaled positions]
   |                                                  |
   v                                                  v
[Attention sees only L-length positions]         [Attention sees positions up to M, M >> L]
   |                                                  |
   v                                                  v
[Test on long data: OOD positions trigger        [Test on long data: model has seen scattered
 catastrophic failure]                             OOD positions, generalizes better]
                                                      |
                                                      v
                                                  [Length curriculum gradually increases M]

From static position assignment to randomized position exposure, the core shift is: don’t let the model memorize “normal” absolute positions; instead, force it to learn position-agnostic attention by seeing randomly scaled versions of the same short inputs.

Expert Assessment

Problem choice: This is a real gap with high practical relevance. The “train short, test long” scenario is the default for most LLM applications—few teams can afford to pre-train on 128K sequences. The paper sits at the intersection of position encoding research (YaRN, PI, etc.) and length generalization, which is one of the hottest topics in LLM deployment. The approach is timely, not manufactured.

Method maturity: It’s a clever insight wrapped in a straightforward implementation—randomized position scaling during training is something any team with a few GPUs could try. There’s no complex architecture change. That said, the paper doesn’t deeply explore *why randomization works (is it the diversity of absolute positions? The relative distance variety? Something else?). There are simpler alternatives (e.g., random scaling of attention logits) that might achieve similar effects and are worth considering. Overall, the method is a neat trick but not a breakthrough—it’s “good engineering” more than “deep science.”

Experimental integrity: The baselines are fair: standard fine-tuning, YaRON only, and a length-curriculum-only ablation. The choice of benchmarks (BABILong for synthetic reasoning, MRCR for realistic coreference) covers both controlled and realistic evaluation. The numbers are clean: for BABILong, Randomized YaRON improves accuracy from ~30% to ~60% at 128K compared to YaRON alone. That’s a big delta. However, I have two concerns: (1) they only test on two benchmarks—is this robust across diverse tasks like long-document QA or summarization? (2) the 8K training limit means the model never sees true long-range dependencies; the improvement might be from better local pattern generalization rather than true long-range reasoning. No red flags, but the experimental scope is narrow.

Writing quality: The paper is clearly written, but the “related work” section is thin—almost no discussion of why other interpolation methods (e.g., Positional Interpolation, NTK-aware scaling) fail where this succeeds. The ablation study on curriculum schedule is good, but the analysis of *why randomization works is shallow. If the authors added a section on “attention pattern analysis” showing that randomized training makes attention heads more uniformly distributed across positions, the paper would jump from “good” to “great.” As it stands, the paper is a solid empirical contribution but leaves mechanistic understanding as future work.

Verdict: weak accept — a practical, well-executed idea that fills a clear gap, but lacks deep mechanistic insight and narrow experimental scope.

Takeaways

1. Randomized position scaling is a cheap drop-in augmentation: If you’re training an LLM with RoPE or YaRON, just add a random position scaling factor during training. It requires no data changes (same short inputs) and minimal code changes (one line of position scaling before the encoding layer). This should generalize beyond the specific method here—try it for any sequence model that needs length robustness.
2. Length curriculum matters more than scaling factor: The paper shows that a curriculum (gradually increasing M) is essential; a fixed large M from the start hurts performance. This principle—don’t shock the model with OOD signals too early—applies broadly to any form of data augmentation that extends the input distribution.
3. The “short data, long generalization” framing is a blueprint: For practitioners, the key insight is that you don’t need long-sequence data to achieve length generalization; you need to simulate the *positional statistics of long sequences within your short-data training. This idea could be extended to other distribution shifts (e.g., simulating OOD vocabulary distributions by randomly masking tokens).