Research map / Architectures and efficiency / Attention and transformers
Dataset Decomposition: Faster LLM Training with Variable Sequence Length Curriculum
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes dataset decomposition, which splits tokenized documents into power-of-two length buckets and trains with variable sequence length and batch size under a length curriculum, yielding faster and more accurate LLM pretraining than concat-and-chunk.
Problem
Fixed-length concat-and-chunk pretraining mixes unrelated documents within sequences, wastes attention compute on cross-document tokens, and breaks documents into short chunks, hurting accuracy and training efficiency at long context lengths.
Approach
Decompose each document into adjacent power-of-two length subsequences stored in length buckets, so every sequence comes from a single document; sample buckets per step with a tunable length curriculum and variable batch size so token count per step stays constant while attention cost tracks actual lengths.
Key results
- Dataset decomposition with the Grow-P2 curriculum (8 cycles) reaches regular average 54.4 and long-context average 30.9 at 103B tokens, versus 51.5 and 27.5 for Baseline-8k.
- DD reaches the baseline's accuracy with less than half the tokens (>2x data efficiency) in OpenLM-1B dataset scaling experiments.
- DD step time is 244 ms versus 304 ms for the 8k baseline, and the paper reports up to 6x faster training to reach target accuracy.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use dataset decomposition with Grow-P2 curriculum and 8 cycles instead of fixed-length concat-and-chunk | OpenLM-1B pretraining · Regular evaluations | average accuracy: 51.5 -> 54.4 | Helped |
| use dataset decomposition with Grow-P2 curriculum and 8 cycles instead of fixed-length concat-and-chunk | OpenLM-1B pretraining · Long-context evaluations | average accuracy: 27.5 -> 30.9 | Helped |
| use dataset decomposition with Grow-P2 curriculum and 8 cycles instead of fixed-length concat-and-chunk | OpenLM-1B pretraining · Training | step time (ms): 304 -> 244 | Helped |
| add document masking | Baseline-8k · Regular evaluations | average accuracy: 51.5 -> 52.4 | Helped |
| add document masking | Baseline-8k · Long-context evaluations | average accuracy: 27.5 -> 27.1 | Hurt |
| train only on sequences of length 1024 rather than the natural mixture (ablation) | OpenLM-1B with dataset decomposition · MDQA | average accuracy: 21.9 -> 0.2 | Hurt |
| exclude sequences longer than 2k (ablation) | OpenLM-1B with dataset decomposition · MDQA | average accuracy: 21.9 -> 8.1 | Hurt |
| exclude sequences shorter than 1k (ablation) | OpenLM-1B with dataset decomposition · MDQA | average accuracy: 21.9 -> 22.1 | Helped |
| use Grow-P2 curriculum with 8 cycles instead of uniform sampling (ablation) | OpenLM-1B with ≥256 mixture · MDQA | average accuracy: 23.0 -> 28.0 | Helped |
| use Grow-P2 curriculum with 8 cycles instead of one cycle (ablation) | OpenLM-1B with ≥256 mixture · MDQA | average accuracy: 25.9 -> 28.0 | Helped |
| reverse the length curriculum from long-to-short instead of short-to-long (ablation) | OpenLM-1B with ≥256 mixture · Regular evaluations | average accuracy: 54.4 -> 50.3 | Hurt |
| use Grow-P100 with 8 cycles instead of Grow-P2 with 8 cycles (ablation) | OpenLM-1B with ≥256 mixture · MDQA | average accuracy: 28.0 -> 24.1 | Hurt |
Surprising findings
- Training on a single 1024-length bucket gives the best regular benchmark scores but performs poorly on long-context MDQA, showing a sequence-length vs. benchmark trade-off.
- Concatenating unrelated short sequences does not mitigate the length correlation; a length-matched chunked bucket underperforms a natively short bucket.
Limitations
- Training speed gains are significant only when the target sequence length is long enough; attention is not a dominant cost otherwise.
- Hyperparameters were not explicitly optimized and the same settings were reused across methods.
- Results are mostly on OpenLM models trained on RefinedWeb and DataComp-LM with up to 137B tokens.
Experimental setup
Datasets: RefinedWeb DataComp-LM Natural Questions-Open PIQA COPA OpenBookQA LAMBADA HellaSwag Winograd WinoGrande SQuAD BoolQ
Benchmarks: PIQA COPA OpenBookQA LAMBADA HellaSwag Winograd WinoGrande SQuAD BoolQ CoQA Jeopardy ARC-Easy
Metrics: accuracy training step time (ms) time-to-accuracy speedup data efficiency
Baselines: Baseline-8k concat-and-chunk Baseline-8k+DM document masking Pack-8k+DM best-fit packing ICLM in-context pretraining Baseline-2k
Builds on: OpenLM FlashAttention2 RoPE in-context pretraining best-fit packing
Closest papers by meaning
- Fixing It in Post: A Comparative Study of LLM Post-Training Data Quality and Model Performance (NeurIPS 2025 · similarity 0.90)
- Long-Short Alignment for Effective Long-Context Modeling in LLMs (ICML 2025 · similarity 0.90)
- Beyond Next Token Prediction: Patch-Level Training for Large Language Models (ICLR 2025 · similarity 0.90)
- Sheared LLaMA: Accelerating Language Model Pre-training via Structured Pruning (ICLR 2024 · similarity 0.90)
- InfLLM: Training-Free Long-Context Extrapolation for LLMs with an Efficient Context Memory (NeurIPS 2024 · similarity 0.90)
- Efficient Long Context Fine-tuning with Chunk Flow (ICML 2025 · similarity 0.90)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.