2.2 Self-Improving pretraining using post-trained models

Our self-improvement framework assumes we already have access to a fully trained (i.e., first pre- and then post-trained) model. This model has effectively absorbed information from across the entire pretrain and post-train datasets already – and this expertise can now be brought to bear on individual examples in the pretraining datasets to train a new model using an effectively superior training signal than the one from which it was trained itself.

We consider using this fixed teacher model in two ways: as a rewriter and as a judge.

Suffix Rewriter Given a prefix x_1,...,j−1 and a suffix x_j, the task of the rewriter is to produce a rewrite of the suffix ◯x_j that is superior to x_j for policy training. Policy training would proceeed using the same suffix x_1,...,j−1 but with the rewrite ◯x_j as the target.

There are various ways that the rewrite ◯x_j can be superior to the suffix x_j during training:

Overall quality: if the suffix is low quality, e.g. comes from a low quality part of the pretraining corpus, the rewriter can improve it, making the training target higher quality.
Safety: if the prefix and suffix are unsafe, the rewriter can steer the model towards a safe suffix given an unsafe prefix. Note this is quite different to simply rewriting the whole original document, which would mean the model is no longer exposed to unsafe inputs.
Augmentation: rewriting the data in various ways can improve performance, as has been shown in the offline setting of rewriting entire documents. This has been shown to improve diversity and knowledge (Hao et al., 2025; Allen-Zhu and Li, 2023), quality (Nguyen et al., 2025), and reasoning ability (Wang et al., 2025; Ishibashi et al., 2025). Our setting allows the model to steer from natural input (prefix) data towards new augmentations (via a rewritten suffix).

To build such a rewriter we can either directly prompt an existing post-trained model or fine-tune it further especially for this task. We detail our approach in section 3.

Suffix Judge Given a prefix x_1,...,j−1 and possible completions ¯x_j, the task of our judge is to discern which completion is superior as a target for policy training.

There are thus various ways that a judge can provide signal to improve the policy model, including:

Overall quality: if the suffix, rewrite or certain rollouts are low quality, they will receive low reward. At the start of training, rollouts are likely to be poor and the suffix or rewrite may receive higher reward. After sufficient training, rollouts are more likely to receive high reward.
Safety: if the prefix and suffix are unsafe, the rewrite or rollouts can steer the model towards a safe suffix given an unsafe prefix. Among the multiple policy rollouts the judge can choose between them to encourage safety amongst model generations.
Factuality: similarly, after sufficient training, selecting the most factual generations among the rollouts can improve the factuality of the policy model.

Similarly to building the rewriter, to build a judge we can either directly prompt an existing post-trained model, or further fine-tune it especially for this task. In our experiments, we consider both settings. We also consider judging each of the above — quality, safety and factuality — by prompting the post-trained judge model for each individually. The prompts we employ are given in Figure 2, Figure 3, and Figure 4. We detail our full approach in section 3.

Policy Model Training Putting it all together, we train our policy model using the sequence pretraining task described in section 2. We assume we have access to a post-trained model that can act as a suffix judge and a suffix rewriter, as described above.

For each prefix, we consider several candidate completions during online training. We can consider (i) the original suffix, (ii) a rewritten suffix; and (iii) K rollouts ¯x^k_j, k = 1, . . . , K from the current policy π.