Research map / Large language models / Language, safety and interpretability
Not All Tokens Are What You Need for Pretraining
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 Selective Language Modeling (SLM), which scores tokens by excess loss against a reference model and trains only on the top-k% tokens, yielding RHO-1 models with better math and general pretraining efficiency.
Problem
Standard causal language modeling applies next-token loss uniformly to all pretraining tokens, even though many are noisy, already learned, or irrelevant to the target distribution, wasting compute and hurting data efficiency.
Approach
Analyze token-level loss trajectories during pretraining to show heterogeneous dynamics, then train a language model only on tokens with high excess loss relative to a reference model trained on high-quality data (Selective Language Modeling), with a token selection ratio of about 60-70%.
Key results
- RHO-1-Math-1B achieves 40.6% and RHO-1-Math-7B 51.8% on MATH after fine-tuning on ToRA-69k, state-of-the-art for their sizes.
- Continual pretraining on 15B OpenWebMath tokens improves average few-shot accuracy by 16.5% (1B) and 10.4% (7B) over CLM baselines; RHO-1-7B matches DeepSeekMath-7B with 15B tokens vs 500B.
- On 80B general tokens, RHO-1 (Tinyllama-1B) improves average performance by 6.8% across 15 benchmarks over continual CLM training.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Apply selective token loss instead of causal language modeling | Tinyllama-1.1B continual pretraining on OpenWebMath · GSM8K few-shot CoT | accuracy: 6.4 -> 29.8 | Helped |
| Apply selective token loss instead of causal language modeling | Tinyllama-1.1B continual pretraining on OpenWebMath · MATH few-shot CoT | accuracy: 2.4 -> 14.0 | Helped |
| Apply selective token loss instead of causal language modeling | Mistral-7B continual pretraining on OpenWebMath · GSM8K few-shot CoT | accuracy: 42.9 -> 66.9 | Helped |
| Apply selective token loss instead of causal language modeling | Mistral-7B continual pretraining on OpenWebMath · MATH few-shot CoT | accuracy: 22.2 -> 31.0 | Helped |
| Apply selective token loss rather than causal language modeling before ToRA-69k fine-tuning | Tinyllama-1B · MATH tool-integrated reasoning | accuracy: 38.4 -> 40.6 | Helped |
| Apply selective token loss rather than causal language modeling before ToRA-69k fine-tuning | Mistral-7B · MATH tool-integrated reasoning | accuracy: 48.4 -> 51.8 | Helped |
| Use reference-model loss to score tokens for self-referenced training (ablation) | Tinyllama-1.1B on OpenWebMath · Average few-shot accuracy | accuracy: 21.5 -> 23.9 | Helped |
| Use reference-model entropy to score tokens for self-referenced training (ablation) | Tinyllama-1.1B on OpenWebMath · Average few-shot accuracy | accuracy: 21.5 -> 23.0 | Helped |
| Select the intersection of tokens scored by reference loss and entropy (ablation) | Tinyllama-1.1B on OpenWebMath · Average few-shot accuracy | accuracy: 21.5 -> 24.8 | Helped |
| Use the intersection of reference-loss and entropy token selections (ablation) | Tinyllama-1.1B on PPile · Average few-shot accuracy | accuracy: 24.7 -> 26.5 | Helped |
| Use a smaller 1B reference model for token selection | Llama-2-7B continual pretraining · Average few-shot accuracy | accuracy: 43.5 -> 44.4 | Helped |
| Vary the token selection ratio (ablation) | Tinyllama-1B SLM on OpenWebMath · GSM8K and MATH few-shot | accuracy | Mixed |
Surprising findings
- Downstream loss tracks selected-token loss rather than overall loss, so reducing loss on all tokens is not necessary for better performance.
- Self-reference (reference model trained on the same pretraining corpus) still improves downstream accuracy by up to 3.3%.
Limitations
- Only verified on models up to 7B and datasets under 100B tokens; scalability to large models is untested.
- Training exclusively with SLM causes the unselected-token loss to rise and may cause overfitting to the reference model's domain; a general loss may be needed.
- Requires a high-quality reference model (or same-vocabulary model), which adds cost and vocabulary constraints.
Experimental setup
Datasets: OpenWebMath SlimPajama StarCoderData Tulu-v2 OpenHermes-2.5 ToRA-69k Proof-Pile-2
Benchmarks: GSM8K MATH SVAMP ASDiv MAWPS TabMWP MathQA MMLU-STEM SAT GSM-Hard MMLU BBH
Metrics: few-shot accuracy Pass@1 Pass@10 F1 validation loss
Baselines: Tinyllama-CT Mistral-CT DeepSeekMath-7B Gemma Qwen1.5 Phi-1.5 LLemma MAmmoTH ToRA
Builds on: Causal Language Modeling RHO-LOSS TinyLlama Mistral-7B ToRA
Closest papers by meaning
- Better & Faster Large Language Models via Multi-token Prediction (ICML 2024 · similarity 0.88)
- Beyond Next Token Prediction: Patch-Level Training for Large Language Models (ICLR 2025 · similarity 0.88)
- Data Engineering for Scaling Language Models to 128K Context (ICML 2024 · similarity 0.87)
- PolyPythias: Stability and Outliers across Fifty Language Model Pre-Training Runs (ICLR 2025 · similarity 0.87)
- On the Generalization Ability of Next-Token-Prediction Pretraining (ICML 2025 · similarity 0.87)
- DataMan: Data Manager for Pre-training Large Language Models (ICLR 2025 · similarity 0.87)
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.