Expected Return Causes Outcome-Level Mode Collapse

In one sentence: standard reinforcement learning (RL) is built to collapse onto a few answers, not just prone to it. Dividing each reward by how often its answer is already produced removes the cause, and the model learns to give every good answer in proportion to how good it is.

The problem: many good answers, one chosen

Many tasks we now train with RL have lots of correct answers. A chemistry model can propose thousands of molecules that bind a target. A reasoning model can write many valid proofs. An assistant can phrase a reply in many helpful ways. In these settings we want a model that keeps all the good options alive, not one that finds a single favourite.

In practice, RL-trained models do the opposite. After training they repeat the same few outputs. This is called outcome-level mode collapse. The usual explanation is that the model “did not explore enough”, so the usual fixes are exploration bonuses, entropy bonuses or a stronger pull towards the original model. These help for a while, but the collapse comes back as training goes on.

Why it happens: a rich-get-richer loop

We show that the cause is the training objective itself. Standard RL maximises expected return, the average reward of what the model produces. When you work out how this objective updates the model, the push each outcome receives is

its current probability × how much better than average it is.

That first factor is the problem. An outcome that is already a little more common gets a bigger push, which makes it more common, which gets it an even bigger push.

Figure 1 · The rich-get-richer loop

The paper proves this exactly. In an idealised setting with perfect exploration and no noise, the ratio between any two outcomes’ probabilities grows or shrinks exponentially, and the more common outcome always wins. Exploration bonuses and regularisation can slow this down, but they cannot stop it, because the pressure to collapse is built into the objective.

Figure 2 · Simulated training, three outcomes

The fix: divide by probability

If the problem is a probability factor in the update, the fix is to cancel it. Inverse probability scaling (IPS) divides each outcome’s reward by the probability that the model currently produces that outcome:

• common outcomes have their reward divided by a large number, so they stop snowballing;
• rare but good outcomes have their reward divided by a small number, so they get a fair share of the learning signal.

With this change, the stable end point of training is no longer “one outcome takes everything”. It becomes a distribution where each outcome’s probability is proportional to its reward. Two equally good answers end up equally likely, and an answer that is twice as good ends up twice as likely.

IPS-GRPO: a drop-in change

GRPO is a popular RL algorithm for language models that samples a group of answers for each question and compares them. That makes IPS easy to add: count how often each distinct answer appears in the group, and divide its reward by that frequency. Nothing else changes. There is no extra network, no new loss, and it works with the usual clipping and KL settings.

Figure 3 · What IPS does inside one group

1. Sample a group of answers, as GRPO already does.
2. Group identical answers and count how often each one appears.
3. Divide each reward by its answer’s frequency, with a small floor so that rare answers don’t produce huge values.
4. Run the normal GRPO update with these scaled rewards.

What we found

A grid with four equal peaks. In a small grid world whose reward has several equally good regions, both GRPO and FlowRL (a flow-based method) end up stopping almost only at one of them. IPS-GRPO covers all of them and gets about 10× closer to the ideal distribution (an L1 distance of 0.003, compared with 0.028 for both baselines).

Reasoning with many valid answers (HypoSpace). These are puzzles that have a known, finite set of correct answers: causal graphs that fit some data, 3D block structures that match a top-down view, and Boolean rules that explain observations. We fine-tuned language models on them and measured how many of the valid answers each model finds. GRPO and FlowRL stop discovering new answers early. IPS-GRPO keeps finding new ones as you sample more.

Figure 4 · Results on HypoSpace

Drug discovery. We replaced the policy update in REINVENT, a leading RL method for chemical language models, with IPS-GRPO. The goal was to find as many unique high-scoring molecules as possible within a fixed number of scoring calls.

Why it matters

The main point is a change of diagnosis. Collapse in RL is not mainly a failure of exploration, so it can’t be fully fixed by exploring harder. It comes from what the objective rewards. Once that factor is removed, regularisers such as KL go back to their intended job, which is keeping the model close to a trusted starting point, instead of fighting collapse. This idea is also the basis for our follow-up work: MIPS handles outcome spaces too large for counting, and DARS handles continuous outputs such as images and free text.