PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Junjie Huang, Jiarui Qin, Di Yin, Weiwen Liu, Yong Yu, Xing Sun, Weinan Zhang

arXiv:2602.03075v1cs.CL

TL;DR

The paper asks whether post-training improvements can retroactively strengthen the pre-training foundation rather than leaving training unidirectional. ReMiT uses an in-pipeline RL model to dynamically reweight reasoning-relevant tokens during mid-training. It improves 10 pre-training benchmarks by 3% on average and sustains gains of over 2% through post-training.

  • Problem

    Research has rarely examined how post-trained models can guide pre-training, despite the importance of a strong base model for effective post-training.

  • Method

    ReMiT reuses an RL-tuned model to compute token-level discrepancies and dynamically reweight mid-training objectives without hard token selection.

  • Results

    3% average improvement on 10 pre-training benchmarks persists as over 2% improvement throughout the post-training pipeline.

  • Takeaways & Limitations

    ReMiT establishes an iterative feedback loop in which improvements circulate between the base and RL models for continued evolution.

  • Takeaways & Limitations

    The approach assumes an underlying reasoning-quality reward function and uses a publicly released RL-tuned model as its reference policy.

Abstract

from arXiv · show

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.

1 Introduction

LLM training is usually unidirectional, leaving limited evidence on using post-training models to improve the pre-training foundation. ReMiT addresses this gap by using an in-pipeline RL model to guide token weighting during mid-training, producing gains that persist through post-training.

  • Standard pre-training weights every token uniformly, whereas RL assigns non-uniform weights from reward signals and improves reasoning capabilities.
  • Prior pre-training improvements mainly use document- or sample-level selection and reweighting, while token-level reasoning importance remains underexplored.
  • Mid-training is identified as a critical turning point because it uses high-quality corpora and rapidly decaying learning rates to acquire high-order capabilities.
  • ReMiT reuses an in-pipeline RL-tuned model to dynamically assign token-level weights during mid-training.
  • ReMiT establishes bidirectional synergy without specially trained or external reference models.
  • 3% average improvement on 10 pre-training benchmarks persists as over 2% improvement throughout post-training.

2 Preliminaries and Empirical Motivation

Mid-training refines models on high-quality reasoning-oriented data under a rapidly decaying learning rate, but standard next-token prediction weights all positions equally. ReMiT is motivated by evidence that capability-relevant token differences are localized and that mid-training produces a major distributional shift.

  • Mid-training concludes pre-training by refining models on high-quality corpora under a rapidly decaying learning rate while retaining the standard next-token prediction objective.
  • An autoregressive model assigns a sequence probability as the product of conditional next-token probabilities.
  • The mid-training dataset contains N token sequences whose parameters are optimized by minimizing average negative log-likelihood.
  • Uniform token weighting fails to distinguish low-entropy tokens from pivotal tokens that drive complex reasoning steps.
  • After mid-training, the token-level log-probability gap between base and RL models narrows substantially, indicating a qualitative transformation in the base model.
  • RL-model advantages are highly localized, with distinct pivotal tokens receiving much higher likelihood than under the base model.
  • ReMiT therefore targets mid-training for dynamically upweighting pivotal tokens in reasoning-oriented, post-training-aligned corpora.

3 The ReMiT Framework

ReMiT forms a bidirectional training loop by reusing the RL model to guide mid-training, dynamically emphasizing tokens where the base model underperforms. Soft, normalized, bounded weighting preserves sequence coherence and stabilizes optimization.

  • ReMiT retroactively transfers reasoning priors from post-training to the pre-training foundation, creating a self-reinforcing flywheel.
  • The framework reuses the RL-tuned model produced by the previous cycle instead of training a separate teacher model.
  • Mid-training is selected because reasoning-oriented corpora and rapid learning-rate decay enable efficient assimilation of the RL reference’s priors.
  • Dynamic weighting prioritizes tokens where the base model significantly underperforms relative to the RL reference.
  • Sequence-level centering normalizes log-probability discrepancies so token weights remain stable and comparable across sequences.
  • A scaled sigmoid maps centered discrepancies to bounded weights, preserving standard NTP when model confidence is equal.
  • Clipping acts as a safety rail against token overfitting and degradation of basic syntactic coherence.
  • ReMiT integrates these weights through soft reweighting, amplifying pivotal-token gradients while preserving sequence semantic coherence.

4 Theoretical Motivation

ReMiT can be understood as weighted training toward an implicit target distribution, with theoretical analysis linking its updates to KL divergence from a KL-regularized optimal policy. The framework uses the RL model to identify reasoning-relevant tokens while preserving the ground-truth training direction rather than mimicking the teacher distribution.

  • 4.2 Directional Consistency with the Optimal Policy: ReMiT is theoretically motivated as an update that locally reduces KL divergence between the base model and the KL-regularized optimal policy.The analysis frames token reweighting as directionally connected to the optimal policy under KL regularization.
  • 4.1 ReMiT as Optimization towards an Implicit Target Distribution: ReMiT’s weighted negative log-likelihood is equivalent to minimizing divergence toward an implicit target distribution shaped by token-level weights.The target is constructed by reweighting the data distribution according to log-probability gaps.
  • 4.1 ReMiT as Optimization towards an Implicit Target Distribution: The implicit target distribution assigns higher probability mass to pivotal tokens, causing ReMiT updates to steer the model distribution toward this enhanced data distribution.The RL model influences optimization by reshaping the target rather than directly replacing the training distribution.
  • 4.2 Directional Consistency with the Optimal Policy: Directional consistency holds when the implicit target assigns greater probability to tokens favored by the optimal policy.The local update is defined as a small interpolation toward the implicit target, and the KL divergence decreases under the stated condition.
  • 4.3 Comparisons with Knowledge Distillation: Unlike standard knowledge distillation, ReMiT uses hard targets and dynamically scales gradient magnitude on informative tokens while preserving the ground-truth direction.Standard KD matches the teacher’s full probability distribution, including low-probability tokens; ReMiT instead reweights the negative log-likelihood.

5 Experiments

Experiments evaluate ReMiT across three open-source model families and ten downstream tasks, then test whether its mid-training gains persist through post-training and iterative refinement. ReMiT consistently outperforms baseline strategies, transfers gains across post-training methods, and maintains advantages over knowledge distillation.

  • Experimental Setup: ReMiT is evaluated on OLMo-1B, SmolLM3-3B, and Youtu-LLM-2B using few-shot downstream benchmarks spanning reasoning and general capabilities.The evaluation includes tasks such as MATH, MBPP, MMLU-Pro, TruthfulQA, and ARC-C.
  • Experimental Setup: 50B tokens are used for mid-training with each model family’s official high-quality corpus and aligned optimization hyperparameters.ReMiT uses clipping threshold ϵ = 0.2.
  • Main Results: ReMiT achieves the best mid-training performance by dynamically prioritizing reasoning-critical tokens while preserving semantic coherence through soft reweighting.It outperforms both RHO-1’s token selection and MiniPLM’s sample selection.
  • Main Results: ReMiT’s mid-training gains transfer consistently to post-training under direct RL and the complete SFT, DPO, and RLVR pipeline.The comparison applies identical post-training procedures to vanilla NTP and ReMiT models on OLMo-1B.
  • Main Results: ReMiT2 further improves over ReMiT when an initial ReMiT-derived RL model guides a second mid-training cycle.This produces a cascading feedback loop in which an improved base yields a stronger RL reference.
  • Analysis and Discussion: ReMiT maintains later-stage advantages because its moderate distributional gap avoids the over-alignment associated with knowledge distillation.KD’s immediate gains fade after post-training, whereas ReMiT preserves adaptation potential.
  • Analysis and Discussion: Clipping prevents gradient collapse, preserves general modeling capabilities, and accelerates convergence while prioritizing reasoning-critical tokens.The mechanism contrasts with aggressive token discarding, which diminishes gradients.

6 RELATED WORK

Related work improves pre-training through data selection and examines token-level signals associated with capability-relevant updates. Existing approaches range from document and sample reweighting to analyses of sparse pivotal tokens and entropy patterns.

  • Data Selection Strategies in Pre-Training: Pre-training data-selection methods use document cleaning, heuristic or model-based filtering, and sample-level reweighting to improve corpus quality.Sample-level approaches are described as relatively coarse-grained when based on heuristic rules.
  • Token-Level Analysis: Recent token-level analyses suggest that low-probability or high-entropy tokens and sparse key tokens disproportionately influence meaningful model updates.These tokens are associated with pivotal decision junctions in reasoning.

7 CONCLUSION

The paper couples pre-training and post-training by identifying mid-training as a critical turning point and introducing ReMiT for RL-guided token reweighting. Across model families, ReMiT improves both stages and establishes an iterative feedback loop between base and RL models.

  • Conclusion: ReMiT assigns dynamic token-wise weights from an RL reference during mid-training, guiding the base model toward pivotal reasoning tokens.Its soft reweighting preserves semantic coherence and training efficiency rather than discarding data.
  • Conclusion: ReMiT consistently enhances mid-training and post-training across multiple model families, forming an iterative feedback loop between the base and RL model.The conclusion characterizes this loop as coupling the two training stages.

A Background: The Mid-Training Phase

Mid-training bridges broad pre-training and specialized capability refinement through higher-quality data and rapid learning-rate decay. The paper frames this stage as a window where token reweighting can align the base model more closely with reasoning-oriented behavior.

  • Mid-Training Characteristics: Mid-training shifts training toward higher-quality STEM, code, and reasoning data while rapidly decreasing the learning rate.These changes distinguish mid-training from stable pre-training and support capability refinement.
  • Distributional Shift: Mid-training narrows the token-distribution gap between the base and RL models, while ReMiT dynamically reweights tokens during this transition.The distributional comparison is visualized through token-level log-probability differences.
  • Reference-Model Analysis: The RL model assigns higher likelihood to selected discourse connectives than the SFT model, revealing distinct token-level preferences linked to reasoning behavior.The comparison uses log-probability margins between the RL and SFT models.

D.3 Evaluation of Pre-trained Base LMs

The evaluation measures base and post-trained models on broad suites of 10 downstream benchmarks covering mathematics, code, reasoning, instruction following, truthfulness, and general knowledge.

  • Base-Model Evaluation: Base models are evaluated few-shot across 10 widely used downstream benchmarks.The suite includes mathematics, reasoning, instruction following, code generation, truthfulness, and general-knowledge tasks.
  • Post-Training Evaluation: Post-trained models are evaluated zero-shot on 10 benchmarks spanning mathematics, code, reasoning, instruction following, and general capability.The reported suite includes GSM8K, Minerva Math, MBPP, HumanEval, ARC-Challenge, GPQADiamond, MMLUPro, and IFEval variants.

E.1 Pre-Training Results

The experiments compare ReMiT with alternative mid-training approaches and examine whether its gains persist through post-training. Results favor RL-guided token reweighting, with robustness to imperfect references and sustained downstream benefits.

  • Reference Policy Ablation: ReMiT consistently outperforms SFT-Guided throughout OLMo-1B mid-training and across the accompanying benchmark comparison.The results empirically support using the RL model rather than the SFT model as the reference policy.
  • Knowledge Distillation: Knowledge distillation improves mid-training performance immediately but its advantages fail to persist after subsequent post-training.The paper attributes this pattern to strict teacher-distribution matching potentially reducing future adaptation plasticity.
  • Post-Training Results: ReMiT achieves lower SFT and DPO losses, a higher RLVR verifiable correct rate, and preserved gains across the post-training pipeline.Table 3 reports post-trained performance across 10 downstream benchmarks.
  • Reference Quality: ReMiT outperforms ReMiT (Intermediate), which outperforms Vanilla NTP, showing a positive relationship between reference quality and downstream performance.The intermediate reference was trained for only 10% of the total RL steps.
  • Robustness and Scalability: ReMiT (Intermediate) still outperforms Vanilla NTP, while clipping limits extreme token weights under imperfect guidance.These findings support robustness to weaker references and the in-pipeline reference strategy.

F.2 Computational Efficiency and Overhead Analysis

ReMiT adds reference-model computation and reduces training throughput, but its faster convergence offsets that overhead in GPU-hour efficiency. The analysis also reports a higher performance ceiling than Vanilla NTP under extended training.

  • Time-to-Convergence Efficiency: ReMiT achieves a 3.5x speedup in time-to-convergence over Vanilla NTP when measured in GPU hours.The reported speedup is presented as outweighing the reduction in training throughput.
  • Overhead vs. Gain: ReMiT incurs a 43% reduction in throughput because of the frozen RL reference’s additional inference overhead.The overhead is amortized by the method’s faster learning trajectory.
  • Performance Ceiling: ReMiT reaches accuracy levels that Vanilla NTP struggles to attain even with extended training.This comparison indicates a qualitative reasoning improvement beyond simply training the baseline longer or on more tokens.
Loading 2602.03075v1…