PaperlayerView exact revision on arXiv

Source-linked AI summary

Rho-1: Not All Tokens Are What You Need

Zhenghao Lin, Zhibin Gou, Yeyun Gong, Xiao Liu, Yelong Shen, Ruochen Xu, Chen Lin, Yujiu Yang, Jian Jiao, Nan Duan, Weizhu Chen

arXiv:2404.07965v4cs.CLcs.AI

TL;DR

Uniform next-token loss can spend computation on noisy or nonessential tokens, even after document-level filtering. RHO-1 uses Selective Language Modeling to select tokens with high excess loss against a reference model; it reports improved performance and token efficiency across math and general continual pretraining.

  • Problem

    Document-level filtering leaves noisy tokens in high-quality datasets, and applying the same loss to every token can spend computation on nonessential tokens.

    Source: §1 · p. 2

  • Method

    Selective Language Modeling scores corpus tokens against a reference model and trains on tokens with high excess loss between the reference and training models.

    Source: §2.2 · pp. 3–4

  • Results

    RHO-1-7B matched DeepSeekMath-7B using 15B tokens versus 500B, while SLM improved results across math and general continual-pretraining evaluations.

    Source: §1 · p. 3

  • Takeaways & Limitations

    The paper concludes that token-level selection is an effective perspective for improving language-model pretraining in mathematics and general settings.

    Source: §4 · pp. 9–10

    Effectiveness was verified only on models up to 7B parameters and datasets below 100B tokens; scaling to very large models and datasets remains to be studied.

    Source: C Limitations and Future Work · p. 20

Abstract

from arXiv · show

Previous language model pre-training methods have uniformly applied a next-token prediction loss to all training tokens. Challenging this norm, we posit that "9l training". Our initial analysis examines token-level training dynamics of language model, revealing distinct loss patterns for different tokens. Leveraging these insights, we introduce a new language model called Rho-1. Unlike traditional LMs that learn to predict every next token in a corpus, Rho-1 employs Selective Language Modeling (SLM), which selectively trains on useful tokens that aligned with the desired distribution. This approach involves scoring pretraining tokens using a reference model, and then training the language model with a focused loss on tokens with higher scores. When continual pretraining on 15B OpenWebMath corpus, Rho-1 yields an absolute improvement in few-shot accuracy of up to 30% in 9 math tasks. After fine-tuning, Rho-1-1B and 7B achieved state-of-the-art results of 40.6% and 51.8% on MATH dataset, respectively - matching DeepSeekMath with only 3% of the pretraining tokens. Furthermore, when continual pretraining on 80B general tokens, Rho-1 achieves 6.8% average enhancement across 15 diverse tasks, increasing both efficiency and performance of the language model pre-training.

1 Introduction

Document-level filtering does not remove all noisy or misaligned tokens, while uniform next-token loss can waste computation on nonessential tokens. RHO-1 uses reference-model token scores to focus training, based on evidence that tokens follow different loss trajectories.

1 Introduction

  • Token losses follow distinct patterns: many tokens are already learned, some remain difficult or fluctuate, and only a subset shows notable loss reduction.The analysis identifies easy tokens and hard tokens with variable losses that resist convergence, potentially producing ineffective gradient updates.Source: §1 · p. 2
  • SLM scores corpus tokens using a reference model and trains on tokens with high excess loss between the reference and training models.The reference model is trained on high-quality corpora to establish utility scores for the desired distribution.Source: §1 · p. 2
  • SLM is reported to improve token efficiency, downstream performance, and perplexity on benchmarks using selected tokens.The paper explains “Rho” as selective modeling of tokens with higher information density (ρ).Source: §1 · p. 2
  • 15B tokens let RHO-1-7B match DeepSeekMath-7B, which used 500B tokens; RHO-1 1B and 7B models also exceeded CLM baselines by over 16% on GSM8k and MATH.After fine-tuning, RHO-1-1B and 7B scored 40.6% and 51.8% on MATH, respectively.Source: §1 · p. 3
  • 6.8% average improvement across 15 benchmarks followed SLM training of Tinyllama-1B on 80B general tokens, with gains over 10% in code and math.Self-referencing without high-quality reference data yielded average downstream-task improvement of up to 3.3%.Source: §1 · p. 3

2 Selective Language Modeling

Token-level losses differ during pretraining: many tokens are already learned, while others remain difficult or fluctuate. Selective Language Modeling uses a reference model to score tokens and concentrates training on those with high excess loss.

2.1 Not All Tokens Are Equal: Training Dynamics of Token Loss

  • 26% of tokens showed notable loss reduction, while 51% remained consistently low-loss, 11% persistently high-loss, and 12% increased in loss.The analysis continued Tinyllama-1B pretraining on 15B OpenWebMath tokens, evaluating about 320,000 validation tokens at checkpoints every 1B tokens.Source: §2.1 · p. 3
  • Many consistently low- and high-loss tokens fluctuated with high variance, and analysis found many such tokens were noisy.The authors suggest selecting appropriate tokens could stabilize training trajectories and improve data efficiency.Source: §2.1 · p. 3

2.2 Selective Language Modeling

  • SLM trains a reference model on curated high-quality data, scores pretraining-corpus tokens, then selectively trains on tokens with high excess loss.The reference model's token losses provide the basis for selecting training tokens.Source: §2.2 · pp. 3–4, Figure 4 · p. 4
  • Excess loss is the training-model token loss minus the reference-model loss, and SLM focuses on tokens with high excess loss.The token-selection ratio k% sets the proportion selected for cross-entropy training; the loss is computed over the top k% by score.Source: §2.2 · p. 4

3 Experiments

Experiments evaluate SLM through math and general-domain continual pretraining, self-reference tests, and analyses of token selection and downstream performance.

3.2 Math Pre-training Results

  • 40.6% on MATH for RHO-1-1B and 51.8% for RHO-1-7B after fine-tuning, reported as state-of-the-art results.Fine-tuning used the 69k-sample ToRA corpus.Source: Table 2 · p. 6, §3.2 · p. 7

3.1 Experimental Setup

  • Experiments used OpenWebMath for math pretraining and an 80B-token mixed general corpus, with curated math and general datasets to train reference models.The main experiments initialized the reference and continual-pretraining models from the same base model.Source: §3.1 · pp. 5–6

3.2 Math Pre-training Results

  • 16.5% on 1B models and 10.4% on 7B models: RHO-1-Math’s average few-shot accuracy gains over direct continued pretraining.RHO-1-7B achieved results comparable to DeepSeekMath-7B with 15B rather than 500B pretraining tokens.Source: §3.2 · pp. 6–7

3.3 General Pre-training Results

  • 6.8% average enhancement across 15 benchmarks: SLM outperformed direct continual pretraining on 80B general-domain tokens, with code and math gains exceeding 10%.Figure 5 compares CLM-trained Tinyllama-CT with SLM-trained RHO-1.Source: Figure 5 · p. 7, §3.3 · p. 7

3.4 Self-Reference Results

  • The intersection of tokens selected by the two scoring functions delivered +3.3% performance with 40% fewer tokens; OWM-only selection improved same-corpus performance by +2.4%.On PPile, using an OWM-trained reference model yielded a 1.8% improvement with 30% fewer tokens.Source: Table 3 · p. 8, §3.4 · p. 8

3.5 Ablation Study and Analysis

  • Across 4B-token pretraining, selected-token loss reductions aligned with lower downstream loss; selected-token loss correlated with task performance, while unselected tokens had a negative impact.The authors conclude that reducing loss across all tokens is not imperative for improved performance.Source: Figure 6 · p. 8, §3.5 · p. 8
  • About 60% of the original tokens was a suitable selection ratio for the tested 1B-model SLM setup.The ratio was tested while training on 5B tokens.Source: Figure 9 · p. 9, §3.5 · p. 9
  • Most selected OpenWebMath tokens were mathematics-related; later checkpoints selected tokens with higher late-stage and lower early-stage perplexity, with a sample-wise double descent in selected-token perplexity.The authors suggest selection based on excess loss may target tokens most in need at each checkpoint.Source: §3.5 · p. 9

4 Conclusion

The paper proposes Selective Language Modeling to select more suitable tokens for each pretraining stage and reports effectiveness in mathematics and general-domain experiments. The conclusion also calls for further research on token-level pretraining.

4 Conclusion

  • Selective Language Modeling trains RHO-1 on tokens selected as more suitable for the current pretraining stage.The authors frame the method as a response to their analysis of token-level pretraining loss.Source: §4 · p. 9
  • Experiments and analysis in mathematics and general domains demonstrate SLM's effectiveness and emphasize token-level considerations in LLM pretraining.Source: §4 · p. 10

A Author Contributions

  • Zhenghao Lin and Zhibin Gou led much of the method, data, training, evaluation, analysis, and writing work, while other listed collaborators provided project leadership, mentorship, and writing contributions.Source: A Author Contributions · p. 19

B.1 Pretraining Data Optimization

Pretraining-corpus optimization seeks to improve training performance and efficiency by improving the quality and scale of the data mixture.

B.1 Pretraining Data Optimization

  • Pretraining-corpus optimization aims to maximize language-model training performance and efficiency through a better data mixture.Source: §B.1 · p. 19
  • Corpus construction includes collecting data through crawling or synthesis and removing duplicates.Source: §B.1 · p. 19
  • Filtering, selection, data composition, and curriculum are also described as ways to optimize pretraining corpora.Source: §B.1 · p. 19

B.2 Data Selection

Prior data-selection work spans document, sample, and token levels, while strict filtering can harm evaluation loss or induce bias. The paper positions its approach as fine-grained token-level selection for large-scale language-model pretraining.

B.2 Data Selection

B.3 Language Model Training Dynamics

The paper identifies varied token-level training trajectories, including easy tokens and hard tokens that resist convergence, and proposes targeting influential tokens. It contrasts this focus with scaling-law work that studies cross-entropy over all training tokens.

B.3 Language Model Training Dynamics

  • Token trajectories span easy tokens and hard tokens that resist convergence, unlike the interpretation of little perplexity change as already learned [Xia et al., 2022].Source: §B.3 · p. 20
  • The authors propose selective language modeling to target influential tokens and optimize learning based on these training-dynamics patterns.Source: §B.3 · p. 20

B.4 Scaling Laws

C Limitations and Future Work

The authors identify limits in SLM’s demonstrated scale and generalization, while outlining ways to broaden its reference-model, data, and application scope.

Generalizability

  • Math-focused SLM training quickly converges toward the reference model’s domain and raises loss on unselected tokens; no adverse effects from this rise had been observed yet.The authors suggest general text-and-code pretraining may prevent overfitting.Source: C Limitations and Future Work · p. 20

Scalability

  • SLM has been verified only on models up to 7B parameters and datasets under 100B tokens, leaving its scalability to larger settings untested.The authors suggest very large models might develop a similar useful-data compression bias, but say this needs study.Source: C Limitations and Future Work · p. 20

Reference model

  • Token scoring requires a high-quality reference model, which could be a base model trained on high-quality data or a performant open-source model.The authors also suggest proprietary model APIs as a possible source of reference log probabilities, leaving that approach for future work.Source: C Limitations and Future Work · pp. 20–21

Future work

  • Proposed extensions include reweighting tokens, using multiple reference models, and token-level curriculum or iterative strategies; SLM may also extend to supervised fine-tuning and multimodal data.Reweighting is proposed as a possible way to improve robustness.Source: C Limitations and Future Work · p. 21

D Analysis and Visualization of Tokens in Pretraining

The appendix details how token-loss trajectories are categorized, describes anomalous tokens, and specifies the math and general-domain evaluation suites.

D.1 More Details of Four Categories Tokens

  • Tokens are classified by fitted loss change and final loss: ΔL < −0.2 marks H→L, ΔL > 0.2 marks L→H, and near-zero change is split into L→L or H→H by final loss versus the mean.Losses are collected every 1B training tokens, and the endpoint difference of a fitted line estimates loss change.Source: §D.1 · p. 21

D.2 Non-Converging Tokens in Pretrainig

  • Many selected anomalous tokens come from chaotic corpora containing custom symbols, gibberish, timetables, or bibliographic references, though fluctuations in ordinary text may be normal.The authors suggest mitigating losses from the chaotic examples could make training more stable and efficient.Source: §D.2 · p. 22

E.1 Math Evalution

  • The math evaluation covers ten benchmarks spanning difficulty levels, mathematical domains, and multiple-choice and open-ended questions.The named benchmarks include GSM8k, MATH, GSM-Hard, SVAMP, ASDIV, MAWPS, TabMWP, MathQA, MMLU-STEM, and SAT.Source: §E.1 · p. 22

E.2 General Evalution

  • General-domain evaluation uses lm-evaluation-harness for language benchmarks, with separate pipelines for HumanEval, TydiQA, and MBPP.The reported measures include Pass@1 and Pass@10 for HumanEval and MBPP, and F1 for TydiQA.Source: §E.2 · p. 22

F Relate the Selected Tokens’ Loss to Downstream Task Performance

The section relates selected and unselected token losses to downstream accuracy, then provides examples and checkpoint visualizations of SLM token selection.

F Relate the Selected Tokens’ Loss to Downstream Task Performance

  • Fitted trends relate selected-token loss positively and unselected-token loss negatively to average downstream accuracy, measured using MATH and GSM8K.The authors infer that selected tokens can improve performance and unselected tokens may be detrimental; the passage describes a fitted relationship.Source: F Relate the Selected Tokens’ Loss to… · p. 22

G.1 Token Selected Examples

  • Figure 13 marks in blue the tokens selected by SLM in example sequences.The passage describes the examples but reports no specific selection outcome.Source: §G.1 · p. 23

G.2 Dynamic Token Selected

  • Figure 14 compares selection tendencies at 0%, 33%, 66%, and 100% of training, with colors encoding preference from high to low.The color sequence runs from deep blue and blue through black and orange to dark orange.Source: §G.2 · p. 23

H Self-Reference Setting

The section describes reference-loss and information-entropy scoring as ways to select tokens, and identifies a weak-to-strong math benchmark results table.

H Self-Reference Setting

  • The reference-loss score selects tokens based on reference-model loss, with higher reference loss corresponding to lower expected selection.The entropy score instead ranks tokens using reference-model uncertainty in context.Source: H Self-Reference Setting · p. 23
  • Table 5 is labeled as reporting weak-to-strong generalization results on math benchmarks.The supplied caption gives no result values or comparison outcome.Source: Table 5 · p. 24
  • Higher reference-model entropy indicates greater token uncertainty; the authors do not expect the pretraining model to learn tokens it remains uncertain about.The passage notes that Table 4 provides additional self-reference results for different selection ratios and score-function combinations.Source: H Self-Reference Setting · p. 24

I Weak-to-Strong Generalization

The section tests whether a smaller reference model can guide pretraining of a larger model, reporting benefits from cross-size token selection. It also identifies differing vocabularies as a setting where token alignment may be considered.

I Weak-to-Strong Generalization

Loading 2404.07965v4…