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

Abstract: LLM post-training methods such as supervised fine-tuning (SFT), reinforcement learning (RL), and distillation are often analyzed through their loss functions: maximum likelihood, policy gradients, forward KL, reverse KL, or related objective-level variants. We study a complementary factor: the state distribution on which supervision is applied. For an autoregressive policy, a state is a prompt plus generated prefix. SFT trains on fixed dataset states, while RL and on-policy distillation (OPD) train on states induced by the current learner. We formalize post-training as state-distribution shaping and run a controlled smallscale study using Qwen3-0.6B-Base on GSM8K, with TruthfulQA and MMLU as retention evaluations. Our results show three phenomena. First, a mild SFT run improves GSM8K with little forgetting, while a stress SFT run causes substantial retention loss. Second, OPD from a degraded SFT teacher surpasses that teacher on GSM8K, TruthfulQA, and MMLU, despite using the teacher as its only supervision source. Third, a lightweight on-policy RL run improves GSM8K while preserving retention. These results support a state-centric view of post-training: the source and locality of training states can be as important as the form of the supervision signal.

• The paper demonstrates that state distribution is key in LLM post-training, influencing both performance improvement and retention.
• It compares off-policy SFT with on-policy techniques like RL and OPD, highlighting distinct trade-offs in stability and capability retention.
• Experimental findings indicate that focusing on local state supervision reduces catastrophic forgetting, suggesting novel hybrid algorithm designs.

State Distribution-Centric Post-Training of LLMs

Overview

This paper introduces a state-distribution framework for analyzing and designing post-training procedures in autoregressive LLMs, emphasizing the critical role played by the distribution of states—prompt plus generated prefix—used during optimization. The work provides a unified lens for supervised fine-tuning (SFT), reinforcement learning (RL), and on-policy distillation (OPD), arguing that the locus of supervision (off-policy vs. on-policy) is often as impactful as the loss function itself. Through systematic experiments with Qwen3-0.6B-Base on GSM8K, TruthfulQA, and MMLU, the paper establishes that state distribution is central to both model improvement and retention of prior capabilities.

The State Distribution Framework

The paper formalizes LLM post-training as the shaping of a model's induced state distribution, $d_\pi(s)$, where a "state" is the conditioning context of a prompt together with the generated prefix. The authors propose two key axes for each post-training algorithm:

• State source: Specifies whether supervision occurs on dataset-derived (off-policy) states or learner-induced (on-policy) states.
• Signal source: Describes whether the learning signal comes from labels, rewards, teacher predictions, or continuations.

By isolating these axes, the analysis transcends objective-centric views and exposes mechanistic explanations for phenomena such as catastrophic forgetting, teacher-student reversal, and local improvement, which are not easily accounted for by token-level losses alone.

Algorithmic Insights

SFT (Supervised Fine-Tuning)

SFT applies dense cross-entropy supervision on fixed dataset states, making it inherently off-policy. When the dataset states are compatible with the learner's own induced states, SFT is efficient and stable. However, aggressive or narrow off-policy SFT causes the model to update in regions of the state space unvisited during inference, precipitating catastrophic forgetting and degraded performance even on the target task.

RL (Reinforcement Learning)

RL optimizes rewards on trajectories sampled from the model's own policy. As an on-policy method, RL updates only on states the current model visits, preserving locality and avoiding global drift. RL may suffer from sample inefficiency due to sparse rewards but exhibits robust retention of non-target capabilities because state drift remains bounded and recoverable.

OPD (On-Policy Distillation)

OPD decouples supervision from state sampling: the student generates states; the teacher gives local guidance only on those states. This framework is analogous to DAgger-style imitation learning and enables students to surpass degraded teachers—a result attributed to learning from teacher repairs in locally reachable states while avoiding inheriting teacher trajectory-level errors.

Experimental Results

• Mild SFT improved GSM8K from 0.448 to 0.512 with negligible retention loss ($<1.5\%$), demonstrating that SFT is not inherently destructive.
• Stress SFT led to substantial retention degradation (retention ratio $0.8258$, forgetting $0.0635$), reducing both general and target task performance.
• OPD from degraded teacher yielded GSM8K $0.466$ (vs. teacher's $0.420$), TruthfulQA $0.275$ (vs. $0.245$), and MMLU $0.430$ (vs. $0.364$), confirming the student outperformed a degraded teacher on all axes.
• On-policy RL improved GSM8K to $<1.5\%$0 while nearly perfectly retaining baseline performance (mean forgetting $<1.5\%$1).

A notable finding is that maximum mean discrepancy (MMD) drift between base and post-trained rollouts is insufficient as a scalar measure; stress SFT and OPD from stress teacher exhibited similar MMD drift but vastly different retention outcomes, emphasizing that state-source locality dominates retention dynamics.

Implications and Theoretical Perspectives

The thesis that state-distribution locality is central to post-training leads to several substantial implications:

• Algorithm Design: The separation of state and signal sources enables new hybrid algorithms, e.g., on-policy dense shaping (continuation-based OPD), which preserves both task specialization and broad retention.
• Stability: On-policy post-training methods (RL, OPD) are less likely to induce catastrophic forgetting because updates are constrained to locally recoverable states, unlike dense off-policy SFT.
• Teacher-Student Dynamics: Students trained with OPD can systematically improve over degraded teachers due to selective learning from teacher repairs—not full trajectory imitation.

The findings motivate richer state-distribution analyses, beyond simple drift metrics, to characterize training dynamics and to inform more principled post-training protocols.

Practical Impact and Future Directions

This state-centric framework offers tangible guidance for practitioners: aggressive specialization with SFT should be avoided where broad retention is needed, and on-policy supervision (RL, OPD with continuation) should be considered for improved stability and task transfer. The approach suggests that future developments in AI post-training may focus on adaptive state-distribution control, advanced supervision shaping (trajectory-level guidance), and new metrics for modeling recoverability and locality in state space.

Further research should scale these methods to larger models, more diverse tasks, richer drift metrics (embedding-based), and integrate reference models for more robust state supervision. Exploration of adaptive state sampling and dynamic signal provision may also yield improved algorithms for safe and effective post-training.

Conclusion

The paper presents a rigorous state-distribution view of LLM post-training, demonstrating experimentally that on-policy locality and state-source selection are primary drivers of retention, forgetting, and improvement. Token-level objectives and scalar drift metrics are incomplete without consideration of training state distributions. The results challenge objective-centric dogma and open new avenues for principled post-training methods that balance specialization and generalization (2605.22731).