Abhijeet Sinha

Projects

MIPS: Mode Collapse in Large Outcome Spaces

Keeping RL diverse when there are far too many possible outcomes to count, by estimating each outcome's probability from a single sample.

In one sentence: inverse probability scaling needs to know how likely each outcome is. When outcomes never repeat, counting fails, so MIPS learns a small “backward” model that recovers an outcome’s probability from just the one route that produced it.

Background: the fix that needs repeats

In earlier work we showed why reinforcement learning collapses onto a few outcomes, and fixed it with inverse probability scaling (IPS): divide each reward by how likely the model is to produce that outcome. Popular outcomes stop snowballing, and the model learns to produce each good outcome in proportion to its reward.

The practical version, IPS-GRPO, estimates “how likely” by counting. It samples a group of, say, 64 outputs and checks how often each distinct one appears. That works well when the same outcomes come up again and again.

The problem: in big spaces, nothing repeats

Now think about designing molecules from 350 building blocks and 17 reaction types, or building evolutionary trees for 10 species. The number of possible outcomes is astronomically larger than any group you can afford to sample. Each outcome shows up once, every count is the same, and IPS has nothing to work with.

Figure 1 · Why counting breaks down

Small outcome spaceHuge outcome spaceAABACBABA: 4/8B: 3/8C: 1/8o1o2o3o4o5o6o7o8every one:1/8Counts differ, so they tell uswhich outcomes are over-sampled.Nothing repeats. Every count is 1/8,so IPS sees no difference at all.
Frequency-based IPS needs outcomes to show up more than once in a group. When there are far more possible outcomes than samples in the group, every outcome tends to appear exactly once, and the frequency estimate tells you nothing.

A tempting shortcut is to use the probability of the actual sequence of steps the model took, which is easy to compute. But in these problems the same final object can be built in many different orders, so one route’s probability can be far smaller than the outcome’s.

Figure 2 · Many routes, one outcome

startoutcome o3 of the 6 routes to the same cellProbability of the outcome= sum of the probabilities of all 6 routesIn real problems (molecules, trees) thenumber of routes explodes, so the sumcannot be computed directly.The probability of the one route you sampledis not the probability of the outcome.
In a constructive task, the same final object can be built in many orders. Moving right twice and up twice can be done in 6 different ways, and all of them land on the same cell.

The idea: work backwards from the outcome

Probability theory gives a neat identity. For any route τ that produces an outcome o:

**P(route) = P(outcome) × P(this particular route outcome)**

Rearranged, the outcome’s probability is the route’s probability divided by “how likely this route is, among all routes to the same outcome”. The first number comes straight from the policy. For the second, MIPS trains a small backward model. Given a finished object, it learns how the current policy tends to build it, one step back at a time. It is trained on the policy’s own samples with ordinary maximum likelihood.

Figure 3 · How MIPS estimates the probability

Sample oneroute τ thatbuilds oPolicy:how likely was this route?P(τ)Backward model: of allroutes to o, how likelyis this one? c(τ | o)Estimated outcome probabilityP(o) ≈ P(τ) ÷ c(τ | o)reward ÷ P(o) → standard GRPO
One identity does the work: the probability of a route equals the probability of its outcome times the chance of taking that particular route to reach it. MIPS learns the second term with a small backward model, then divides it out. One sample is enough, and nothing needs to repeat.

With that estimate, every sampled outcome gets its own probability, even if it appears only once. The scaled reward then goes into the unchanged GRPO update. GFlowNets use a similar backward model for a different purpose. MIPS stays inside a standard RL pipeline and only uses the backward model to correct the reward.

What we found

Hypergrids: a test where we know the right answer. In a grid world, we can compute exactly how far the trained model is from the ideal reward-proportional distribution. We swept 8 grid sizes and 8 group sizes (192 settings in total). GRPO is far from the target almost everywhere. Frequency-based IPS is good on small grids but breaks down as the grid grows, and bigger groups only delay the breakdown. MIPS stays close to the target across the range. When we remove the backward model, the error jumps sharply (for example, a Jensen–Shannon divergence of 0.004 versus 0.56), which confirms that the backward model is what makes it work.

Figure 4 · The trend on Hypergrids (schematic)

GRPOIPS-GRPO (frequency counting)MIPS-GRPOmismatch with targetoutcome space small compared with the groupoutcome space huge
A simplified sketch of what the paper measures over 192 grid and group sizes (exact values are in Figure 2 and Tables 4–5 of the paper). Making the group bigger helps frequency-based IPS, but only delays the point where it breaks down. MIPS stays close to the target across the whole range.

Molecular synthesis. The agent builds a molecule step by step through chemical reactions, and is rewarded by a protein-binding score (sEH). GRPO collapses to a single molecule. Frequency-based IPS barely helps, because molecules almost never repeat. MIPS-GRPO produces about 45,800 unique molecules out of 50,000 samples and finds nearly 4× as many distinct high-scoring families as RGFN, a GFlowNet built for this task. Its average quality is the same or slightly better.

Figure 5 · Molecular synthesis

GRPOIPS-GRPORGFN (GFlowNet)MIPS-GRPO04080120160count001870Distinct high-reward modes0025149High-reward scaffolds
Distinct high-scoring molecule families found in 50,000 samples, averaged over 3 seeds (Table 2 of the paper). A scaffold is the core ring structure of a molecule, so a higher count means the molecules really are different and not small variations of one design.

Evolutionary (phylogenetic) trees. The agent builds a tree for 10 species by repeatedly joining branches, and is rewarded by how well the tree explains the DNA data. Here we want broad coverage, and we also want better trees to be sampled more often.

Method Distinct trees in 1M samples Probability tracks quality (correlation)
GRPO 1 n/a
IPS-GRPO 263,532 0.14
PhyloGFN 1,000,000 0.84
MIPS-GRPO 1,000,000 0.97

MIPS matches the specialised GFlowNet on coverage. It also has the strongest link between how often a tree is sampled and how good the tree is.

Limits and what’s next

MIPS is only as good as its backward model, and the current theory assumes finite, step-by-step construction processes with positive rewards. A natural next step is outcomes that are continuous or only “similar” rather than identical. That is the setting our DARS work addresses with a different tool.