Randomized YaRN Improves Length Generalization for Long-Context Reasoning

Abstract: LLMs are typically pretrained on short sequences and then extended to work on longer sequences with additional training. However, such LLMs still struggle to further generalize to very long sequences. We propose Randomized YaRN, a training method that improves length generalization by combining YaRN-based positional extrapolation with randomized positional encoding and a length curriculum. During training on short context data, tokens are assigned YaRN positional encodings sampled from a larger position range, exposing the model to out-of-distribution positional representations even on short-context inputs. We evaluate Randomized YaRN on two challenging long-context reasoning benchmarks, BABILong and Multi-Round Coreference Resolution (MRCR). When training on data with <8K context, Randomized YaRN consistently improves reasoning performance on context lengths from 16K to 128K and outperforms standard fine-tuning, with the largest gains appearing at far out-of-distribution lengths. Our results suggest that progressively exposing models to OOD positional distributions provides an effective recipe for generalizable long-context reasoning.

• The paper introduces Randomized YaRN, which samples progressively wider positional indices during LoRA training so models encounter out-of-distribution encodings before long-context evaluation.
• Randomized YaRN achieves 90.3% average BABILong accuracy on Qwen2.5-7B from 16K to 128K tokens, compared with 83.6% for LoRA and 81.3% for Trained YaRN.
• The length curriculum is essential: removing it reduces MRCR performance by up to 18.3 points, showing that gradual exposure to larger positional ranges improves robust long-context reasoning.

Randomized YaRN addresses a persistent gap between context-window extension and genuine length generalization: models fine-tuned on short-context data (<8K tokens) typically degrade sharply when evaluated at 16K–128K contexts, even when positional extrapolation methods such as YaRN are applied at inference time. The paper proposes a training method that combines YaRN-based positional scaling with randomized positional encoding (RPE) and a length-generalization curriculum, showing that exposing models to out-of-distribution (OOD) positional encodings during short-context training yields substantial gains on long-context reasoning benchmarks.

Motivation and positioning

The authors distinguish between two strategies for extending context: data-centric continued pre-training or instruction tuning on long corpora, which makes specific long sequences in-distribution but does not necessarily confer generalization; and positional-encoding modifications, which extrapolate from short-context data but have shown mixed results on reasoning tasks beyond the training distribution, despite strong performance on retrieval-style tests like Needle-in-a-Haystack. The paper targets the latter setting under an explicitly constrained regime: fewer than 5K training examples, all shorter than 8K context, evaluated up to 128K — a 16× extrapolation from the training distribution.

Method

The method builds on two components applied atop LoRA fine-tuning of all linear layers:

• Randomized position sampling over YaRN encodings: For each batch drawn from sequences of length $L$, a set of $L$ indices is sampled uniformly without replacement from $\{1,\dots,L_t\}$ where $L_t \gg L$, sorted ascending, and the $j$-th token receives the YaRN encoding $\mathrm{YaRN}(i_j; s)$ of the sampled index rather than its true position. Token order is preserved while absolute positions span a much larger range. This exposes the model to rotation angles it would otherwise only encounter at evaluation time in longer contexts.
• Length curriculum: The maximum sampling length $L_t$ increases across epochs ($L_1 \leq L_2 \leq \dots \leq L_T$, e.g., 8K → 16K → 24K), so the model first learns to use moderately extrapolated encodings before facing larger OOD indices.

At inference, standard YaRN is used with scale $s'$, typically equal to the training scale $s$. Notably, the paper reports that settings with $L$0 can generalize beyond $L$1, indicating that training with a smaller extrapolation range can transfer to a larger one at test time; $L$2 and $L$3 are therefore tuned independently as hyperparameters.

Experimental setup

Evaluations cover two reasoning-heavy benchmarks: BABILong's three-hop QA3 subset (with state-tracking multi-entry queries requiring temporal-order reasoning, not just single-fact retrieval), scored by exact match; and Multi-Round Coreference Resolution (MRCR), which requires verbatim retrieval plus an additional transformation amid distractor responses, scored by character-level similarity. Base models are Qwen2.5-7B-Instruct (grouped query attention, native 32K context) and Olmo3-7B-Instruct (sliding-window attention in 24 of 32 layers with full attention elsewhere; natively pre-trained with YaRN at scale 8 for 64K context). For Olmo3, YaRN/RPE modifications are applied only to full-attention layers, and vanilla LoRA and RPE-without-YaRN baselines are excluded as inapplicable. Baselines include zero-shot inference with and without inference-time YaRN, vanilla LoRA, Trained YaRN (LoRA + YaRN at train and test), and RPE without YaRN during training.

Results

Inference-time YaRN alone is insufficient. Adding YaRN only at inference improves average OOD accuracy on BABILong by roughly 1 point over zero-shot, and similarly fails to help on MRCR. Training with modified positional encodings is necessary for meaningful generalization.

Randomized YaRN outperforms all baselines at OOD lengths. On BABILong, Randomized YaRN reaches 90.3% average OOD accuracy (16K–128K) on Qwen2.5-7B versus 83.6% for vanilla LoRA and 81.3% for Trained YaRN, and 88.0% on Olmo3-7B versus 83.2% for Trained YaRN. Gains are largest at far-OOD lengths: at 128K, Qwen2.5 achieves 83.9% versus 63.0% for LoRA and 67.2% for Trained YaRN — a gap of nearly 17 points over the strongest baseline. On MRCR, Randomized YaRN attains 72.7% average OOD accuracy on Qwen2.5 versus 64.4% for LoRA, with the largest margin in the 64–128K bin (68.8% vs. 31.7% for LoRA); on Olmo3 it reaches 43.8% versus 38.1% for Trained YaRN. A consistent pattern is that Trained YaRN degrades steeply with length (e.g., Olmo3 drops from 98.0% at 16K to 67.2% at 128K on BABILong), whereas Randomized YaRN degrades more gracefully.

The curriculum is critical. Removing the length curriculum reduces OOD performance substantially for both RPE and Randomized YaRN on MRCR — RPE falls from 75.1% to 56.8% average OOD accuracy and Randomized YaRN from 77.3% to 69.4%, a degradation of up to 18.3 points. This establishes that gradually expanding the positional sampling range, rather than randomization alone, drives robust generalization.

Hyperparameter sensitivity is moderate. A learning-rate sweep on Olmo3/BABILong shows both Trained YaRN and Randomized YaRN peak at $L$4, with Randomized YaRN ahead by 4.8 points at the peak and generally better across learning rates. Total compute is modest: approximately 250 GPU-hours, with training on a single 80 GB GPU.

Limitations and open questions

The paper concedes several constraints. The method is designed for scenarios where evaluation context greatly exceeds training context, and its benefit profile outside that regime is unexamined. Experiments are limited to 7B-scale models due to compute constraints, leaving scaling behavior open. Evaluation is English-only. Additionally, the observation that inference scale $L$5 can outperform matched scales is reported empirically without a mechanistic account, and the interaction between randomized position sampling and hybrid attention architectures (sliding-window plus full attention) is handled heuristically by modifying only full-attention layers. Whether the curriculum schedule itself (epoch counts, step sizes) is optimal, or merely sufficient, remains an open question.

Conclusion

Randomized YaRN provides a simple, low-compute recipe for length generalization in long-context reasoning: sample positions from a progressively growing range, assign their YaRN encodings to short-context training data, and apply standard YaRN at inference. Across BABILong and MRCR and two architecturally distinct 7B models, it consistently beats LoRA, Trained YaRN, and RPE baselines at OOD lengths, with the largest gains at 64K–128K, and ablations confirm that the length curriculum is essential. The results support the broader claim that progressive exposure to OOD positional distributions — rather than more long-context data — is an effective path to generalizable long-context reasoning when only short-context supervision is available.