Beyond Mode Collapse: Distribution Matching for Diverse Reasoning

On-policy reinforcement learning methods like GRPO suffer from \emph{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 (\textbf{D}istribution-\textbf{M}atching \textbf{P}olicy \textbf{O}ptimization), 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.

Reinforcement learning algorithms that training LLMs by trial and error, such as GRPO, often get stuck on a single solution once they find something that works. This limits their ability to explore alternative strategies and find truly optimal answers, especially for complex problems with many possible solutions. To address this, we propose DMPO, a method that encourages LLMs to consider a broader set of high-quality solutions. DMPO does this by constructing a target distribution over promising solutions and guiding the model to align with it, maintaining diversity during learning. We tested DMPO on challenging NP-hard combinatorial problems with exponentially many possibilities but only a few optimal solutions. Compared to standard methods, DMPO consistently finds better solutions and preserves diversity, preventing the model from collapsing onto a single answer. Beyond optimization tasks, DMPO improves performance on math reasoning and out-of-domain challenges, showing that diversity-preserving training generalizes across tasks. By enabling LLMs to explore more broadly and avoid local traps, DMPO helps models make better decisions, learn more effectively, and perform more reliably across a wide range of reasoning tasks.