💻 computer science

Inference-Time Diversity in RL-Trained Lean Theorem Provers: A Diagnostic Study

This study reveals that RL-trained Lean theorem provers suffer from inference-time mode collapse, which can be effectively mitigated by applying a fixed schedule of diverse tactic skeletons to recover significant performance gains, whereas prompt paraphrasing and irrelevant comments prove ineffective or detrimental.

Original authors: Zachary Burton

Published 2026-05-19

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Zachary Burton

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine a brilliant student who has been trained to solve very difficult math puzzles using a special computer language called Lean. This student, let's call them "DeepSeek," is incredibly smart but has a strange habit: when asked to solve a puzzle, they get stuck in a mental rut.

This paper is a diagnostic report on why this happens and how to fix it. Here is the story of what the researchers found, told in everyday terms.

1. The Problem: The "Broken Record" Effect

The researchers noticed that when they asked DeepSeek to solve the same math puzzle over and over again (just with slightly different random seeds, like rolling dice), the student kept trying the exact same first step.

The Analogy: Imagine you are trying to open a locked door. You try the key in the lock. It doesn't work. You try again. It still doesn't work. You try a third time, and a hundredth time. Even if you try 64 times, you are just jamming the same key into the same lock in the same way. You aren't exploring other keys or other ways to open the door.
The Result: In the study, doubling the number of attempts from 32 to 64 didn't solve a single new puzzle. The student had "mode collapsed"—they were stuck on a narrow path and couldn't think of anything else.

2. The Fix: Giving a "First Step" Hint

The researchers realized that the student wasn't lazy; they were just too stuck in their routine. So, they tried a new trick: before asking the student to solve the puzzle, they gave them a specific instruction on how to start.

The Analogy: Instead of just saying, "Solve this puzzle," the researchers said, "Start by simplifying the numbers," or "Start by introducing a new variable." It's like giving the student a specific key to try first, rather than letting them guess.
The Result: This simple nudge broke the rut. The student suddenly started solving many more puzzles. The "rut" wasn't because the student was bad at math; it was because they were too rigid in how they started.

3. The "Why": It's a Training Issue, Not a Brain Issue

The researchers wanted to know: Is the student naturally bad at math, or did the training make them rigid?

The Experiment: They tested the student before they got the special math training (the "Base" model).
The Finding: The untrained student couldn't solve any puzzles at all. They just gave up or repeated the question.
The Conclusion: The special training (Reinforcement Learning) is what taught the student how to solve puzzles in the first place. But, that same training is also what made them rigid and stuck in a rut. The training created the skill, but it also created the "broken record" problem.

4. Who Else Has This Problem?

The researchers tested other "students" (different AI models) to see if this was a universal problem.

The "SFT" Student: One student was trained differently (Supervised Fine-Tuning). This student was actually better at solving puzzles overall, but they didn't have the "rut" problem. When the researchers gave them hints, it actually made them slightly worse. They were already flexible enough to find their own path.
The "RL" Students: The students trained with the special "Reinforcement Learning" method (like DeepSeek) were the ones who got stuck. They needed the "first step" hints to break out of their narrow thinking.

5. The "First Move" Evidence

To prove this wasn't just a fluke, the researchers looked at the actual "first moves" the students made.

The Finding: Out of 64 attempts on a single puzzle, the "rutty" student used the exact same first move in nearly half the cases. They were incredibly predictable.
The Metaphor: If you asked a human to write a story 64 times, they might start with "Once upon a time," "The sun was shining," or "It was a dark night." But this AI student would start with "Once upon a time" 32 times, and "The sun was shining" 30 times. They had almost no variety in their opening.

Summary

The paper concludes that Reinforcement Learning (a specific way of training AI) creates a double-edged sword:

It teaches the AI how to solve formal math proofs (which it couldn't do before).
But it also makes the AI so focused on a few "winning" strategies that it stops exploring new ones.

The solution isn't to give the AI more computing power or more time to think; it's to give it structural guidance—a simple nudge on how to start the proof. This forces the AI to break its routine and explore new paths, unlocking solutions it was previously too stubborn to find.

Technical Summary: Inference-Time Diversity in RL-Trained Lean Theorem Provers

Problem Statement

Recent neural theorem provers, particularly those trained with Reinforcement Learning (RL) and verifiable rewards (RLVR), have achieved strong performance on formal proof benchmarks like miniF2F. The prevailing assumption is that sufficient RL training and increased inference-time compute (via i.i.d. sampling) allow models to internalize the structural priors necessary for formal proofs.

This paper challenges that assumption, diagnosing a phenomenon of mode collapse at inference time in RL-trained Lean provers. Specifically, the authors observe that increasing the sampling budget ( $k$ ) yields diminishing returns far earlier than expected. On the miniF2F-test benchmark with DeepSeek-Prover-V1.5-RL, doubling the sampling budget from $k=32$ to $k=64$ produces zero additional solved theorems (stuck at 42/244). The model converges on a narrow set of proof strategies and fails to explore alternatives, even when provided with additional stochastic samples.

Methodology

The study employs a controlled diagnostic framework using DeepSeek-Prover-V1.5-RL and its non-RL base variant, alongside cross-model comparisons with DeepSeek-Prover-V2-7B (RL) and Goedel-Prover (SFT-only).

1. The Skeleton Schedule Intervention

To test if the plateau is due to a lack of structural diversity rather than a lack of compute, the authors introduce a fixed, deterministic schedule of 15 tactic skeletons.

Structure: A $15 \times 8 = 120$ grid combining 15 hand-selected Lean tactic skeletons (e.g., simp, intro, constructor, aesop) with 8 goal-hint comments (natural language suggestions).
Execution: For a budget $k$ , the model attempts the first $k$ entries of this grid. This forces the model to start proofs with specific strategic openings rather than relying on random sampling.
Ablations: To isolate the source of improvement, the authors compare:
- B-RL (Skeleton): The structured 15-skeleton schedule.
- C1 (Paraphrase): 16 semantically equivalent instruction paraphrases (no tactic structure).
- C2 (Irrelevant Comments): 16 content-free Lean comments.
- C3 (NL-only): A dedicated condition using only natural language hints without tactic prefixes.

2. Difficulty Stratification

The authors define an empirical formalization-difficulty score ( $\sigma$ ) using a leave-one-out construction. For a target theorem, $\sigma$ counts how many other baseline models (Goedel-Prover, V2-7B) solve it without intervention. This stratifies theorems into buckets: Frontier (unsolved by others), Hard, Medium, Easy, and Trivial.

3. Direct Distributional Measurement

To move beyond behavioral plateaus, the authors analyze the first-tactic head diversity across 13,159 stochastic samples. They count the number of distinct leading tactic identifiers (e.g., have, rw, norm_num) generated per theorem.

Key Results

1. Breaking the Plateau via Structural Guidance

Baseline Collapse: DeepSeek-Prover-V1.5-RL with i.i.d. sampling solves 38 theorems at $k=16$ , 42 at $k=32$ , and remains at 42 at $k=64$ .
Skeleton Efficacy: The fixed skeleton schedule breaks this plateau, solving 55 theorems at $k=16$ (+45% relative improvement) and continuing to find new proofs up to $k=64$ (60 total).
Ablation Findings:
- Structure vs. Diversity: Instruction paraphrases (C1) yield the same pass@16 as the baseline (38/244), while irrelevant comments (C2) degrade performance to 25/244. This confirms gains are driven by structural guidance, not mere prompt diversity.
- NL vs. Tactics: The NL-only condition (C3) solves 48/244, intermediate between baseline and full skeletons. This indicates that natural language hints recover roughly 58% of the gap, while literal tactic prefixes contribute the remainder.

2. Localization of the Gain

The structural gain is concentrated in the Easy and Trivial buckets (theorems solvable by other open-source provers but missed by V1.5-RL's i.i.d. sampling).

In the combined Easy+Trivial column, the skeleton intervention yields a net gain of +15, while paraphrases yield 0 and irrelevant comments yield -14.
No perturbation unlocks the "Frontier" bucket (theorems unsolved by any baseline), suggesting the intervention remedies underexploration of known solvable problems rather than expanding the frontier of capability.

3. RL-Specificity of Mode Collapse

Base Model Contrast: DeepSeek-Prover-V1.5-Base (non-RL) proves zero theorems under any condition (i.i.d. or skeletons). This establishes that RL is the locus of both proof capability and the subsequent mode collapse. The base model lacks the capability to generate valid proofs; the RL stage creates the capability but collapses it into a narrow policy.
Cross-Model Contrast:
- V2-7B (RL): Shows a flat aggregate gain (+1) but unlocks +3 frontier theorems unreachable by any i.i.d. baseline.
- Goedel-Prover (SFT): Shows a negative gain (-10.0 $\pm$ 4.4) with skeletons. SFT-trained models appear to already explore a broad set of tactic stems, making the fixed skeleton schedule redundant or restrictive.

4. Direct Evidence of Collapse

Distributional analysis reveals that the median theorem receives only 2 distinct first-tactic heads across 64 samples, with 49% of the benchmark receiving a deterministic single strategic opening. The observed Shannon entropy is $\le 8\%$ of the theoretical capacity, confirming severe mode collapse.

Significance and Claims

The paper reframes the narrative from "tactic skeletons improve performance" to "RL-trained provers severely underexplore the tactic space at inference time."

RL as Dual Agent: The authors claim that RL plays a dual role: it is the stage that creates proof capability (as the base model has none) and simultaneously the stage that collapses that capability onto a narrow policy.
Inference-Time Scaling: The results suggest that simply scaling inference compute via i.i.d. sampling has diminishing returns much earlier than assumed. Structural diversity (varying proof strategies via prompts) is a more efficient axis for scaling.
Remediability: The mode collapse is specific to RL-trained models. SFT-trained models do not exhibit the same recoverable collapse via structural hints; in fact, they suffer from it. This implies that future interventions (training-time entropy regularization or inference-time skeleton retrieval) must navigate the trade-off between creating capability and maintaining diversity.

The study concludes that the "mode collapse" is a real, pipeline-contingent phenomenon inherent to the RL fine-tuning process in formal verification domains, and that fixed structural interventions can effectively recover the lost exploration without requiring additional training.