Explain the following research paper like I am a layperson: 

""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ϕ
χXiamen University πTsinghua University ϕMicrosoft
https://aka.ms/rho
Abstract
Previous language model pre-training methods have uniformly applied a next-token
prediction loss to all training tokens. Challenging this norm, we posit that “Not all
tokens in a corpus are equally important for language model training”. Our initial
analysis delves into 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 excess
loss. 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 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.
0 5 10 15
Tokens (B)
5
10
15
20
Math Acc (%)
16.3% better
10x faster
Avg Few-shot Acc on 1B LMs
DeepSeekMath-1B (150B Tokens)
Rho-1-1B
Baseline
0 5 10 15
Tokens (B)
20
25
30
35
40
45
50
16.4% better
5x faster
Avg Few-shot Acc on 7B LMs
DeepSeekMath-7B (500B Tokens)
Rho-1-7B
Baseline
Figure 1: We continual pretrain 1B and 7B LMs with 15B OpenWebMath tokens. RHO-1 is trained
with our proposed Selective Language Modeling (SLM), while baselines are trained using causal
language modeling. SLM improves average few-shot accuracy on GSM8k and MATH by over 16%,
achieving the baseline performance 5-10x faster.
⋆Equal contribution. See author contributions for details. Work done during their internships at Microsoft
Research Asia. B: <PRESIDIO_ANONYMIZED_EMAIL_ADDRESS>; <PRESIDIO_ANONYMIZED_EMAIL_ADDRESS>
⋄Correspondence authors.
Preprint. Under review.
arXiv:2404.07965v1 [cs.CL] 11 Apr 2024
The farm has 35 hens <Apr12 1:24> and 12 pigs. ##davidjl123 says totaling 47 animals.
�! �" �# �$ �% �& �' EOS
Desired Tokens Undesired Tokens
Causal Language Modeling
✓ Keep loss ✘ Remove loss
�( �! �" �# �$ �% �& �'
Selective Language Modeling
�( �! �" �# �$ �% �& �'
✓ ✘ ✓ ✘ ✘ ✓ ✓
Noisy Pretraining Corpus
�! �" �# �$ �% �& �' EOS
✓
Figure 2: Upper: Even an extensively filtered pretraining corpus contains token-level noise. Left:
Previous Causal Language Modeling (CLM) trains on all tokens. Right: Our proposed Selective
Language Modeling (SLM) selectively applies loss on those useful and clean tokens.
1 Introduction
Scaling up model parameters and dataset size has consistently elevated the next-token prediction
accuracy in large language models, yielding significant advancements in artificial intelligence [Kaplan
et al., 2020, Brown et al., 2020, OpenAI, 2023, Team et al., 2023]. However, training on all available
data is not always optimal or feasible. As a result, the practice of data filtering has become crucial,
using various heuristics and classifiers [Brown et al., 2020, Wenzek et al., 2019] to select training
documents. These techniques significantly improve data quality and boost model performance.
However, despite thorough document-level filtering, high-quality datasets still contain many noisy
tokens that can negatively affect training, as illustrated in Figure 2 (Upper). Removing such tokens
might alter the text’s meaning, while overly strict filtering could exclude useful data [Welbl et al., 2021,
Muennighoff et al., 2024] and lead to biases [Dodge et al., 2021, Longpre et al., 2023]. Furthermore,
research indicates that the distribution of web data does not inherently align with the ideal distribution
for downstream applications [Tay et al., 2022, Wettig et al., 2023]. For example, common corpus
at the token level may include undesirable content like hallucinations or highly ambiguous tokens
that are hard to predict. Applying the same loss to all tokens can result in wasted computation on
non-beneficial tokens, possibly limiting LLM’s potential to merely mediocre intelligence.
To explore how language models learn at the token level, we initially examined training dynamics,
particularly how the token-level loss evolves during usual pretraining. In §2.1, we evaluated the
model’s token perplexity at different checkpoints and categorized tokens into different types. Our
findings reveal that significant loss reduction is limited to a select group of tokens during training.
Many tokens are “easy tokens” that are already learned, and some are “hard tokens” that exhibit
variable losses and resist convergence. These tokens can lead to numerous ineffective gradient
updates.
Based on these analyses, we introduce RHO-1 models trained with a novel Selective Language
Modeling (SLM) objective. As shown in Figure 2 (Right), this approach inputs the full sequence into
the model and selectively removes the loss of undesired tokens. The detailed pipeline is depicted
in Figure 4: First, SLM trains a reference language model on high-quality corpora. This model
establishes utility metrics to score tokens according to the desired distribution, naturally filtering
out unclean and irrelevant tokens. Second, SLM uses the reference model to score each token in a
corpus using its loss (§2.2). Finally, we train a language model only on those tokens that exhibit a
high excess loss between the reference and the training model, selectively learning the tokens that
best benefit downstream applications (§2.2).
We show through comprehensive experiments that SLM significantly enhances token efficiency during
pretraining and improves performance on downstream tasks. Furthermore, our findings indicate that
SLM effectively identifies tokens relevant to the target distribution, resulting in improved perplexity
scores on benchmarks for models trained with the selected tokens. §3.2 shows the effectiveness of
2
0 5 10 15
Trained Tokens(B)
0
1
2
3
4
Loss
(a) Loss for different token types
H H (11%)
L H (12%)
H L (26%)
L L (51%)
0 5 10 15
Trained Tokens(B)
0.0
0.1
0.2
0.3
Loss
(b) Example L L tokens
L L Token 1
L L Token 2
L L Token 3
0 5 10 15
Trained Tokens(B)
1.5
2.0
2.5
3.0
3.5
Loss
(c) Example H H tokens
H H Token 1
H H Token 2
H H Token 3
Figure 3: The loss of four categories of tokens during pretraining. (a) shows the loss of H→H,
L→H, H→L, and L→L tokens during pretraining. (b) and (c) show three cases of fluctuating tokens’
loss in L→L and H→H during pretraining, respectively.
SLM on math continual pretraining: both 1B and 7B RHO-1 outperform CLM-trained baselines
by over 16% on the GSM8k and MATH datasets. SLM reaches baseline accuracy up to 10x
faster, as shown in Figure 1. Remarkably, RHO-1-7B matches the state-of-the-art performance of
DeepSeekMath-7B using only 15B tokens, compared to the 500B tokens required by DeepSeekMath.
Upon fine-tuning, RHO-1-1B and 7B achieve 40.6% and 51.8% on MATH, respectively. Notably,
RHO-1-1B is the first 1B LM to exceed 40% accuracy, nearing the early GPT-4’s CoT performance
of 42.5%. §3.3 confirms the efficacy of SLM in general pretraining: Training Tinyllama-1B on 80B
tokens with SLM improves 6.8% on average across 15 benchmarks, with gains over 10% in code and
math tasks.
2 Selective Language Modeling
2.1 Not All Tokens Are Equal: Training Dynamics of Token Loss
Our investigation begins with a critical look at how individual tokens’ losses evolve during standard
pre-training. We continue pre-training Tinyllama-1B with 15B tokens from OpenWebMath, saving
checkpoints after every 1B tokens. We then evaluate token-level loss at these intervals using the
validation set of approximately 320,000 tokens. Figure 3(a) reveals a striking pattern: tokens
fall into four categories based on their loss trajectory—persistent high loss (H→H), increasing
loss (L→H), decreasing loss (H→L), and consistent low loss (L→L). For further details on these
categories, see §B.1. Our analysis uncovers that a mere 26% of tokens show a notable loss reduction
(H→L), while the majority (51%) remain in the L→L category, indicating they have already been
learned. Interestingly, 11% of the tokens are persistently challenging (H→H), likely due to high
aleatoric uncertainty [Hüllermeier and Waegeman, 2021]. Additionally, 12% of tokens experience an
unexpected loss increase (L→H) during training.
Our second observation is that a significant number of token losses exhibit persistent fluctuations,
and resist convergence. The loss of many L→L and H→H tokens, as depicted in Figure 3 (b) and (c),
show high variance during training. In §B.2, we visualize and analyze the content of these tokens and
find that many of them are noisy, which is consistent with our hypothesis.
Consequently, we learn that the loss associated with each token during training does not decrease
smoothly like the overall loss; instead, there is a complex training dynamic among different tokens.
If we can select the appropriate tokens for the model to focus on during training, we may be able to
stabilize the trajectory of the model’s training and enhance its efficiency.
2.2 Selective Language Modeling
Overview Inspired by the practice of reference model in document-level filtering, we propose a
simple pipeline of token-level data selection, termed “Selective Language Modeling (SLM)”. Our
method comprises three steps, as depicted in Figure 4. We begin by training a reference model on a
curated, high-quality dataset. This model then assesses the loss of each token within the pretraining
corpus. In the final phase, we train the language model selectively, focusing on tokens with high
excess loss between the training and reference model. The intuition is that tokens with high excess
3
Step 1
Train a
reference
model on
high-quality
text.
Reference
Model
High-quality
Corpus
Pretraining
Corpus
Step 2
Calculate
each token’s
ppl in the
pretraining
corpus. Language
Model
Step 3
Train an
LLM with
loss focused
on high-score
tokens.
Figure 4: The pipeline of Selective Language Modeling (SLM). SLM optimizes language model
performance by concentrating on valuable, clean tokens during pre-training. It involves three steps:
(Step 1) Initially, train a reference model on high-quality data. (Step 2) Then, score each token’s
loss in a corpus using the reference model. (Step 3) Finally, train the language model selectively on
tokens that show higher excess loss compared to the reference loss.
loss are more learnable and better aligned with the desired distribution, naturally excluding tokens
that are either irrelevant or of low quality. Below, we provide a detailed description of each step.
Reference Modeling We begin by curating a high-quality dataset that reflects the desired data
distribution. We train a reference model (RM) using standard cross-entropy loss on the curated data.
The resulting RM is then used to assess the token loss within a larger pretraining corpus. We compute
the reference loss (Lref) of a token xi based on the probability that the RM assigns to this token. The
calculation is formalized as follows:
Lref(xi) = − log P(xi
|x< i) (1)
By evaluating Lref for each token, we establish the reference loss for selective pretraining, allowing
us to focus on the most influential tokens in language modeling.
Selective Pretraining Note that causal language modeling (CLM) employs the cross-entropy loss:
LCLM(θ) = −
1
N
X
N
i=1
log P(xi
|x<i; θ) (2)
Here, LCLM(θ) represents the loss function parameterized by model θ. N is the length of the sequence,
xi
is the i-th token in the sequence, and x<i represents all tokens before the i-th token. In contrast,
Selective Language Modeling (SLM) trains the language model with a focus on tokens that exhibit
a high excess loss when compared to the reference model. The excess loss (L∆) for a token xi
is
defined as the difference between the current training model loss (Lθ) and the reference loss:
L∆(xi) = Lθ(xi) − Lref(xi) (3)
We introduce a token selection ratio k%, which determines the proportion of tokens to be included
based on their excess loss. The cross-entropy loss for the selected tokens is computed as follows:
LSLM(θ) = −
1
N ∗ k%
X
N
i=1
Ik%(xi) · log P(xi
|x<i; θ) (4)
Here, N ∗ k% defines the number of tokens that fall within the top k% of excess loss. The indicator
function Ik%(xi) is defined as:
Ik%(xi) = 
1 if xi
is in the top k% of L∆
0 otherwise (5)
This ensures that the loss is applied only to the tokens that are deemed most beneficial for the language
model to learn from. In practice, token selection can be implemented by ranking the tokens in a
batch according to their excess loss and using only the top k% of tokens for training. This process
4
Table 1: Few-shot CoT reasoning results of math pretraining. All models are tested with few-shot
prompting. Previous best results are highlighted in blue, while our best results are in purple.
∗Only
unique math-related tokens are calculated. For RHO-1, we calculate only the selected tokens that
are used for training. †We use OpenAI’s MATH subset [Lightman et al., 2023] for evaluation, since
some original test samples have been used in public training sets such as PRM800k. ‡The SAT only
has 32 four-choice problems, so we average our results over the last three checkpoints, if available.
Model |θ| Data Uniq.
Toks∗
Train
Toks GSM8K MATH† SVAMP ASDiv MAWPS TAB MQA MMLU
STEM SAT‡ AVG
1-2B Base Models
Tinyllama 1.1B - - - 2.9 3.2 11.0 18.1 20.4 12.5 14.6 16.1 21.9 13.4
Phi-1.5 1.3B - - - 32.4 4.2 43.4 53.1 66.2 24.4 14.3 21.8 18.8 31.0
Qwen1.5 1.8B - - - 36.1 6.8 48.5 63.6 79.0 29.2 25.1 31.3 40.6 40.0
Gemma 2.0B - - - 18.8 11.4 38.0 56.6 72.5 36.9 26.8 34.4 50.0 38.4
DeepSeekLLM 1.3B OWM 14B 150B 11.5 8.9 - - - - - 29.6 31.3 -
DeepSeekMath 1.3B - 120B 150B 23.8 13.6 - - - - - 33.1 56.3 -
Continual Pretraining on Tinyllama-1B
Tinyllama-CT 1.1B OWM 14B 15B 6.4 2.4 21.7 36.7 47.7 17.9 13.9 23.0 25.0 21.6
RHO-1-Math 1.1B OWM 14B 9B 29.8 14.0 49.2 61.4 79.8 25.8 30.4 24.7 28.1 38.1
∆ -40% +23.4 +11.6 +27.5 +24.7 +32.1 +7.9 +16.5 +1.7 +3.1 +16.5
RHO-1-Math 1.1B OWM 14B 30B 36.2 15.6 52.1 67.0 83.9 29.0 32.5 23.3 28.1 40.9
≥ 7B Base Models
LLaMA-2 7B - - 14.0 3.6 39.5 51.7 63.5 30.9 12.4 32.7 34.4 31.4
Mistral 7B - - 41.2 11.6 64.7 68.5 87.5 52.9 33.0 49.5 59.4 52.0
Minerva 8B - 39B 164B 16.2 14.1 - - - - - 35.6 - -
Minerva 62B - 39B 109B 52.4 27.6 - - - - - 53.9 - -
Minerva 540B - 39B 26B 58.8 33.6 - - - - - 63.9 - -
LLemma 7B PPile 55B 200B 38.8 17.2 56.1 69.1 82.4 48.7 41.0 45.4 59.4 50.9
LLemma 34B PPile 55B 50B 54.2 23.0 67.9 75.7 90.1 57.0 49.8 54.7 68.8 60.1
Intern-Math 7B - 31B 125B 41.8 14.4 61.6 66.8 83.7 50.0 57.3 24.8 37.5 48.7
Intern-Math 20B - 31B 125B 65.4 30.0 75.7 79.3 94.0 50.9 38.5 53.1 71.9 62.1
DeepSeekMath 7B - 120B 500B 64.1 34.2 74.0 83.9 92.4 63.4 62.4 56.4 84.4 68.4
Continual Pretraining on Mistral-7B
Mistral-CT 7B OWM 14B 15B 42.9 22.2 68.6 71.0 86.1 45.1 47.7 52.6 65.6 55.8
RHO-1-Math 7B OWM 14B 10.5B 66.9 31.0 77.8 79.0 93.9 49.9 58.7 54.6 84.4 66.2
∆ -30% +24.0 +8.8 +9.2 +8.0 +7.8 +4.8 +11.0 +2.0 +18.8 +10.4
eliminates the loss for undesired tokens without incurring additional costs during pretraining, making
our approach both efficient and easily integrated.
3 Experiments
We continually pretrained models in both mathematical and general domain and designed ablation
and analysis experiments to understand the effectiveness of SLM.
3.1 Experimental Setup
Reference Model Training To train our mathematical reference model, we gathered a dataset of
0.5B high-quality, math-related tokens. This dataset is a blend of synthetic data from GPT [Yu et al.,
2024, Huang et al., 2024] and manually curated data [Yue et al., 2024, Ni et al., 2024]. For the general
reference model, we compiled a corpus of 1.9B tokens from open-source datasets, such as Tulu-v2
[Ivison et al., 2023] and OpenHermes-2.5 [Teknium, 2023]. We trained the reference models for 3
epochs. The maximum learning rate was set at 5e-5 for 1B models and 1e-5 for 7B models, applying
a cosine decay schedule. We set the maximum sequence lengths to 2048 for 1B models and 4096 for
7B models, packing multiple samples into these lengths for model input. In all main experiments, we
initialized the continual pretraining model and the reference model with the same base model.
Pretraining Corpus For mathematical reasoning, we utilize the OpenWebMath (OWM) dataset
[Paster et al., 2023], which comprises approximately 14B tokens sourced from math-related web
pages in the Common Crawl. In the general domain, we combine the SlimPajama [Daria et al., 2023]
5
Table 2: Tool-integrated reasoning results of math pretraining.
Model Size Tools SFT Data GSM8k MATH SVAMP ASDiv MAWPS TAB GSM-H AVG
Used for SFT? ✓ ✓ ✗ ✗ ✗ ✗ ✗
Previous Models
GPT4-0314 - ✗ - 92.0 42.5 93.1 91.3 97.6 67.1 64.7 78.3
GPT4-0314 (PAL) - ✓ - 94.2 51.8 94.8 92.6 97.7 95.9 77.6 86.4
MAmmoTH 70B ✓ MI-260k 76.9 41.8 82.4 - - - - -
ToRA 7B ✓ ToRA-69k 68.8 40.1 68.2 73.9 88.8 42.4 54.6 62.4
ToRA 70B ✓ ToRA-69k 84.3 49.7 82.7 86.8 93.8 74.0 67.2 76.9
DeepSeekMath 7B ✓ ToRA-69k 79.8 52.0 80.1 87.1 93.8 85.8 63.1 77.4
Our Pretrained Models
TinyLlama-CT 1B ✓ ToRA-69k 51.4 38.4 53.4 66.7 81.7 20.5 42.8 50.7
RHO-1-Math 1B ✓ ToRA-69k 59.4 40.6 60.7 74.2 88.6 26.7 48.1 56.9
∆ +8.0 +2.2 +7.3 +7.5 +6.9 +6.2 +5.3 +6.2
Mistral-CT 7B ✓ ToRA-69k 77.5 48.4 76.9 83.8 93.4 67.5 60.4 72.6
RHO-1-Math 7B ✓ ToRA-69k 81.3 51.8 80.8 85.5 94.5 70.1 63.1 75.3
∆ +3.8 +3.4 +3.9 +1.7 +1.1 +2.6 +2.7 +2.7
and StarCoderData [Li et al., 2023a] (both part of the Tinyllama corpus) with OpenWebMath, training
on a total of 80 billion tokens with a mix ratio of 6:3:1.
Pretraining Setting For math pretraining, we continue pretraining on the Tinyllama-1.1B model
[Zhang et al., 2024] and the Mistral-7B model [Jiang et al., 2023] with learning rates of 8e-5 and
2e-5, respectively. For general domain, we set the learning rate for Tinyllama-1.1B model to 1e-4.
The batch size is uniformly set to 1M tokens for both domains. Regarding the token selection ratio,
we use 60% for the Tinyllama-1.1B model and 70% for the Mistral-7B model.
Baseline Setting We use models that have been continually pretrained (Tinyllama-CT and MistralCT) through regular causal language modeling as baselines. Moreover, we compare RHO-1 with
well-known and top-performing baselines, including Gemma [Team et al., 2024], Qwen1.5 [Bai et al.,
2023], Phi-1.5 [Li et al., 2023b], DeepSeekLLM [DeepSeek-AI, 2024], DeepSeekMath [Shao et al.,
2024a], CodeLlama [Roziere et al., 2023], Mistral [Jiang et al., 2023], Minerva [Lewkowycz et al.,
2022], Tinyllama [Zhang et al., 2024], LLemma [Azerbayev et al., 2023], and InternLM2-Math [Ying
et al., 2024]. For fine-tuning results, we also compare with previous best models MAmmoTH[Yue
et al., 2024] and ToRA[Gou et al., 2024].
Evaluation Setup To comprehensively evaluate pretrained models, we compare their few-shot
capabilities and fine-tuning performance across a variety of tasks. We adopt the lm-eval-harness3
for
general tasks, and math-eval-harness4
for math tasks. We use vllm (v0.3.2) [Kwon et al., 2023] to
speed up inference. Further details on our evaluation can be found in Appendix C.
3.2 Math Pre-training Results
Few-shot CoT Reasoning Results We evalute base models prompting with few-shot chain-ofthought (CoT) [Wei et al., 2022a] examples following previous works [Lewkowycz et al., 2022,
Azerbayev et al., 2023, Shao et al., 2024b]. As results shown in Table 1, in comparison to continue
pretraining directly, RHO-1-Math has achieved the average few-shot accuracy improvement of
16.5% on 1B models and 10.4% on 7B models. Furthermore, after training for multiple epochs on
OpenWebMath, we find that RHO-1 could further increase the average few-shot accuracy to 40.9%.
Compared to DeepSeekMath-7B, which pretrained on 500 billion math-related tokens, RHO-1-7B
pretrained on only 15 billion tokens (selecting 10.5 billion tokens) achieved comparable results,
demonstrating the efficiency of our approach.
Tool-Integrated Reasoning Results We fine-tune RHO-1 and baseline models on 69k ToRA corpus
[Gou et al., 2024], consisting of 16k GPT-4-generated trajectories in a tool-integrated reasoning
format, and 53k answer-augmented samples using LLaMA. As presented in Table 2, RHO-1-1B
3
https://github.com/EleutherAI/lm-evaluation-harness
4
https://github.com/ZubinGou/math-evaluation-harness
6
and RHO-1-7B achieved a state-of-the-art 40.6% and 51.8% on MATH dataset, respectively. On
some unseen tasks (e.g., TabMWP and GSM-Hard), RHO-1 also demonstrates a certain degree of
generalizability, with an average few-shot accuracy improvement of 6.2% on the RHO-1-Math-1B
and 2.7% on RHO-1-Math-7B.
MMLU
BBH
MATH
GSM8k
MBPP(p@1)
MBPP(p@10)
HumEval(p@1)
HumEval(p@10)
0
10
20
30
40
Metrics (%)
+11.3
+3.9
+5.0
+28.2
+6.5
+7.8
+6.9
+10.6
Performance of General Pretrained Base Model
AGIEval
ARC-C
ARC-E
BoolQ
PIQA
HellaSwag
WinoGrande OBQA
TydiQA
20
30
40
50
60
70
80
Metrics (%)
+1.1 +5.0
+8.6
+11.3 +0.9
+1.4 +0.2
+3.4
+8.9
Tinyllama
Tinyllama-CT
Rho-1-1B
Figure 5: General pretraining results. We continual pretraining Tinyllama-1B on 80G general
tokens. Tinyllama-CT is etrained with CLM, while RHO-1 is trained with our proposed SLM.
3.3 General Pre-training Results
We confirm the efficacy of the SLM in general pretraining by continual training Tinyllama-1.1B
on 80 billion tokens. The results depicted in Figure 5 indicate that although Tinyllama has already
undergone extensive training on the majority of these tokens, the application of SLM yields an
average enhancement of 6.8% across 15 benchmarks compared to direct continual pretraining. The
improvements were especially pronounced in code and math tasks, exceeding 10%.
3.4 Analysis
Selected Token Loss Aligns Better with Downstream Performance We use the reference model
to filter tokens and explore the changes in validation loss after training on all/selected tokens, while
observing their relationship with downstream loss. As shown in Figure 6, we pretrain approximately
4B tokens and displayed the variation curves of loss on different pretraining method and validation
sets during the pretraining process. We can observe that on the tokens selected by the reference
model, the decrease in average loss of the RHO-1 is more significant compared to regular pretraining.
On the contrary, on the unselected tokens, the decrease in average loss of the regular pretraining is
more significant. If we associate fig (a), fig(b) with fig(c), it is not difficult to find that the model
trained on the selected tokens has a more significant decrease in downstream loss, while ordinary
pretraining, although reducing the average loss of all tokens during the training phase, is difficult
to have a significant decrease in downstream loss. Therefore, we expect that selecting tokens for
pretraining is more efficient.
Moreover, We relate the selected tokens’ loss to its downstream task performance via a power law
in Figure 7, which is similar to a concurrent study [Gadre et al., 2024]. Observing the curve fitted
from the data points in the graph, the average loss of tokens selected by the SLM shows a positive
correlation with the performance of downstream tasks, whereas the average loss of tokens not selected
exhibits a negative correlation with downstream task performance. Therefore, it is not necessary for
the all tokens’ loss to decrease to benefit the model’s ultimate performance. See Appendix D for
more details.
7
0 1 2 3 4
Tokens (B)
0.98
1.02
1.06
1.10
Loss
(a) Selected Token Loss
Baseline
Rho-1
0 1 2 3 4
Tokens (B)
0.85
0.90
0.95
1.00
1.05
1.10
Loss
(b) Downstream Token Loss
Baseline
Rho-1
0 1 2 3 4
Tokens (B)
2.5
2.9
3.3
3.7
4.1
Loss
(c) Unselected Token Loss
Baseline
Rho-1
Figure 6: The dynamics of pretraining loss and downstream loss. (a) and (c) represent the loss
of tokens selected/unselected by SLM during pretraining in both SLM and CLM methods, while
(b) represents the loss of the SLM and CLM methods on downstream corpora. We tested the above
results through the process of pretraining with a total of 4 billion tokens.
0.98 0.95 0.92 0.89 0.86
Loss
5.0
7.5
10.0
12.5
15.0
17.5
20.0
22.5
Accuracy(%)
(a) Accuracy vs. Selected Tokens' Loss
Selected Tokens at 2B
Selected Tokens at 5B
Selected Tokens at 8B
Selected Tokens at 11B
Selected Tokens at 14B
3.66 3.62 3.58 3.54 3.50
Loss
5.0
7.5
10.0
12.5
15.0
17.5
20.0
22.5
Accuracy(%)
(b) Accuracy vs. Unselected Tokens' Loss
Unselected Tokens at 2B
Unselected Tokens at 5B
Unselected Tokens at 8B
Unselected Tokens at 11B
Unselected Tokens at 14B
Figure 7: The relationship between the selected tokens / unselected tokens loss in SLM and
downstream task performance. The y-axis represents the average few-shot accuracy on GSM8k and
MATH. The x-axis represents the average loss on selected tokens / unselected tokens at corresponding
checkpoint(2B, 5B, 8B, 11B, and 14B).
What Tokens are Selected with SLM? We aim to analyze the tokens selected by the SLM method
in pretraining to further explore its working mechanism. To this end, we visualize the token selection
process during the training of RHO-1 using the OpenWebMath. In §E.1, we have highlighted in blue
the tokens that were retained during actual pretraining. We observe that the majority of tokens chosen
by the SLM method are closely related to mathematics, effectively training the model on the parts of
the original corpus that are pertinent to mathematical content.
Furthermore, we investigated the differences in token filtering across various checkpoints during the
training process and tested the perplexity of these tokens on different checkpoints. As illustrated
in Figure 8, we found that the tokens selected by later checkpoints tend to have higher perplexity
towards the later stages of training and lower perplexity in the earlier stages. This may suggest that
the model first optimizes tokens with a larger learnable space, thereby increasing learning efficiency.
Moreover, we noticed a sample-wise “double descent” [Nakkiran et al., 2021] on the loss of selected
tokens, where the select token’s perplexity initially increases before decreases. This might be an
effect of selecting tokens based on excess loss, targeting those most in need at each checkpoint.
Effect of Token Select Ratio We investigate the impact of token selecting ratios of the SLM.
Generally, the selecting ratio is defined by heuristic rules, similar to the approach previously employed
in the training of Masked Language Models (MLMs) [Devlin et al., 2019, Liu et al., 2019]. As shown
in Figure 9, the selected tokens is suitable for accounting for about 60% of the original tokens.
8
2 5 8 11 14
Tokens(B)
2.35
2.45
2.55
2.65
PPL
PPL of Tokens Selected by Different CKPT
Selected Tokens at 2B
Selected Tokens at 5B
Selected Tokens at 8B
Selected Tokens at 11B
Selected Tokens at 14B
Figure 8: The PPL of tokens selected by different checkpoint. We test the PPL of the tokens
selected at 2B, 5B, 8B, 11B, and 14B.
40 50 60 70 80 90 100
Token Select Ratio (%)
0
5
10
15
20
Accuracy (%)
Accuracy vs. Select Ratio
GSM8K
Math
Figure 9: Effect of token select ratio. We train
1B LM with SLM objective on 5B tokens.
Table 3: Weak-to-Strong generization result on math benchmark.
Model GSM8K MATH† SVAMP ASDiv MAWPS TAB MQA MMLU
STEM SAT‡ AVG
Llama-2-7B-CT 28.4 13.6 50.3 62.8 79.5 37.6 34.1 41.6 43.5 43.5
Llama-2-7B-CT w/ 1B RM 29.8 16.0 55.5 63.7 80.4 37.9 34.3 38.2 43.8 44.4
Weak-to-Strong Generization Apart from the main experiments where we use the same base
model for the reference and continual pretraining, we also investigate if a smaller reference model
can effectively guide the pretraining of a larger model. We use Tinyllma-1.1B as reference model and
continual pretraining Llama-2-7B on math. Results presented in Table 3 indicate that, despite the
considerable gap between the small and large models [Li et al., 2023c], employing the small reference
model to token selection can still yield benefits to the pre-training of the larger model. If reference
and training models have different vocabularies, one can consider performing token alignment [Wan
et al., 2024, Fu et al., 2023], which we leave for future work.
4 Related Works
Pretraining Data Optimization The objective of optimizing pre-training corpora is to maximize
the performance and efficiency of language model training by improving the quality and scale of the
pretrain data mixture. This typically includes data collecting through crawling [Raffel et al., 2020]
or synthesis [Polu and Sutskever, 2020, Gunasekar et al., 2023], de-duplication [Lee et al., 2021,
Kandpal et al., 2022, Tirumala et al., 2023], filtering and selection [Albalak et al., 2024], as well as
data composition [Xie et al., 2024] and curriculum [Chen et al., 2024, MA et al., 2024].
Data Selection Data selection for fine-tuning has been extensively studied, focusing on improving
quality [Li et al., 2023d], diversity [Liu et al., 2024], and distribution matching [Li et al., 2023e,
Xia et al., 2024, Ni et al., 2024]. For pretraining, various lightweight filters are utilized [Albalak
et al., 2024], including heuristic-based (e.g., language and item count filtering), classifier-based
[Brown et al., 2020], and perplexity-based approaches [Wenzek et al., 2019]. The massive public
RedPajama-Data-v2 dataset [Computer, 2023], for example, leverages over 40 quality indicators
for data filtering and reweighting. Nevertheless, strict filtering like blocklist [Raffel et al., 2020]
and Safety API filtering [Welbl et al., 2021], have been found to hurt evaluation loss or induce bias
[Dodge et al., 2021]. To our knowledge, we are the first to explore token-level data selection, aimed
at enhancing data quality and information density at the most fundamental granularity.
Language Model Training Dynamics Investigating the training dynamics of language models is
essential for understanding their behavior throughout the training process. This research includes
studying internal representations [Saphra and Lopez, 2018], the acquisition of linguistic knowledge
[Choshen et al., 2021, Liu et al., 2021], and the phenomenon of grokking [Power et al., 2022].
9
The analysis by Xia et al. [2022] is the most related to ours, which examines token-level training
trajectories in models of varying sizes. Our findings, however, diverge from those of Xia et al. [2022],
who posit that tokens with little change in perplexity are “already learned”. We identify a spectrum of
token patterns, including “easy tokens” and “hard tokens” that resist convergence. Recognizing this,
we propose a method of selective language modeling that targets the influential tokens, optimizing
the learning process.
Scaling Laws Scaling laws guide us in discovering the impact of factors such as parameter count,
data size, and compute on language model performance and behavior. These studies usually focus
on predicable scaling though power law [Kaplan et al., 2020, Hernandez et al., 2021], optimal
resource allocation [Hoffmann et al., 2022], downstream tasks [Wei et al., 2022b, Isik et al., 2024,
Gadre et al., 2024], architectures [Tay et al., 2022], memorization [Tirumala et al., 2022, Carlini
et al., 2022, Henighan et al., 2023, Biderman et al., 2024], and repeating data [Hernandez et al.,
2022, Muennighoff et al., 2024, Xue et al., 2024]. Most scaling laws on model performance study
cross-entory loss on all training tokens, while we focus on the tokens loss of desired distributions.
5 Discussion and Future Work
Generalization In math continual pretraining, as depicted in Figure 6, training exclusively with
SLM leads to quickly convergence to the domain focused by the reference model, accompanied by
a significant rise in the loss of unselected tokens. Although no adverse effects, like biases, have
been observed from the increased loss yet, a general pretraining loss on text and code may prevent
overfitting [Goodhart and Goodhart, 1984], as suggested by Ouyang et al. [2022] and Azerbayev
et al. [2023]. Furthermore, future efforts could broaden the corpus scope of the reference model, and
enlarge the pretraining data size, as exemplified by DeepSpeedMath [Shao et al., 2024a].
Scalability Due to budget constraints, we have only verified the effectiveness of our method on
smaller models (<=7B parameters) and smaller datasets (<100B tokens). Smaller models benefit
significantly from removing the loss of irrelevant tokens and focusing on important ones. However,
it’s possible that very large models trained on extensive corpora may naturally develop this inductive
bias to compress useful data (i.e., compressing everything), although it may sounds inefficient for
now. Therefore, future works should study whether this selective language modeling technique can
scale to very large models and data [Kaplan et al., 2020].
Is training a reference model necessary? To score tokens, we need a high-quality reference
model. This could be a base model trained with a small amount of high-quality data, or a performant
open-source model. In fact, since we only need input logprobs or perplexity from reference model,
we could even utilize more powerful proprietary model APIs. We can input tokens and use the log
probabilities of the input returned by the API as reference scores. We leave this for future works.
How to improve upon SLM? There are many natural extensions of SLM, e.g., reweighting tokens
instead of selecting may improve robustness; using a reference model as a reward model to guide
pretraining with reinforcement learning; adopting multiple reference models to reduce overfitting;
designing token-level curriculum learning and iterative strategies for continuous improvements, etc.
Expanding the Use of SLM SLM may be extended to supervised fine-tuning to address the noise
and distribution mismatches in many SFT datasets. Another potential application is alignment, e.g.,
by training a reference model to emphasize helpfulness, truthfulness, and harmlessness, we may
obtain a base model that is natively aligned during the pretraining stage.
""