Self-Improving Pretraining: using post-trained models to pretrain better models — part 2/5

Cleaned-text code points [15595, 31470). Bibliography omitted here; source.md includes it.

Data. We use the SlimPajama ( SP , Soboleva et al. (2023) ) and RedPajama pretraining datasets ( RP , Weber et al. (2024) ). SP is a derivative of RP , created by applying more aggressive safety and quality filtering to produce a “slimmer” higher-quality dataset. Thus training only on SP can be considered as a baseline where the training only uses safe and high-quality samples. We use RP for training our method in the safety experiments. To ensure fairness in training, policy, judge, and rewriter models were trained and evaluated on non-overlapping subsets of the data.

Judge training. To fine-tune Llama3-8B-Instruct for a judge role, we generate synthetic data from subsets of SP and RP with known rewards (i.e., safe vs. unsafe completions and higher vs. lower quality completions). For the quality task, we create the data by asking a Llama3.3-70B-Instruct ( Dubey et al., 2024 ) model to spoil the original suffix (see Appendix Figure 13 ) extracted from SP . A pair of original and corrupted suffixes is then used to create two samples, by wrapping the pair or its flipped version in the quality judge prompt, given in Figure 3 . For the safety task, we use the same model to filter safe and unsafe suffixes from RP by prompting it (see Figure 3 ) to evaluate the safety of the suffix with 8 random seeds, and only using samples where all 8 judgments are safe or unsafe. We then use the same prompt to wrap suffixes for training. We generate 75,432 training and 4,096 validation samples for the quality task, and 3,192 and 512 for the safety task respectively.

Rewriter training. We similarly fine-tune a rewriter from the Llama3-8B-Instruct model (for safety experiments only). We found this was necessary otherwise Llama would refuse to rewrite unsafe prompts. To generate synthetic data we follow the same filtering procedure as for judge safety task training. 73,080 safe and unsafe suffixes then used with the rewriter prompt template, provided in Appendix Figure 14 .

Policy training. Training on SP is conducted on 983,520 samples. Training on the RP dataset is conducted on 257,154 samples, that were filtered to include particularly unsafe content. Specifically, we applied tag-based filtering to extract unsafe documents. Validation and test data were further filtered to retain unsafe data by GPT-OSS-120B to ensure the prefix contains unsafe content. Filtering details are provided in Appendix B.1 .

#### 1.2.2 Experimental Setup

##### Safety Experimental Setup

Our pipeline involves the following three models: a judge, rewriter, and the policy model. Below we will summarize the setup for the components.

Suffix judge. Recent studies provide strong evidence that LLM judges become more robust and effective when they generate their own Chain-of-Thought (CoT) analyses before producing final judgments ( Zhang et al., 2024a ; Chen et al., 2025c ; Whitehouse et al., 2025 ) . To fine-tune Llama3.1-8B-Instruct to be a safety and quality judge incorporating reasoning, we use GRPO ( Shao et al., 2024 ) as our optimization algorithm. Unlike SFT, GRPO does not require generating high-quality synthetic CoT data, but fully relies on a signal from the final judgment, while incentivizing reasoning traces that result in correct judgments. To reward the judge model during training we rely on labels from synthetically generated data of judgments, rewarding a correctly categorized suffix with 1.0, and 0.0 for mismatching the label.

We run GRPO training on the synthetically generated data, where the judge is simultaneously trained on two tasks: quality and safety. We set the global batch size to 256, with 16 generations per prompt under temperature T = 0.6 T=0.6 and t ​ o ​ p ​ _ ​ p = 0.6 top\_p=0.6 . We train on 64 GPUs for 500 steps with 2.0 ​ e − 07 2.0e-07 constant learning rate. The maximum prompt length is set to 3584 tokens, and the model can generate up to 512 new tokens.

During training we observe that initially the safety task is easier to learn than quality, as the model plateaus at 0.94 average reward score at approximately 100 steps, while the reward for quality keeps growing until the end of the training, see Figure 5 . Manually analyzing judgments, we found that the initial model tends to favor suffixes that feel more complete , rather than those that are more coherent with respect to the context . Training helps to fix this problem.

At inference time to make a judgment we query the model twice, for safety using Figure 3 , and for quality using Figure 3 , and combine the results. For the safety judgment, this is a pointwise score, but for quality this is a pairwise judgment given two candidate responses, which outputs which is better of the two. For the latter during policy training we run all pairwise comparisons amongst candidates in the batch, assigning reward 0 or 1 in each case, and take the average of their rewards to obtain pointwise scores. For each rollout, our judge is prompted to evaluate safety and quality 5 times each with temperature T = 1.0 T=1.0 and t ​ o ​ p ​ _ ​ p = 0.6 top\_p=0.6 .

Suffix rewriter. Similarly, we train Llama3.1-8B-Instruct model with the GRPO algorithm. Our goal is to build a suffix rewriter that leaves safe high quality suffixes unchanged (hence the generative output would typically copy the suffix that is given in the input context), whereas for unsafe suffixes, they should be rewritten to be safe. Hence, to train the rewriter, the reward is assigned with the following method: • If the model was prompted to rewrite a safe suffix, we return reward 1.0 1.0 if the rewritten suffix x ¯ j \bar{x}_{j} is an exact match of the given suffix x j x_{j} , otherwise we reward it with 0.0: R safe = { 1.0 if ​ x ¯ j = x j ​ , 0.0 otherwise . R_{\text{safe}}=\begin{cases}1.0&\text{ if }\bar{x}_{j}=x_{j}\text{\;,}\\ 0.0&\text{ otherwise\;.}\end{cases} (1)

• If the model was prompted to rewrite an unsafe suffix, the rewritten suffix is evaluated with the suffix judge based on quality J qual J_{\text{qual}} and safety J safe J_{\text{safe}} , averaging judgments across 5 5 random seeds: R unsafe = 1 2 ​ ( J qual ​ ( x ¯ j , x j | x 1 , … , j − 1 ) + J safe ​ ( x ¯ j ) ) ​ . R_{\text{unsafe}}=\frac{1}{2}\left(J_{\text{qual}}(\bar{x}_{j},x_{j}|x_{1,\dots,j-1})+J_{\text{safe}}(\bar{x}_{j})\right)\text{ \;.} (2)

To train the suffix rewriter model we use same setup as for the suffix judge. We modify the maximum prompt length to 3968 tokens, and the model generation length to 128 new tokens to match our suffix length. We validate model performance on safe and unsafe subsets. We observe steady improvement on the copy task (exact match reward score on safe suffixes, Figure 6 ), and use the final checkpoint that achieves token overlap percent plateaued at 98 % 98\% , as shown in Figure 7 .

##### Factuality Experimental Setup

Suffix judge. In the factuality training setting, we only consider using a judge, and not a rewriter. For the factuality judge, this is a pointwise judgment given one candidate response, and a reference answer. We use the original suffix from the training data as the reference. We use GPT-OSS-120B with the prompt given in Figure 4 . In subsection A.2 , we conduct a detailed study using different strong post-trained models as the judge, and various prompt designs, comparing their performance.

The suffix judge outputs whether the continuation has no hallucination (reward 1), possible hallucination (reward 0.5) or definite hallucination (reward 0). As in the safety experiments, we combine this reward with an overall quality score of the generation, by adding the quality scoring judge rewards. This is done in the same way as in subsubsection 1.2.2 . GPT-OSS-120B is prompted with temperature T = 1.0 T=1.0 and t ​ o ​ p ​ _ ​ p = 1.0 top\_p=1.0 . We also consider another variant of using a single pivot candidate for pairwise comparisons instead, resulting in K K judgments for each update, rather than ( K 2 ) \binom{K}{2} .

##### Quality Experimental Setup

Suffix judge. In the quality training setting, we also only consider using a judge, and not a rewriter. For the quality judge, this is a pairwise judgment given two candidate responses, which outputs which is better. For this we use GPT-OSS-120B with the prompt given in Figure 3 . We run all pairwise comparisons amongst candidates, assigning reward 0 or 1 in each case, and take the mean of their rewards to obtain pointwise scores.

We also consider two other variants: (1) using the trained model from subsubsection 1.2.2 but only prompted for quality; and (2) using a single pivot candidate for pairwise comparisons instead, resulting in K K judgments for each update, rather than ( K 2 ) \binom{K}{2} .

##### Policy training variants and ablations

We conduct a series of variants and ablations of policy training primarily in the safety pretraining setting. First, we conduct both from scratch and continued pretraining in this setting.

##### Continual pretraining experiments

Self-Improving Pretraining models are trained with online DPO (unless said otherwise in ablations) with the global batch size 256, sampling 16 rollouts per prompt using temperature T = 1.0 T=1.0 and t ​ o ​ p ​ _ ​ p = 1.0 top\_p=1.0 . We train on 64 GPUs for 2000 steps with cosine learning rate l ​ r = 5.0 ​ e − 06 lr=5.0e-06 , min ratio 0.1 0.1 , and 100 100 warmup steps. The maximum sequence length is set to 2048 tokens, and the model generates N = 128 N=128 new tokens for each rollout. For the safety task, the fine-tuned Llama3.1-8B-Instruct judge is used to select DPO pairs from 16 rollouts and the original suffix, while GPT-OSS-120B is used to judge 16 rollouts for the quality and factuality tasks.

##### From-scratch pretraining

To pretrain from scratch, we use a similar setup, but increase the number of training steps to 21,000, increase the learning rate to 5.0 ​ e − 04 5.0e-04 , and the number of warmup steps to 2000 2000 . In these experiments, we only use 1 rollout for training.

##### Ablations

We also ablate various ways of doing the training with different loss functions and candidate generation pools during online training, all compared to next token prediction baselines.

In particular, firstly we compare to: SFT on either (i) rewrites or (ii) (single) rollouts; which do not require a judge during training. For RL training, we use online DPO, which has shown performance comparable to GRPO ( Lanchantin et al., 2025 ) . As mentioned before, DPO is an off-policy algorithm that allows learning from sequences not generated from the current policy, such as the original suffix or rewrites, making it suitable for our approach. First, a baseline simple option is to use the rewrite as the chosen and the current rollout as the rejected in online DPO, which also does not require a judge, inspired by the approach in Chen et al. (2024) .

For our full Self-Improving Pretraining method using a judge, we compare online DPO with reward filtered (RF)-NLL. For RF-NLL we consider two flavors: rollout vs rewrite as candidates to be judged, or rollout vs. original suffix vs rewrite. For online DPO, we consider: (i) suffix vs 1 rollout, (ii) rewrite vs. 1 rollout, (iii) suffix vs. 16 rollouts; and (iv) 16 rollouts only. We also conduct a separate study of the effect of scaling the number of rollouts. For policy model generations during training we use a temperature of 1.

For quality and factuality ablations, we study the effects of (i) a single rollout which does not require a judge during training, (ii) 2, 4, 8, 16 rollouts, (iii) suffix as pivot for 8 rollouts. We also compare using the trained judge from subsubsection 1.2.2 , with GPT-OSS-120B as an online judge in the quality pretraining setting.

#### 1.2.3 Evaluations

We evaluate our models on a broad set of benchmarks, including standard evaluations and additional benchmarks focused on coherence, safety and factuality. For generation tasks, we use GPT-OSS-120B as a judge and judgments across 8 random seeds. For the policy model we use greedy generations.

Generation quality. To evaluate the generation quality we use 1k samples from the test split of SP as data with safe prefixes, and 1k samples from the test split of filtered RP as data with unsafe prefixes. Generation quality is evaluated by comparing a sequence of length N N against baseline generations of Llama Base of the same length. We use GPT-OSS-120B as a suffix judge using the prompt given in Figure 3 . We average judgments across 8 random seeds using a temperature of 0.7. In addition, we measure coherence, particularly in terms of repetition, independently using the prompt given in Figure 16 . Note that the generation quality score (win rate) is hence 50.0 for Llama Base given it is used as the baseline in the pairwise comparison.

Standard Evaluations. We use a set of standard evaluation tasks to measure the pretrained policy model’s general reasoning abilities. In particular, we average performance across the following datasets: BoolQ ( Clark et al., 2019 ) , PIQA ( Bisk et al., 2020 ) , SIQA ( Sap et al., 2019 ) , HellaSwag ( Zellers et al., 2019 ) , ARC easy and challenge ( Clark et al., 2018 ) , OpenBookQA ( Mihaylov et al., 2018 ) , and 5-shot performance on the aggregated MMLU benchmark ( Hendrycks et al., 2020 ) .

Safety. The policy model’s safety is evaluated as a weighted average across five datasets: the RP test split, RealToxicityPrompts ( Gehman et al., 2020 ) , ToxiGen ( Hartvigsen et al., 2022 ) , and the XStest safe and unsafe sets ( Röttger et al., 2024 ) . In each case, safety is evaluated with GPT-OSS-120B as a judge using the prompt given in Figure 3 . We use majority vote over N predictions with a temperature of 1.

Factuality. The policy model’s factuality is evaluated as a weighted average across five datasets: the RP test split, FActScore ( Min et al., 2023 ) , HaluEval ( Li et al., 2023 ) , which are generation tasks, and the TruthfulQA multiple-choice tasks MC1 and MC2 ( Lin et al., 2022 ) . We evaluate on the QA, dialogue, summarization tasks in HaluEval with the provided ground-truth answers as reference. For FActScore, the provided wikipedia text is used as ground-truth reference for the GPT judge. For the RP test split, FActScore, and HaluEval, the evaluation is done with the corresponding judge prompts given in Figure 4 , Figure 17 , Figure 18 , respectively. We again use GPT-OSS-120B as a judge, using a temperature of 0.7.

#### 1.2.4 Results

##### Main results

Table 1 summarizes our main results in the continued pretraining setting when optimizing for quality, factuality and safety. We find that all three objectives significantly improve over the initial and continually pretrained baselines in several metrics. Self-Improving Pretraining provides superior generation quality over standard (SlimPajama test set) prefixes, and higher scores on standard pretraining evaluations in all three cases. A breakdown of the standard evaluations can be found in Table 3 .

When optimizing for quality, we see the largest gains in generation quality on standard prefixes, with a win rate of 86.3% over the baseline generations, and a 87.9% win rate in terms of coherence.

When optimizing for factuality, we also see significant gains in quality (84.0% win rate), and more importantly, an improvement in factuality evaluations from 42.3 to 57.6. The breakdown in to individual factuality tasks can be found in Table 4 , where we observe wins in every individual benchmark tested.

When optimizing for safety, we also see significant gains in quality for unsafe prefixes (77.7% win rate), as well as significant improvements in safety evaluations with an average increase from 76.9 to 91.1. The breakdown into individual safety tasks is given in Table 5 . Again, we observe wins in most individual benchmarks tested.