Research map / Large language models / Reasoning and chain of thought
Beyond Next Token Prediction: Patch-Level Training for Large Language Models
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 patch-level training, which averages every K tokens into one patch and trains the LLM to predict the next patch for a fraction of the data, then finishes with token-level training, cutting training cost roughly in half without loss in quality.
Problem
Reducing the prohibitive compute cost (C ≈ 6ND) of pretraining LLMs without sacrificing performance, by increasing the information density of the text units processed per training step.
Approach
Aggregate K consecutive token embeddings (by averaging) into a patch, train a standard Transformer with a single output head to predict all K tokens of the next patch via a summed cross-entropy loss, then initialize a token-level model from these weights and continue token-level training on the remaining data.
Key results
- With K=4 and λ=2/3, Transformer-370M reaches 0.5× training cost with average zero-shot accuracy 42.5 vs 42.2 for the token-level baseline and comparable perplexity (10.7 vs 10.9).
- Transformer-2.7B with patch training at 0.5× cost reaches average zero-shot accuracy 53.5 vs 53.1 for the baseline.
- Measured speed-up on 8 A100 GPUs is about 3.5× with patch size 4, implying roughly 0.523× actual runtime.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use patch-level training with K=4 and λ=2/3 | Transformer-370M · WikiText-103 test | PPL: 10.9 -> 10.7 | Helped |
| use patch-level training with K=4 and λ=2/3 | Transformer-370M · zero-shot average across 6 benchmarks | accuracy: 42.2 -> 42.5 | Helped |
| use patch-level training with K=4 and λ=2/3 | Transformer-780M · WikiText-103 test | PPL: 9.2 -> 9.1 | Helped |
| use patch-level training with K=4 and λ=2/3 | Transformer-780M · zero-shot average across 6 benchmarks | accuracy: 45.5 -> 46.3 | Helped |
| use patch-level training with K=4 and λ=2/3 | Transformer-1.3B · zero-shot average across 6 benchmarks | accuracy: 48.5 -> 49.0 | Helped |
| use patch-level training with K=4 and λ=2/3 | Transformer-2.7B · WikiText-103 test | PPL: 7.1 -> 7.2 | Hurt |
| use patch-level training with K=4 and λ=2/3 | Transformer-370M · Pile training, compute cost | training cost: 1.0× -> 0.5× | Helped |
| reduce patch-training context from KT=8192 to T=2048 (ablation) | Transformer-370M + Patch (λ=1/2) · WikiText-103 test | PPL: PPL ↑ 0.06 | Hurt |
| reduce patch-training context from KT=8192 to T=2048 (ablation) | Transformer-370M + Patch (λ=1/2) · zero-shot evaluation | accuracy: zero-shot accuracy ↓ 0.2 | Hurt |
| add input and output linear projections to the patch-level model (ablation) | Transformer-370M patch-level training · subsequent token-level model | PPL: 11.46 -> 12.33 | Hurt |
| vary patch size from K=2 or K=4 to K=8 or K=16 (ablation) | Transformer-370M patch-level training · 90B training tokens | test loss | Hurt |
| increase λ beyond 3/4 (ablation) | Transformer-370M patch-level training · fixed data size | test loss | Hurt |
Surprising findings
- Patch-level initialization often matches or beats token-level training from scratch despite halving compute, which the authors attribute to possible regularization and reduced long-range token distance.
- Adding linear projections at both input and output of the patch model substantially hurts downstream token-level performance even though it lowers patch-level loss, so patch loss is not a reliable predictor of final performance.
Limitations
- Scalability is untested beyond 2.7B parameters and 360B tokens; the paper states state-of-the-art models are trained on much larger scales.
- No general scaling law for patch size K and fraction λ; optimal values are tuned on the specific settings studied.
- Speed-up falls short of theoretical K× due to data loading and gradient synchronization overhead.
Experimental setup
Datasets: Pile WikiText-103 Alpaca MMLU HellaSwag PIQA WinoGrande ARC-E ARC-C
Benchmarks: WikiText-103 MMLU HellaSwag PIQA WinoGrande ARC-E ARC-C MT-Bench
Metrics: perplexity normalized accuracy average zero-shot accuracy test NLL loss MT-Bench score tokens per second
Baselines: token-level training from scratch model growth methods
Builds on: LLaMA speculative decoding model growth mixup
Closest papers by meaning
- Better & Faster Large Language Models via Multi-token Prediction (ICML 2024 · similarity 0.92)
- Dataset Decomposition: Faster LLM Training with Variable Sequence Length Curriculum (NeurIPS 2024 · similarity 0.90)
- Sheared LLaMA: Accelerating Language Model Pre-training via Structured Pruning (ICLR 2024 · similarity 0.90)
- Knowledge Fusion of Large Language Models (ICLR 2024 · similarity 0.90)
- Compact Language Models via Pruning and Knowledge Distillation (NeurIPS 2024 · similarity 0.89)
- Model Merging in Pre-training of Large Language Models (NeurIPS 2025 · similarity 0.89)
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.