Post-Training is About States, Not Tokens: A State Distribution View of SFT, RL, and On-Policy Distillation

Technical Summary: Post-Training as State Distribution Shaping

Problem Statement

Current analyses of Large Language Model (LLM) post-training methods—specifically Supervised Fine-Tuning (SFT), Reinforcement Learning (RL), and Distillation—primarily focus on the optimization objectives (e.g., maximum likelihood, policy gradients, KL divergence). However, this objective-centric view struggles to explain several empirical phenomena:

• Why mild SFT can improve performance with minimal forgetting, while aggressive ("stress") SFT causes catastrophic forgetting and degraded target performance.
• Why RL, despite using sparse or noisy rewards, often yields stable improvements.
• Why a student model in distillation can sometimes outperform a degraded teacher, even when the teacher is the sole supervision source.

The paper argues that these discrepancies arise because standard analyses overlook a critical axis of variation: the state distribution on which supervision is applied. In autoregressive models, a "state" is defined as the prompt plus the generated prefix ($s_t = (x, y_{<t})$). Different post-training methods apply their signals to fundamentally different state distributions (fixed dataset prefixes vs. learner-induced rollouts), which significantly impacts model behavior, locality, and retention.

Methodology

The authors propose a unified framework viewing post-training as state-distribution shaping. They decompose any post-training method into two distinct axes:

1. State Source ($q_k$): Where the training contexts (states) originate.
2. Signal Source ($S$): What supervision (tokens, rewards, logits) is provided at those states.

Theoretical Formulation:

• SFT: Off-policy state fitting. The model minimizes token loss on fixed dataset states ($d_{data}$). This creates a mismatch if dataset trajectories diverge from the model's own rollouts, leading to exposure bias and potential instability under stress.
• RL: On-policy local improvement. The model samples trajectories from its current policy ($d_{\pi_\theta}$) and updates based on rewards. Updates are naturally localized to states the model actually visits.
• On-Policy Distillation (OPD): A hybrid where the student samples states ($d_{\pi_S}$) but the teacher provides the signal (logits or continuations). This decouples the state source from the signal source, allowing the student to learn local repairs without inheriting the teacher's global trajectory distribution.

Experimental Setup:

• Model: Qwen3-0.6B-Base with LoRA adapters.
• Hardware: Single RTX 3090 GPU.
• Tasks:
  • Target: GSM8K (math reasoning).
  • Retention: TruthfulQA and MMLU (to measure catastrophic forgetting).
• Metrics: Target accuracy, retention ratio, forgetting magnitude, and state drift (measured via Maximum Mean Discrepancy, MMD, on lexical features of rollout states).

Experimental Conditions:

1. Mild SFT: Standard training on GSM8K.
2. Stress SFT: Aggressive training (5 epochs, high LR, high LoRA rank) to induce degradation.
3. OPD: Training a student using a degraded "Stress SFT" teacher as the signal source, but sampling states from the student's own policy.
4. On-Policy RL: Lightweight RL using group-relative exact-answer rewards on GSM8K.

Key Results

1. SFT is Context-Dependent: Mild SFT improved GSM8K (0.448 $\to$ 0.512) with negligible retention loss. Conversely, Stress SFT degraded both the target task (0.420) and retention tasks (TruthfulQA/MMLU dropped significantly), demonstrating that SFT is not inherently destructive but becomes so under off-policy pressure that misaligns with the model's natural state distribution.
2. OPD Surpasses Degraded Teachers: A student trained via OPD using the "Stress SFT" model as a teacher outperformed the teacher on all metrics (GSM8K: 0.466 vs. 0.420; MMLU: 0.430 vs. 0.364). This confirms that students do not inherit the teacher's global trajectory failures if they are trained on their own states, receiving only local conditional guidance.
3. RL Preserves Locality: Lightweight on-policy RL improved GSM8K (0.472) while maintaining near-base retention levels. This supports the view that RL acts as a local improvement procedure grounded in the learner's reachable states.
4. Drift Metrics are Insufficient: Methods with nearly identical scalar state drift (MMD) exhibited vastly different retention outcomes. For instance, Stress SFT and OPD from the Stress teacher had similar MMD values (~0.0109) but vastly different retention ratios (0.82 vs. 0.95). This indicates that the source and locality of training states matter more than the magnitude of distributional shift alone.

Contributions

1. Unified Framework: Formulated SFT, RL, and OPD under a common "state-distribution view," explicitly separating the state source from the signal source.
2. Controlled Empirical Pipeline: Implemented a single-GPU experimental setup measuring target accuracy, retention, forgetting, and rollout-state drift to isolate the effects of state distribution.
3. Evidence for State-Centric Claims: Provided empirical evidence that:
   • OPD can recover performance and surpass a degraded teacher by decoupling state sampling from signal generation.
   • On-policy RL improves target performance with minimal retention loss due to its locality.
   • Scalar drift metrics (like MMD) cannot fully characterize post-training dynamics without considering the state source.

Significance and Claims

The paper claims that post-training behavior is governed by the interaction between the supervision signal and the training state distribution. The authors argue that the "source, locality, and learner-dependence of training states" are as critical as the form of the supervision signal itself.

• Reframing SFT: Catastrophic forgetting is not solely a function of maximum likelihood but arises when dense off-policy updates push the model into regions incompatible with its own future state distribution.
• Reframing Distillation: A student can outperform a teacher because the teacher's errors are often coupled to the teacher's own trajectory distribution. By querying the teacher on the student's states, the student accesses useful local repairs without inheriting global trajectory failures.
• Reframing RL: RL's stability is partly attributed to its on-policy nature, ensuring updates are applied only to states the model can actually reach and recover from.

The paper concludes that a state-centric view is necessary to explain stability and performance reversals that objective-level analysis alone cannot account for. The authors emphasize that these findings are based on a small-scale study and should be viewed as evidence for a mechanistic hypothesis rather than a definitive benchmark claim.