New method improves reinforcement learning for many distinct solutions
MIPS: Mitigating Mode Collapse in Reinforcement Learning for Large Outcome Spaces
Machine LearningArtificial Intelligence
Summary
Sometimes, when teaching computers to solve problems with many good answers, they only learn a few solutions and ignore the rest. This problem is called mode collapse. The authors came up with a new way called Marginal IPS (MIPS) that helps the computer properly weigh all possible solutions, even when there are many and they’re hard to count. They tested it on different tasks like molecule building and family tree predictions, where it kept a good variety of strong solutions while older methods failed.
What this means in practice
- •For drug discovery teams: Generate diverse high-quality molecular structures that meet complex design goals during computational drug design tasks.
- •For bioinformatics engineers: Improved phylogenetic inference by sustaining a wide range of plausible evolutionary trees during analysis workflows.
Authors
Armaan Sandhu, Abhijeet Sinha, Dianbo Liu
Abstract
Expected-reward maximization in reinforcement learning concentrates probability mass on a few outcomes, even in tasks that require many distinct, high-quality solutions, a failure known as mode collapse. Inverse Probability Scaling (IPS) counters this in group-based policy gradients such as GRPO by dividing each reward by the outcome's probability, estimated from its frequency within the sampled group. When the outcome space is large relative to the group, outcomes rarely repeat and these estimates become uninformative. Since an outcome can be constructed along many trajectories, its probability is a sum over all of them, and computing it directly is intractable. We propose Marginal IPS (MIPS), which performs marginal evaluation by combining the sampled trajectory with a learned model of the trajectories that lead to each outcome. On Hypergrids, where the exact sampling distribution is computable, MIPS stays near the reward-distribution target as the outcome space grows relative to the rollout group, while frequency-based IPS degrades. In molecular synthesis and phylogenetic inference, MIPS sustains diverse, high-quality sampling where GRPO and frequency-based IPS collapse, showing that MIPS serves as a reliable solution to mode collapse in RL.