Beyond Mode Collapse: Distribution Matching for Diverse Reasoning

Fuente: arXiv

Saved in:

Bibliographic Details
Main Authors:	Li, Xiaozhe, Li, Yang, Fang, Xinyu, Ding, Shengyuan, Li, Peiji, Chen, Yongkang, Ma, Yichuan, Lyu, Tianyi, Li, Linyang, Lin, Dahua, Guo, Qipeng, Liu, Qingwen, Chen, Kai
Format:	Preprint
Published:	2026
Subjects:	Artificial Intelligence
Online Access:	Acceder al recurso
Tags:	Add Tag No Tags, Be the first to tag this record!

_version_	1866914579537199104
author	Li, Xiaozhe Li, Yang Fang, Xinyu Ding, Shengyuan Li, Peiji Chen, Yongkang Ma, Yichuan Lyu, Tianyi Li, Linyang Lin, Dahua Guo, Qipeng Liu, Qingwen Chen, Kai
author_facet	Li, Xiaozhe Li, Yang Fang, Xinyu Ding, Shengyuan Li, Peiji Chen, Yongkang Ma, Yichuan Lyu, Tianyi Li, Linyang Lin, Dahua Guo, Qipeng Liu, Qingwen Chen, Kai
contents	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.
format	Preprint
id	arxiv_https___arxiv_org_abs_2605_19461
institution	arXiv
publishDate	2026
record_format	arxiv
spellingShingle	Beyond Mode Collapse: Distribution Matching for Diverse Reasoning Li, Xiaozhe Li, Yang Fang, Xinyu Ding, Shengyuan Li, Peiji Chen, Yongkang Ma, Yichuan Lyu, Tianyi Li, Linyang Lin, Dahua Guo, Qipeng Liu, Qingwen Chen, Kai Artificial Intelligence 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.
title	Beyond Mode Collapse: Distribution Matching for Diverse Reasoning
topic	Artificial Intelligence
url	https://arxiv.org/abs/2605.19461