Randomized YaRN Improves Length Generalization for Long-Context Reasoning

2026-06-22

AI Summary

This work addresses the challenge that large language models pretrained on short contexts struggle to generalize to extremely long sequence inference. The authors propose a training strategy that integrates YaRN-based positional extrapolation, randomized distant position encoding, and length curriculum learning. This approach enables the model to effectively encounter out-of-distribution positional information during training—even when trained exclusively on contexts shorter than 8K—thereby substantially enhancing its length generalization capability. Evaluated on the BABILong and MRCR benchmarks, the method achieves significant performance gains across context lengths ranging from 16K to 128K, demonstrating state-of-the-art results particularly on very long out-of-distribution sequences.

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.

Problem

• length generalization
• long-context reasoning
• positional encoding
• out-of-distribution

Innovation

• Randomized YaRN
• length generalization
• positional extrapolation
• out-of-distribution generalization
• long-context reasoning