ReMiT: RL-Guided Mid-Training for Iterative LLM Evolution

TL;DR

This work tackles the challenge of bidirectional improvement in LLM training by introducing ReMiT, an RL-guided mid-training framework that reweights tokens to transfer reasoning priors from an RL-refined model back to the base pre-training stage. By identifying mid-training as a critical window and using an in-pipeline RL reference, ReMiT creates a soft, token-level reweighting scheme that drives the base model toward pivotal reasoning tokens, formalized as optimization toward an implicit target distribution and KL-regularized transfer. Theoretical analysis shows that these updates align with reducing divergence to the optimal policy and avoid overfitting the RL teacher, unlike standard knowledge distillation. Empirically, ReMiT yields average gains around 3% across 10 pre-training benchmarks and maintains over 2% improvements through post-training, demonstrating a co-improving flywheel that enables continuous, self-reinforcing evolution of LLMs with efficiency advantages and without external teachers.

Abstract

Standard training pipelines for large language models (LLMs) are typically unidirectional, progressing from pre-training to post-training. However, the potential for a bidirectional process--where insights from post-training retroactively improve the pre-trained foundation--remains unexplored. We aim to establish a self-reinforcing flywheel: a cycle in which reinforcement learning (RL)-tuned model strengthens the base model, which in turn enhances subsequent post-training performance, requiring no specially trained teacher or reference model. To realize this, we analyze training dynamics and identify the mid-training (annealing) phase as a critical turning point for model capabilities. This phase typically occurs at the end of pre-training, utilizing high-quality corpora under a rapidly decaying learning rate. Building upon this insight, we introduce ReMiT (Reinforcement Learning-Guided Mid-Training). Specifically, ReMiT leverages the reasoning priors of RL-tuned models to dynamically reweight tokens during the mid-training phase, prioritizing those pivotal for reasoning. Empirically, ReMiT achieves an average improvement of 3\% on 10 pre-training benchmarks, spanning math, code, and general reasoning, and sustains these gains by over 2\% throughout the post-training pipeline. These results validate an iterative feedback loop, enabling continuous and self-reinforcing evolution of LLMs.

Figures (16)

• Figure 1: ReMiT on OLMo-1B substantially outperforms baselines trained with the standard mid-training method. (a) At the mid-training stage, ReMiT improves average accuracy across 10 widely-used benchmark tasks by 5.2% and reaches the baseline performance $6\times$ faster. (b) The improvements can carry over to post-training: during RL, ReMiT maintains a higher verifiable correct rate than the baseline and achieves better performance.
• Figure 2: We identify the mid-training stage as a critical turning point, as it rapidly shifts the base model’s token distribution toward that of a more capable RL model. ReMiT enhances this stage by dynamically reweighting tokens in the mid-training corpus.
• Figure 3: Visualization of the log-probability divergence between the pre-trained base model and the RL model. Background intensity reflects the margin $\Delta\log p=\log p_{\mathrm{RL}}-\log p_{\mathrm{base}}$, where deeper red highlights pivotal tokens on which the RL model demonstrates significantly higher confidence than the base model.
• Figure 4: Overview of the proposed ReMiT framework. The pipeline connects pre-training and post-training, establishing a self-reinforcing flywheel: improvements from the RL stage are retroactively transferred to strengthen the base model foundation, which in turn amplifies performance in subsequent post-training stages.
• Figure 5: Performance gains of ReMiT on OLMo-1B acquired during mid-training transfer consistently to post-training, regardless of the post-training process (SFT, DPO, or RLVR). Figure (a): applying RL directly to the mid-trained base model. Figure (b): applying the complete post-training procedure to the mid-trained base model.