Expected Return Causes Outcome-Level Mode Collapse in Reinforcement Learning and How to Fix It with Inverse Probability Scaling
cs.LG, cs.AI
Submitted: 2026-01-29
Updated: 2026-01-29
Journal ref: Proceedings of the 43rd International Conference on Machine Learning (ICML 2026), Proceedings of Machine Learning Research, vol. 306, pp. 113750-113765, 2026
Code: https://github.com/schwallergroup/saturn
License: http://creativecommons.org/licenses/by/4.0/
The gist: Many reinforcement learning (RL) problems admit multiple terminal solutions of comparable quality, where the goal is not to identify a single optimum but to represent a diverse set of high-quality
Terminology
Abstract
Many reinforcement learning (RL) problems admit multiple terminal solutions of comparable quality, where the goal is not to identify a single optimum but to represent a diverse set of high-quality outcomes. Nevertheless, policies trained by standard expected return maximization routinely collapse onto a small subset of outcomes, a phenomenon commonly attributed to insufficient exploration or weak regularization. We show that this explanation is incomplete: outcome level mode collapse is a structural consequence of the expected-return objective itself. Under idealized learning dynamics, the log-probability ratio between any two outcomes evolves linearly in their reward difference, implying exponential ratio divergence and inevitable collapse independent of the exploration strategy, entropy regularization, or optimization algorithm. We identify the source of this pathology as the probability multiplier inside the expectation and propose a minimal correction: inverse probability scaling, which removes outcome-frequency amplification from the learning signal, fundamentally changes the learning dynamics, and provably yields reward-proportional terminal distributions, preventing collapse in multimodal settings. We instantiate this principle in Group Relative Policy Optimization (GRPO) as a drop-in modification, IPS-GRPO, requiring no auxiliary models or architectural changes. Across different reasoning and molecular generation tasks, IPS-GRPO consistently reduces outcome-level mode collapse while matching or exceeding baseline performance, suggesting that correcting the objective rather than adding exploration heuristics is key to reliable multimodal policy optimization.
Sources
- Concrete Problems in AI Safety
- Constitutional AI: Harmlessness from AI Feedback
- Maximum Entropy Reinforcement Learning via Energy-Based Normalizing Flow
- Reasoning with Exploration: An Entropy Perspective
- The Entropy Mechanism of Reinforcement Learning for Reasoning Language Models
- PreND: Enhancing Intrinsic Motivation in Reinforcement Learning through Pre-trained Network Distillation
- Winner Takes It All: Training Performant RL Populations for Combinatorial Optimization
- Saturn: Sample-efficient Generative Molecular Design using Memory Manipulation
- Generative Molecular Design with Steerable and Granular Synthesizability Control
- KL-Regularized Reinforcement Learning is Designed to Mode Collapse
- The Curious Case of Neural Text Degeneration
- AlphaEvolve: A coding agent for scientific and algorithmic discovery
- Does Writing with Language Models Reduce Content Diversity?
- Direct Preference Optimization: Your Language Model is Secretly a Reward Model
- DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models
- Program Synthesis Through Reinforcement Learning Guided Tree Search
- Diverse Preference Learning for Capabilities and Alignment
- Outcome-based Exploration for LLM Reasoning
- Learning to summarize from human feedback
- Beyond Reverse KL: Generalizing Direct Preference Optimization with Diverse Divergence Constraints
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks