Selecting key tokens improves long document training for large language models
Which Long Documents Matter? A Controlled Study of Data Selection for Context Extension
Computation and LanguageArtificial Intelligence
Summary
When teaching large language models to understand really long texts, selecting which parts of those texts to use is very important. The authors compared different ways to pick these long documents and found that choosing documents based on a special measure called LongPPL key-token density works better than picking them randomly. They also developed a faster way to calculate this measure that keeps most of the benefits. Models trained with this method performed better on various tests measuring how well they handle long text.
What this means in practice
- •For language model trainers: Improve training efficiency and downstream performance of large language models by selecting long documents using truncated LongPPL key-token density scores.
- •For ai service engineers: Enhance the context understanding of AI applications dealing with long documents by incorporating models trained with more effectively selected long document data.
Authors
Junlong Jia, Jiang Zhou, Xing Wu, Chaochen Gao Tinghao Yu, Feng Zhang, Lilin Wang
Abstract
Natural long documents are widely used for context extension, but which selection pipelines are effective remains unclear. We conduct a controlled comparison on Qwen2.5-3B at 128K under matched 6B-token continued-training budgets. Random sampling of natural long documents (random-long) outperforms repacking at longer evaluation lengths. Within the same candidate pool, selecting documents with high LongPPL key-token density yields further gains, unlike the evaluated generic-quality and LLM-judge selectors. We introduce truncLongPPL to reduce scoring cost by computing key-token density over a truncated prefix. At 128K, scoring the first half of the reference span yields downstream performance close to full-span scoring with a measured $2.28\times$ end-to-end scoring speedup. Across three training seeds, truncLongPPL achieves the highest mean scores among evaluated pipelines, exceeding random-long by 4.54 points on validation HELMET and 3.28 on held-out RULER, with gains most consistent in retrieval. At a 256K target context, it gains 2.82 average RULER points over random-long on Qwen2.5-3B. In short/long mixed training across Qwen2.5-3B, Qwen3-8B, and OLMo~3-7B, it gains 1.47--7.31 average RULER points over random-long mixing, with seven-task short-context averages within 0.23 points of the corresponding baselines. After supervised fine-tuning (SFT), Qwen3-8B and OLMo~3-7B models trained with selected long documents outperform their random-long counterparts by 1.13--1.67 points on LongBench v2 and 0.62--1.44 points on MRCR.