Proximal Supervised Fine-Tuning

1. Problem Statement

The paper addresses two critical limitations of standard Supervised Fine-Tuning (SFT) for Large Language Models (LLMs):

• Poor Generalization: Standard SFT often leads to "behavior cloning," where the model memorizes specific training data rather than learning generalizable reasoning patterns. This results in performance degradation on out-of-domain (OOD) tasks and a loss of pre-trained capabilities.
• Entropy Collapse: Aggressive SFT can cause the model's output distribution to become overly concentrated (low entropy). This "entropy collapse" limits the model's ability to explore during subsequent Reinforcement Learning (RL) stages, effectively constraining the potential gains from RL optimization (e.g., PPO, GRPO).

The authors argue that while RL algorithms like PPO successfully use "trust regions" to prevent policy drift, standard SFT lacks such constraints, leading to unstable updates and overfitting.

2. Methodology: Proximal Supervised Fine-Tuning (PSFT)

The authors propose PSFT, a novel fine-tuning objective that adapts the Proximal Policy Optimization (PPO) mechanism to the supervised learning setting.

Theoretical Foundation

• SFT as Policy Gradient: The authors frame SFT as a special case of policy gradient methods where the advantage function ($\hat{A}_t$) is fixed to 1 for ground-truth tokens, and sampling occurs from a fixed offline dataset rather than the current policy.
• Trust Region Adaptation: Inspired by TRPO and PPO, PSFT introduces a clipped surrogate objective to constrain the ratio between the new policy ($\pi_\theta$) and the old policy ($\pi_{\theta_{old}}$).

The Objective Function

The PSFT loss function is defined as:
$$\mathcal{L}_{PSFT}(\theta) = \mathbb{E}_{(s_t, a_t) \sim D} \left[ \min \left( r_t(\theta), \text{clip}(r_t(\theta), 1-\epsilon, 1+\epsilon) \right) \right]$$
Where:

• $r_t(\theta) = \frac{\pi_\theta(a_t | s_t)}{\pi_{\theta_{old}}(a_t | s_t)}$ is the importance sampling ratio.
• $\epsilon$ is a hyperparameter (typically 0.2 or 0.28) defining the trust region.
• Unlike PPO, there is no advantage estimation; the "advantage" is implicitly 1 for correct tokens.

Key Mechanisms

1. Dynamic Old Policy ($\pi_{\theta_{old}}$): Unlike standard PPO which often uses a fixed reference model, PSFT dynamically updates $\pi_{\theta_{old}}$ during training. This allows the trust region to shift gradually, enabling the model to learn from the dataset without being confined to a static reference point.
2. Gradient Clipping: The clipping mechanism acts as a regularizer. If the new policy deviates too much from the old one (i.e., $r_t(\theta) > 1+\epsilon$), the gradient for that token is zeroed out. This prevents large, destructive updates that cause entropy collapse.
3. Warm-up Phase: The authors suggest an optional initial warm-up phase where standard SFT is performed to align the initial $\pi_{\theta_{old}}$ with the dataset distribution before switching to the PSFT objective, further stabilizing training.

3. Key Contributions

1. Novel Optimization Objective: PSFT is the first method to successfully translate the "trust region" concept from RL (PPO) into the supervised fine-tuning domain to prevent policy drift and entropy collapse.
2. Entropy Preservation: PSFT effectively maintains high output entropy during training, ensuring the model retains diversity and exploration capabilities for subsequent RL stages.
3. Superior Generalization: By preventing overfitting to the training distribution, PSFT models demonstrate significantly better performance on out-of-domain tasks compared to standard SFT.
4. Enhanced RL Foundation: Models initialized with PSFT serve as a stronger "cold start" for RL algorithms (like GRPO or DPO), leading to faster convergence and higher final performance.

4. Experimental Results

The authors evaluated PSFT across Mathematical Reasoning, Human Value Alignment, and Multimodal (VLM) domains using models like Qwen2.5-7B, Llama3.1-8B, and Qwen3-30B.

A. SFT Stage Performance

• In-Domain: PSFT achieves performance comparable to (and sometimes exceeding) standard SFT on target tasks (e.g., AIME, MATH-500).
• Out-of-Domain: PSFT significantly outperforms standard SFT on OOD benchmarks (e.g., GPQA, TruthfulQA, IFEval). For instance, on IFEval (instruction following), standard SFT caused a severe drop in performance, while PSFT maintained or improved it.
• Entropy Stability: Entropy curves for PSFT remain smooth and stable, whereas standard SFT exhibits sharp declines (entropy collapse) after each epoch.

B. RL Stage Performance (Post-Training)

• Cold Start Advantage: When used as the starting point for RL (GRPO), PSFT-initialized models showed a "slow start" but achieved rapid catch-up, ultimately surpassing models initialized with standard SFT.
• Final Performance: In RL experiments, PSFT $\to$ GRPO combinations achieved higher scores on both in-domain and OOD tasks compared to SFT $\to$ GRPO.
• Alignment: In human alignment tasks (using DPO), PSFT reduced the "alignment tax" (the trade-off between alignment and general capability), allowing the model to learn preferences without losing reasoning abilities.

C. Multimodal (VLM) Results

• PSFT successfully injected long Chain-of-Thought (CoT) reasoning capabilities into Vision Language Models (VLMs) using text-only data without degrading visual understanding, a common failure point for standard SFT.

5. Significance and Impact

• Paradigm Shift: The paper challenges the notion that SFT is merely "behavior cloning." By incorporating RL-inspired constraints, it transforms SFT into a more robust optimization process that balances task-specific learning with general capability preservation.
• Practical Utility: PSFT offers a simple, drop-in replacement for standard SFT in the post-training pipeline. It requires minimal code changes (replacing the loss function) but yields substantial improvements in model robustness and downstream RL potential.
• Scalability: The method is validated across various model sizes (7B to 30B) and architectures (dense and MoE), suggesting it is a scalable solution for industrial-scale model training.

In conclusion, Proximal Supervised Fine-Tuning (PSFT) provides a theoretically grounded and empirically validated solution to the generalization and entropy collapse problems in SFT, establishing a more stable and capable foundation for the subsequent reinforcement learning stages of LLM development.