Beyond Mode Collapse: Distribution Matching for Diverse Reasoning

• URL: http://arxiv.org/abs/2605.19461v1
• Date: Tue, 19 May 2026 07:13:00 GMT
• Title: Beyond Mode Collapse: Distribution Matching for Diverse Reasoning
• Abstract summary: On-policy reinforcement learning methods like GRPO suffer from mode collapse.<n>We show this stems from reverse KL minimization's mode-seeking behavior.<n>We propose DMPO, which prevents mode collapse through principled approximation of forward KL minimization.
• Score: 69.88237286885065
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: On-policy reinforcement learning methods like GRPO suffer from mode collapse: they exhibit reduced solution diversity, concentrating probability mass on a single solution once discovered and ceasing exploration of alternative strategies. We show this stems from reverse KL minimization's mode-seeking behavior, which reinforces the first high-reward trajectory found rather than maintaining a distribution over multiple diverse solutions. We propose DMPO (Distribution-Matching Policy Optimization), which prevents mode collapse through principled approximation of forward KL minimization. DMPO constructs a group level target distribution over sampled trajectories proportional to their rewards, then aligns the policy distribution to this target. This provides mode-covering behavior without requiring sampling from the intractable global target distribution, enabling sustained exploration throughout training. We validate DMPO on NP-hard combinatorial optimization, where exponentially many feasible solutions exist but only a few approach optimality, an ideal testbed for evaluating exploration. DMPO achieves 43.9% Quality Ratio on text-based NP-Bench (vs. GRPO's 40.1%) and 43.1% on vision-based NP-Bench (vs. 38.4%), demonstrating 9% and 12% relative improvements respectively. These gains generalize to mathematical reasoning (+2.0%) and out-of-domain tasks (+2.3%), showing that diversity-preserving training enhances general reasoning capabilities across modalities. Our work establishes distribution matching as a practical, principled approach to preventing mode collapse in on-policy RL, with consistent quality improvements demonstrating sustained exploration across diverse reasoning tasks.

Related papers

• DiffusionOPD: A Unified Perspective of On-Policy Distillation in Diffusion Models [55.01951088768769]
  DiffusionOPD is a new multi-task training paradigm for diffusion models based on Online Policy Distillation (OPD)<n>We show that DiffusionOPD consistently surpasses both multi-reward RL and cascade RL baselines in training efficiency and final performance.
  arXiv  Detail & Related papers  (2026-05-14T16:49:09Z)
• Uniform-Correct Policy Optimization: Breaking RLVR's Indifference to Diversity [23.454044066183027]
  Reinforcement Learning with Verifiable Rewards (RLVR) has achieved substantial gains in single-attempt accuracy (Pass@1) on reasoning tasks.<n>We identify a structural cause for this degradation: common RLVR objectives are indifferent to how probability mass is distributed among correct solutions.<n>We propose Uniform-Correct Policy Optimization (UCPO), a modification to GRPO that adds a conditional penalty on the policy's distribution over correct solutions.
  arXiv  Detail & Related papers  (2026-05-01T03:02:44Z)
• Too Correct to Learn: Reinforcement Learning on Saturated Reasoning Data [55.84428098924793]
  We propose Constrained Uniform Top-K Sampling (CUTS), a parameter-free decoding strategy enforcing structure-preserving exploration.<n>We integrate this into Mixed-CUTS, a training framework synergizing exploitative and exploratory rollouts to amplify intra-group advantage variance.<n> Notably, Mixed-CUTS improves Pass@1 accuracy on the challenging AIME25 benchmark by up to 15.1% over standard GRPO.
  arXiv  Detail & Related papers  (2026-04-20T16:43:28Z)
• Back to Basics: Revisiting Exploration in Reinforcement Learning for LLM Reasoning via Generative Probabilities [10.235183326885794]
  Reinforcement Learning with Verifiable Rewards (RLVR) has emerged as an indispensable paradigm for enhancing reasoning in Large Language Models (LLMs)<n>We analyze this issue from the perspective of sampling probability dynamics, identifying that the standard objective disproportionately reinforces the highest-likelihood paths.<n>We propose a novel Advantage Re-weighting Mechanism (ARM) designed to equilibrate the confidence levels across all correct responses.
  arXiv  Detail & Related papers  (2026-02-05T04:06:55Z)
• SetPO: Set-Level Policy Optimization for Diversity-Preserving LLM Reasoning [50.93295951454092]
  We introduce a set level diversity objective defined over sampled trajectories using kernelized similarity.<n>Our approach derives a leave-one-out marginal contribution for each sampled trajectory and integrates this objective as a plug-in advantage shaping term for policy optimization.<n>Experiments across a range of model scales demonstrate the effectiveness of our proposed algorithm, consistently outperforming strong baselines in both Pass@1 and Pass@K across various benchmarks.
  arXiv  Detail & Related papers  (2026-02-01T07:13:20Z)
• Distribution-Centric Policy Optimization Dominates Exploration-Exploitation Trade-off [34.80019950191864]
  We introduce a textbfdistribution-centric perspective for reinforcement learning.<n>We propose Distribution-Centric Policy Optimization (DCPO), which reformulates entropy regulation as distribution-level regularization.<n>Overall, DCPO replaces sample-levels with distribution-level principles, offering a theoretically grounded and flexible framework for exploration and a stronger EE trade-off.
  arXiv  Detail & Related papers  (2026-01-19T05:20:46Z)
• Stratified GRPO: Handling Structural Heterogeneity in Reinforcement Learning of LLM Search Agents [90.45197506653341]
  Large language model (LLM) agents increasingly rely on external tools such as search engines to solve complex, multi-step problems.<n>The trajectories of search agents are structurally heterogeneous, where variations in the number, placement, and outcomes of search calls lead to fundamentally different answer directions and reward distributions.<n>Standard policy gradient methods, which use a single global baseline, suffer from what we identify and formalize as cross-stratum bias.<n>We propose Stratified GRPO, whose central component, Stratified Advantage Normalization (SAN), partitions trajectories into homogeneous strata based on their structural properties and computes advantages locally within each
  arXiv  Detail & Related papers  (2025-10-07T17:59:13Z)
• MMR1: Enhancing Multimodal Reasoning with Variance-Aware Sampling and Open Resources [113.33902847941941]
  Variance-Aware Sampling (VAS) is a data selection strategy guided by Variance Promotion Score (VPS)<n>We release large-scale, carefully curated resources containing 1.6M long CoT cold-start data and 15k RL QA pairs.<n> Experiments across mathematical reasoning benchmarks demonstrate the effectiveness of both the curated data and the proposed VAS.
  arXiv  Detail & Related papers  (2025-09-25T14:58:29Z)
• Centralized Adaptive Sampling for Reliable Co-Training of Independent Multi-Agent Policies [3.5253513747455303]
  Independent on-policy policy gradient algorithms are widely used for multi-agent reinforcement learning (MARL) in cooperative and no-conflict games.<n>They are known to converge suboptimally when each agent's policy gradient points toward a suboptimal equilibrium.<n>We introduce an adaptive action sampling approach to reduce joint sampling error.
  arXiv  Detail & Related papers  (2025-08-01T20:07:25Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.