ArXiv

Randomized YaRN Improves Length Generalization for Long-Context Reasoning

Authors
Manas Mehta, Fangcong Yin, Greg Durrett
Categories
cs.CL
arXiv
https://arxiv.org/abs/2606.23687v1
PDF
https://arxiv.org/pdf/2606.23687v1

Brief

Randomized YaRN is a training method that improves large language models' length generalization by assigning YaRN positional encodings sampled from a larger position range and applying a length curriculum so short-context (<8K) inputs expose the model to OOD positional distributions. Evaluated on BABILong and MRCR, it yields consistent gains at 16K–128K context lengths and outperforms standard fine-tuning (results reported in the abstract).

Why it matters

Randomized YaRN combines YaRN-based positional extrapolation with randomized positional encodings (sampled from a larger position range) and a length curriculum to expose models to out-of-distribution positional representations during training on short contexts.

Key details

  • When trained on data with <8K context, Randomized YaRN consistently improves reasoning on BABILong and Multi‑Round Coreference Resolution (MRCR) at evaluation lengths from 16K to 128K, outperforming standard fine-tuning with the largest gains at far out-of-distribution lengths.
Source evidence

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.