Randomized YaRN Improves Length Generalization for Long-Context Reasoning

Source: arXiv:2606.23687 · Published 2026-06-22 · By Manas Mehta, Fangcong Yin, Greg Durrett

TL;DR

This paper addresses the challenge of length generalization in large language models (LLMs), where models pretrained on short sequences struggle to generalize to much longer context lengths during inference. The authors propose Randomized YaRN, a novel training method that combines YaRN-based positional encoding scaling with a randomized position sampling scheme and a progressive length curriculum. By randomly assigning tokens positional encodings sampled from a larger position range during training on short sequences, the model is exposed to out-of-distribution (OOD) positional representations, stimulating better positional extrapolation ability. Increasing the maximum sampled position length gradually over training epochs further improves robustness. The method is evaluated on two challenging long-context reasoning benchmarks, BABILong and Multi-Round Coreference Resolution (MRCR), using 7B-parameter models fine-tuned on data with <8K context. Randomized YaRN consistently improves accuracy at OOD context lengths ranging from 16K to 128K tokens and outperforms multiple baselines including vanilla LoRA fine-tuning, standard YaRN at inference, and Randomized Positional Encoding (RPE) alone. Ablations highlight the critical contribution of the length curriculum to performance gains. Overall, the results demonstrate that progressively randomized positional encoding combined with YaRN yields substantial improvements in long-context reasoning without requiring full long-context post-pretraining.

Key findings

• Randomized YaRN achieves up to 90.3% average accuracy on BABILong OOD lengths (16K–128K) with Qwen2.5-7B, surpassing Trained YaRN (81.3%) and RPE (79.7%) by nearly 9-11 percentage points (Table 1).
• On MRCR, Randomized YaRN yields 72.7% average OOD accuracy with Qwen2.5-7B, improving over Trained YaRN (59.3%) and RPE (61.1%) by over 10 points, and similarly outperforms baselines on Olmo3-7B (88.0% vs 83.2%/N/A) (Table 2).
• Performance gains of Randomized YaRN are largest at far out-of-distribution lengths like 128K for BABILong (up to +12.4 points over baselines) and 64-128K for MRCR (Table 1 & 2).
• Ablation of the length curriculum shows OOD accuracy drops by up to 18.3% on MRCR, demonstrating that gradually increasing sampled max positional length during training is critical for robust extrapolation (Table 3).
• Zero-shot inference with YaRN positional extrapolation alone barely improves OOD accuracy (around +1%) over zero-shot without YaRN, underscoring the need for training-time application (Tables 1 & 2).
• Training with Randomized YaRN requires only <5K training examples on context lengths <8K but can generalize convincingly to context lengths up to 128K tokens, a 16x increase.
• Randomized YaRN is effective on distinct architectures with different attention types (Qwen2.5 with grouped query attention and Olmo3 with sliding window attention), showing general applicability.
• Learning rate sweeps indicate Randomized YaRN outperforms trained YaRN consistently across learning rates, with peak gains (4.8 points) at optimal rates for Olmo3 on BABILong (Figure 2).

Methodology — deep read

1. Threat Model & Assumptions: The adversary is typically implicit here; this is a model robustness/training method paper focused on enabling LLMs to extrapolate to unseen long context lengths during inference, not adversarial attacks. The setting assumes access only to short-context training data (<8K tokens), and the challenge is to generalize to longer out-of-distribution (OOD) context lengths (up to 128K tokens). The model does not have ground truth examples at such long lengths during training.

2. Data: Two benchmarks are used: BABILong (a multi-hop reasoning dataset with long passages padded with irrelevant text from PG19) and MRCR (multi-round coreference in long multi-turn conversations). The training sets are small—e.g., ~4K training examples for BABILong with context length ≤8K, and 60 examples for MRCR in 4-8K lengths. Evaluation samples are organized into multiple length bins including very long contexts from 16K to 128K tokens (e.g., up to 305 samples per length bin for BABILong, 147 samples for MRCR). Data preprocessing involves embedding facts into long contexts and padding.

3. Architecture/Algorithm: The base models are Qwen2.5-7B-Instruct and Olmo3-7B-Instruct, large transformer-based LLMs with different attention mechanisms. Both models use Rotary Position Embeddings (RoPE) as default. YaRN (Peng et al. 2024) is a positional encoding scaling method that extends RoPE through an interpolation approach controlled by scale factor s, enabling extrapolation up to s · Lpre in context length. Randomized YaRN builds on YaRN by first applying randomized position encoding (RPE) during training: for each training sequence of length L, actual token positions are randomly sampled without replacement from a larger index set {1,...,Lt} where Lt > L and sorted to maintain order. The token’s positional encoding is then assigned as YaRN(positional_index; s) where positional_index ∈ {1,...,Lt}. This serves to expose the model to positional embeddings out-of-distribution relative to the native training length. A progressive length curriculum is employed where Lt starts near Lpre and grows slowly over epochs to the maximum target L’ to stabilize learning.

4. Training Regime: Models are fine-tuned with LoRA applied to all linear layers with rank 16 and alpha 32, dropout 0.1. Training uses BF16 mixed precision, cosine learning rate with warm-up steps. Epochs vary (e.g. 2 epochs for BABILong, 5 for MRCR). Batch sizes of 4 are used. Different YaRN scale factors s are considered during training and s’ during inference (allowing further extrapolation at inference time). The length curriculum increments maximum sampling length every epoch (e.g., 8K,16K,24K...). Experiments run on single or double high-memory NVIDIA GPUs. Total compute ~250 GPU hours.

5. Evaluation Protocol: Models are evaluated zero-shot and after fine-tuning on short context training data, and tested on held-out sets with much longer context lengths. Metrics are exact match accuracy for BABILong and a similarity-based metric for MRCR aligned with official benchmarks. Results report average accuracy across length bins and highlight performance at in-distribution vs out-of-distribution lengths. Baselines include zero-shot with/without YaRN at inference, vanilla LoRA fine-tuning without YaRN, trained YaRN without randomized position sampling, and RPE training without YaRN. Ablations on training curriculum and learning rates are performed.

6. Reproducibility: Code and data links provided (https://github.com/Manas-Mehta/Randomized-YaRN). Datasets BABILong and MRCR are publicly available. Detailed hyperparameters and configurations are documented including learning rates, batch size, epoch count, and YaRN scale settings. The positional encoding modifications and curriculum schedule are fully described.

Example walk-through: During training on a 4K-length example (short context), instead of assigning token positions 1 to 4000 directly, positions are randomly sampled without replacement from a larger set, say up to 16K, then sorted. Tokens receive YaRN positional encodings computed with scale s for these new sampled positions. Over epochs, the maximum position sampling range expands (e.g., 8K to 16K to 24K...). After training, the model can be tested on sequences up to 128K tokens by computing YaRN encodings with a larger inference scale s’. This encourages the model to handle unseen positional embeddings and perform reasoning over very long passages.

Technical innovations

• Combining YaRN positional encoding scaling with randomized positional encoding, allowing exposure of LLMs during training to out-of-distribution positional representations even on short-context data.
• Introducing a length generalization curriculum that gradually increases the maximum position sampling range in randomized YaRN during training, stabilizing and improving long-context extrapolation.
• Demonstrating that training with randomized YaRN positional encodings outperforms inference-time-only methods and standard fine-tuning baselines on challenging multi-hop reasoning and coreference benchmarks up to 128K tokens.
• Separately tuning YaRN scale factors for training and inference (s and s') to enable further extrapolation beyond the training range at test time.

Datasets

• BABILong — 4K training examples, up to 305 evaluation examples per length bin — publicly available at https://github.com/booydar/babilong
• MRCR — 60 training examples (4–8K context), 147 test examples across multiple length bins — publicly available at https://huggingface.co/datasets/openai/mrcr

Baselines vs proposed

• Zero-shot (no YaRN) on BABILong OOD accuracy = 23.3% vs Randomized YaRN = 90.3% (Qwen2.5-7B)
• Trained YaRN on BABILong OOD = 81.3% vs Randomized YaRN = 90.3% (Qwen2.5-7B)
• RPE on BABILong OOD = 79.7% vs Randomized YaRN = 90.3% (Qwen2.5-7B)
• Zero-shot + YaRN inference on MRCR OOD average = 26.3% vs Randomized YaRN = 72.7% (Qwen2.5-7B)
• Trained YaRN on MRCR OOD average = 59.3% vs Randomized YaRN = 72.7% (Qwen2.5-7B)
• RPE on MRCR OOD average = 61.1% vs Randomized YaRN = 72.7% (Qwen2.5-7B)
• Without length curriculum, Randomized YaRN OOD accuracy drops to 69.4% from 77.3% on MRCR (Qwen2.5-7B)

Figures from the paper

Figures are reproduced from the source paper for academic discussion. Original copyright: the paper authors. See arXiv:2606.23687.

Fig 2: Learning-rate sweep results on BABILong

Limitations

• Method requires that inference context length is substantially larger than training context length; not tested on equal or shorter lengths.
• Experiments limited to 7B-scale LLMs due to computational constraints; larger models not evaluated.
• Evaluations confined to English-language long-context reasoning tasks; multilingual or cross-domain generalization untested.
• Randomized YaRN introduces additional training complexity and hyperparameters (scale factors, curriculum schedule) requiring careful tuning.
• Performance on non-reasoning or simpler long-context tasks is not explored; utility beyond reasoning benchmarks is unclear.
• No adversarial robustness or security-related evaluation is provided; focus is on positional encoding generalization only.

Open questions / follow-ons

• How does Randomized YaRN scale to larger models (100B+ parameters) and even longer contexts beyond 128K tokens?
• Can the approach be combined with data-centric long-context pretraining to further improve generalization?
• What are the effects of randomized YaRN on other task types such as summarization or code generation over long contexts?
• Can similar curriculum and randomized positional encoding concepts be applied to other positional encoding schemes beyond RoPE/YaRN?

Why it matters for bot defense

For bot-defense and CAPTCHA practitioners, this work highlights cutting-edge progress in enabling large language models to maintain reasoning accuracy over very long input contexts, which could be pertinent when modeling or detecting sophisticated automated agents that leverage long context to mimic human behavior or solve reasoning challenges. Improving length generalization enables more reliable evaluation of models when adversaries attempt to manipulate or input unusually long or structured queries. Additionally, the idea of exposing models during training to positional encodings outside the training distribution might inspire analogous data augmentation or curriculum strategies in robust bot detection. However, this work primarily focuses on enhancing the fundamental positional generalization capabilities of LLMs rather than bot detection or CAPTCHA generation directly. Practitioners interested in building defenses may find inspiration on how position-based extrapolation impacts model behavior under distribution shifts but would need additional work to map these findings to adversarial settings or challenge-response verification systems.

Cite

@article{arxiv2606_23687,
  title={ Randomized YaRN Improves Length Generalization for Long-Context Reasoning },
  author={ Manas Mehta and Fangcong Yin and Greg Durrett },
  journal={arXiv preprint arXiv:2606.23687},
  year={ 2026 },
  url={https://arxiv.org/abs/2606.23687}
}

Read the full paper

• arXiv:2606.23687 (abstract)
• arXiv:2606.23687 (PDF)