Technical Summary: Beyond Mode Collapse: Distribution Matching for Diverse Reasoning

1. Problem Statement

The paper identifies mode collapse as a fundamental limitation in on-policy reinforcement learning (RL) methods, particularly those used for training Large Reasoning Models (LRMs) via Reinforcement Learning with Verifiable Rewards (RLVR). While methods like Group Relative Policy Optimization (GRPO) have achieved success in mathematical problem solving and code generation, they exhibit a tendency to concentrate probability mass on a single high-reward trajectory once discovered. This behavior suppresses alternative solution strategies, halts exploration, and leads to brittle reasoning that fails to generalize.

The authors argue that this phenomenon is not merely an implementation artifact but a direct consequence of the optimization objective: standard on-policy RL implicitly minimizes the reverse Kullback-Leibler (KL) divergence ($D_{KL}(\pi_\theta \| p^*)$). Reverse KL is inherently "mode-seeking," amplifying the first successful trajectory encountered rather than maintaining a distribution over multiple diverse high-reward solutions. In contrast, forward KL divergence ($D_{KL}(p^* \| \pi_\theta)$) exhibits "mode-covering" behavior, encouraging the policy to match the full support of the target distribution, but it is typically intractable because it requires sampling from the global reward-weighted distribution.

2. Methodology: DMPO

To address mode collapse without requiring intractable global sampling, the authors propose Distribution-Matching Policy Optimization (DMPO).

Core Insight

DMPO approximates forward KL minimization at the group level rather than globally. Instead of sampling from a global target distribution, the algorithm constructs a target distribution over a specific group of sampled trajectories for a given prompt.

Algorithmic Formulation

1. Group-Level Target Distribution: For a group of $G$ trajectories $\mathcal{O} = \{o_1, \dots, o_G\}$ sampled from the current policy $\pi_{\theta_{old}}$, a target distribution $p(o_i | \mathcal{O})$ is constructed proportional to the exponential of their rewards (a Boltzmann distribution):
   $$p(o_i | \mathcal{O}) = \frac{\exp(r(o_i)/\alpha)}{\sum_{j=1}^G \exp(r(o_j)/\alpha)}$$
   This is computationally tractable as it only requires the group-level partition function.

2. Group-Level Policy Distribution: To align the policy with this target, the method defines a policy distribution $q_\theta(o_i | \mathcal{O})$ based on length-normalized log-likelihoods to account for trajectory length bias:
   $$q_\theta(o_i | \mathcal{O}) = \frac{\exp(\varphi(o_i))}{\sum_{j=1}^G \exp(\varphi(o_j))}$$
   where $\varphi(o_i)$ is the average log-probability per token.

3. Distribution-Matching Loss: The policy is aligned to the target by minimizing the divergence between $q_\theta$ and $p$. The authors utilize Mean Squared Error (MSE) instead of KL divergence for numerical stability, preventing gradient explosion when the policy probability approaches zero:
   $$\mathcal{L}_{DM}(\theta) = \frac{1}{G} \sum_{i=1}^G (p(o_i | \mathcal{O}) - q_\theta(o_i | \mathcal{O}))^2$$

4. Unified Objective: The final objective combines the standard GRPO reward maximization term with the distribution-matching regularization:
   $$\mathcal{L}_{DMPO}(\theta) = \mathcal{L}_{GRPO}(\theta) + \lambda \mathcal{L}_{DM}(\theta)$$
   Here, $\mathcal{L}_{GRPO}$ drives exploitation (reward maximization), while $\mathcal{L}_{DM}$ enforces exploration by preventing the policy from collapsing to a single mode within the high-reward region.

3. Key Contributions

A. DMPO Algorithm

The paper presents DMPO as a principled, practical solution to mode collapse. By approximating forward KL at the group level, it achieves mode-covering behavior without requiring a learned partition function or global sampling. The method is compatible with existing GRPO frameworks, requiring only the addition of a single regularization term.

B. MM-NP-Bench

The authors introduce MM-NP-Bench, a multimodal benchmark designed to rigorously evaluate exploration in reasoning models.

• Scope: It extends the text-based NP-Bench to include visual representations of 10 NP-hard combinatorial optimization tasks (e.g., TSP, Graph Coloring, Vertex Cover).
• Dual-Metric Evaluation: To make mode collapse observable, the benchmark uses two metrics:
  • Success Rate (SR): Measures constraint satisfaction (feasibility).
  • Quality Ratio (QR): Measures optimization quality relative to a heuristic baseline.
  • Significance: Mode collapse manifests as high SR but low QR (finding valid but suboptimal solutions).
• Infrastructure: Includes parametric generators, rule-based verifiers, and heuristic solvers to provide near-optimal baselines.

C. Empirical Validation

The paper provides extensive experiments demonstrating that DMPO outperforms strong baselines (GRPO, GSPO, GPG, FlowRL, etc.) across multiple domains.

4. Results

Performance on Optimization Tasks

• MM-NP-Bench (Multimodal): DMPO achieved a 43.1% Quality Ratio, a 12% relative improvement over GRPO (38.4%). It also achieved the highest Success Rate (61.9% vs. 55.7%).
• NP-Bench (Text-only): DMPO achieved a 43.9% Quality Ratio, a 9% relative improvement over GRPO (40.1%).
• Ablation: The method showed robustness to hyperparameter choices ($\lambda \in [0.3, 0.7]$ and $\alpha \in [0.8, 1.2]$).

Generalization and Transfer

• Mathematical Reasoning: Training on optimization tasks with DMPO improved performance on mathematical reasoning benchmarks (AIME, AMC, etc.) by 2.0% compared to GRPO.
• Out-of-Domain Tasks: DMPO demonstrated a 2.3% improvement on out-of-domain visual reasoning benchmarks (LogicVista, MathVista, etc.), suggesting that diversity-preserving training enhances general reasoning capabilities.

Training Dynamics

Analysis of training logs revealed that while GRPO suffers from "length collapse" (degenerating to short, trivial responses) and premature reward plateaus, DMPO sustains continuous improvement and preserves robust reasoning chains (~400 tokens). DMPO maintains lower entropy than GRPO, indicating "structured diversity" (confident execution of different logical strategies) rather than stochastic confusion.

5. Significance and Claims

The paper claims to establish distribution matching as a practical and principled approach to preventing mode collapse in on-policy RL. The authors argue that:

1. Root Cause Identification: Mode collapse in RLVR is fundamentally caused by the mode-seeking nature of reverse KL minimization.
2. Feasible Solution: Forward KL minimization, traditionally considered intractable, can be effectively approximated at the group level to preserve solution diversity.
3. General Benefit: Preserving diversity during training on specific optimization tasks transfers to improved general reasoning capabilities, suggesting that exploration is a critical component of robust reasoning in Large Reasoning Models.

The work positions DMPO not merely as a diversity regularizer but as a mechanism that enables the discovery of higher-quality solutions that mode-seeking objectives miss, thereby enhancing both optimization performance and the robustness of reasoning across modalities.