Self-Improving Pretraining: using post-trained models to pretrain better models

TL;DR

Self-Improving Pretraining reframes pretraining as a prefix-suffix sequence generation task and leverages a strong post-trained model to rewrite and judge suffix completions. By training a policy via online reinforcement learning to improve the next $K$ tokens, the method optimizes quality, safety, and factuality from the ground up, rather than relying solely on post-hoc alignment. Across continual and from-scratch pretraining, the approach yields significant gains: up to $36.2\%$ relative improvement in factuality, $18.5\%$ in safety, and up to an $86.3\%$ win-rate increase in overall generation quality. The framework demonstrates that upstream, goal-directed supervision can complement standard pretraining, suggesting broader applicability to multi-skill alignment and safer AI deployments.

Abstract

Ensuring safety, factuality and overall quality in the generations of large language models is a critical challenge, especially as these models are increasingly deployed in real-world applications. The prevailing approach to addressing these issues involves collecting expensive, carefully curated datasets and applying multiple stages of fine-tuning and alignment. However, even this complex pipeline cannot guarantee the correction of patterns learned during pretraining. Therefore, addressing these issues during pretraining is crucial, as it shapes a model's core behaviors and prevents unsafe or hallucinated outputs from becoming deeply embedded. To tackle this issue, we introduce a new pretraining method that streams documents and uses reinforcement learning (RL) to improve the next K generated tokens at each step. A strong, post-trained model judges candidate generations -- including model rollouts, the original suffix, and a rewritten suffix -- for quality, safety, and factuality. Early in training, the process relies on the original and rewritten suffixes; as the model improves, RL rewards high-quality rollouts. This approach builds higher quality, safer, and more factual models from the ground up. In experiments, our method gives 36.2% and 18.5% relative improvements over standard pretraining in terms of factuality and safety, and up to 86.3% win rate improvements in overall generation quality.

Figures (14)

• Figure 1: Self-Improving pretraining: Our proposed model training streams pretraining documents and improves the next $K$ generated tokens (suffix, given prefix) at each step with RL. A strong previously post-trained model is used to judge generation candidates at each RL step for quality, safety and hallucination, where the candidates are: (i) $N$ rollouts from the current policy; (ii) the original suffix; and (iii) a rewrite of the suffix by the strong post-trained model. The rewrite can improve the pretrain data's quality or safety; in the latter case as the prefix remains unsafe the model is always learning how to steer away to a safe suffix. At the start of training model rollouts (i) are low quality, so training relies on candidates (ii) and (iii); later in training the judge starts rewarding winning rollouts.
• Figure 2: Suffix judge safety prompt.
• Figure 4: Factuality judge prompt.
• Figure 5: Suffix judge validation rewards on safety and quality tasks. Initial performance of the model is close to random chance on either task, achieving scores above $90\%$ by the end of training.
• Figure 6: Suffix rewriter validation rewards on safe and unsafe suffixes of the RedPajama dataset. Initial performance of the model is close to random chance on the safety task (0.5 score on unsafe suffixes), and near zero on copying safe suffixes (exact match reward score of 0.1), but still increasing after 500 steps.