Randomized YaRN Improves Length Generalization for Long-Context Reasoning
2026-06-22 • Computation and Language
Computation and Language
AI summaryⓘ
The authors propose a new training method called Randomized YaRN to help large language models better understand very long sequences of text. Instead of only training on short sequences, their method mixes in position information from much longer sequences during training, even when the actual input is short. This helps models get used to unusual position patterns so they can handle much longer texts later. They tested this on tough reasoning tasks and found it improved performance, especially on contexts much longer than those seen during training.
Large Language ModelsPositional EncodingLength GeneralizationRandomized YaRNCurriculum LearningOut-of-distributionLong-context ReasoningFine-tuningCoreference ResolutionSequence Modeling
Authors
Manas Mehta, Fangcong Yin, Greg Durrett
Abstract
Large language models (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.