Randomized YaRN Improves Length Generalization for Long-Context Reasoning

arXiv:2606.23687v2 Announce Type: replace 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 three challenging long-context reasoning benchmarks, BABILong, Multi-Round Coreference Resolution (MRCR), and LongBench v2. When training on data with short 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.

Related

• Retrospective Sparse Attention for Efficient Long-Context Generation
• GRKV: Global Regression for Training-Free KV Cache Compression in Long-Context LLMs
• E2LLM: Encoder Elongated Large Language Models for Long-Context Understanding and Reasoning
• SeDeM: Selective Decompression of Hidden-State Memories for Long-Context Question Answering
• Scaling Reasoning Hop Exposes Weaknesses: Demystifying and Improving Hop Generalization in Large Language Models

Source: arXiv cs.CL | 2026-09-01