Stabilizing On-Policy Distillation for MLLM Reasoning with Global Normalization

Abstract: On-policy distillation (OPD) has recently emerged as an important post-training paradigm. By using a stronger teacher model to provide dense, fine-grained supervision for sampled trajectories, OPD offers a clear advantage over reinforcement learning with verifiable rewards (RLVR), which typically depends on sparse binary or outcome-based environmental feedback. However, naive token-level distillation can suffer from gradient instability, due to magnitude misalignment in outlier states. To address this issue, we propose Globally Normalized Distillation Policy Optimization (GNDPO), a practical method that stabilizes optimization by transforming raw KL scores into batch-level relative advantages. This normalization effectively mitigates gradient explosions while retaining the benefits of token-level guidance. Experimental results show that GNDPO substantially improves training robustness and downstream performance across multimodal reasoning tasks. The code is released at https://github.com/OPPO-Mente-Lab/GNDPO.

• The paper presents GNDPO, a method that globally normalizes token-level KL rewards to stabilize on-policy distillation for MLLM reasoning.
• It introduces batch-level normalization to adapt gradient scales, preventing spikes and ensuring smooth, context-invariant optimization.
• Empirical results across InternVL3.5 variants reveal up to 1.6% improvement on reasoning benchmarks compared to standard OPD methods.

Stabilizing On-Policy Distillation for MLLM Reasoning with Global Normalization

Introduction

The paper addresses the instability in on-policy distillation (OPD) for multimodal LLMs (MLLMs) by introducing Globally Normalized Distillation Policy Optimization (GNDPO). OPD, which leverages fine-grained, token-level teacher feedback in post-training, provides denser signal than RL approaches reliant on sparse, terminal rewards. However, direct application of KL-based token-level rewards leads to gradient spikes during on-policy exploration, particularly in out-of-distribution contexts, resulting in unstable or divergent optimization. This paper proposes global normalization of KL-based advantage signals across training batches, ensuring stability while preserving fine-grained supervision.

Background: On-Policy Optimization and Distillation in MLLMs

Standard supervised fine-tuning (SFT) is inadequate for aligning MLLMs with downstream reasoning tasks due to exposure bias and distribution shift. Reinforcement learning with verifiable rewards (RLVR) such as GSPO provides on-policy distribution correction but is limited by sparse, sequence-level signal, leading to inefficient credit assignment and suboptimal reasoning.

OPD addresses these limitations by providing token-level, teacher-derived KL guidance. However, magnitude mismatches and OOD actions can cause local logit discrepancies and gradient explosion. Such instability prevents reliable optimization, particularly as model scale and multimodal reasoning complexity increase.

Figure 1: GSPO, OPD, and GNDPO workflows highlighting the dense but unstable signals in OPD and the batch-level normalization in GNDPO that stabilizes credit assignment.

Method: Globally Normalized Distillation Policy Optimization

GNDPO reformulates token-level KL-guided distillation by normalizing rewards within each training batch. For a minibatch $\mathcal{B}$ and each token $o_{i,t}$, the normalized advantage is computed as:

$$\hat{A}_{i,t}^{\text{norm}} = \frac{\hat{A}_{i,t} - \mu_\mathcal{B}}{\sigma_\mathcal{B} + \epsilon}$$

where $\hat{A}_{i,t}$ is the raw negative reverse KL divergence, $\mu_\mathcal{B}$ is the batch mean, and $\sigma_\mathcal{B}$ is the batch standard deviation. This procedure transforms absolute teacher-student signal into a context-agnostic, relative advantage, ensuring that singular outlier tokens cannot dominate the optimization trajectory.

GNDPO's normalization has two principal effects:

• Adaptive gradient scaling: Catastrophic mismatches inflate $\sigma_\mathcal{B}$, automatically reducing gradient step magnitude for outliers, thus suppressing gradient explosions.
• Context invariant optimization: By subtracting the batch mean, GNDPO shifts the target from absolute distribution recovery to maximizing relative prediction quality, balancing contributions across diverse contexts and tasks.

Experiments and Numerical Results

The method is evaluated with three InternVL3.5 model variants (1B, 2B, 4B parameters) using InternVL3.5-8B as a teacher. The training corpus is Geometry3K, and evaluation covers MathVista, MMMU, MathVision, MathVerse, DynaMath, WeMath, and LogicVista, providing a comprehensive assessment of multimodal mathematical and logical reasoning.

Across all benchmarks, GNDPO outperforms both GSPO and standard OPD. Key results include:

• InternVL3.5-4B: GNDPO achieves up to +1.1% on MMMU, +1.1% on MathVision, and +1.6% on WeMath compared to OPD.
• InternVL3.5-2B: +1.3% on MathVista, +0.5% on MathVision, +1.7% on MathVerse.
• InternVL3.5-1B: Gains are robust even for constrained models, with +0.7% on MathVista and +0.5% on MMMU.

In all scenarios, the globally normalized scheme leads to consistent, aggregate improvements in reasoning accuracy.

Training Dynamics and Stability

Training stability is further corroborated via loss, gradient norm, and entropy traces. OPD is characterized by volatile loss and catastrophic gradient norm spikes, whereas GNDPO maintains smooth loss dynamics and suppresses gradient explosions.

Figure 2: GNDPO exhibits lower, smoother loss dynamics and suppresses gradient norm spikes during training compared to standard OPD.

GNDPO’s normalized gradient flow supports stable exploration as reflected in entropy curves, ensuring effective learning signal propagation without sacrificing search diversity.

Implications and Future Directions

GNDPO provides a principled mechanism to harness token-level supervision in policy optimization for MLLMs, crucial for high-fidelity reasoning and scalable generalization. Practically, it unlocks robust training for large-scale models in domains requiring precise, dense alignment, such as mathematical and logical reasoning or visuolinguistic understanding. The global normalization scheme is readily extensible to other on-policy distillation paradigms or domains with dense but volatile teacher feedback.

Theoretical implications include the necessity to move beyond pointwise supervision in RL-inspired post-training, emphasizing robust, batch-level reward modeling to address outlier effects. Future directions include exploration of batch statistics under very small batch regimes, application to open-ended text generation, and adaptation for streaming or non-i.i.d. data settings intrinsic to interactive or multi-turn dialogue.

Conclusion

The paper presents Globally Normalized Distillation Policy Optimization (GNDPO), a method that stabilizes on-policy, KL-guided distillation for multimodal LLMs by converting raw token-level rewards into robust, globally normalized batch-level advantages. Empirical results on diverse reasoning benchmarks demonstrate clear improvements in both accuracy and training stability over standard OPD and RLVR approaches. GNDPO's design principles and results point to broader trends in large model post-training, where balance between rich dense supervision and training stability is essential for advanced reasoning capabilities in AI systems.