Whatever Remains Must Be True: Filtering Drives Reasoning in LLMs, Shaping Diversity

AI summary

Overview

Research area: machine learning, large language model reasoning, reinforcement learning from verifiable rewards, distribution matching, and formal theorem proving. Technical level: Advanced. One-sentence scope: The paper argues that reinforcement learning with verifiable rewards loses diversity because it implicitly optimizes a mode-seeking Reverse KL objective, and proposes alpha-DPG, an alpha-divergence distributional matching method that explicitly targets a filtered distribution over correct answers and trades off precision against coverage on a Lean theorem-proving benchmark.

What This Paper Is About

Reinforcement Learning with Verifiable Rewards (RLVR) is widely used to tune LLMs for reasoning, but tuned models often become less diverse. The paper argues this happens because RLVR implicitly optimizes the mode-seeking Reverse KL divergence, concentrating probability on a narrow set of high-reward solutions. The goal is to define the desired target distribution explicitly, filter out incorrect answers while preserving the relative probabilities of correct ones, and approximate that target with a controllable family of divergences called alpha-DPG.

Key Contributions

1. Introduces the Distributional Matching with Verifiable Rewards (DMVR) framework, which trains a model by approximating an explicitly defined verifier-based target distribution.
2. Clarifies how the implicit dynamics of RL-based methods lead to reduced diversity, showing that RLVR optimizes toward a filtered version of the base distribution in a mode-seeking way.
3. Proposes alpha-DPG, based on alpha-divergences, to smoothly interpolate between Forward KL, which is mass-covering, and Reverse KL, which is mode-seeking, thereby controlling the precision-diversity trade-off.
4. Shows on the Lean theorem-proving benchmark that alpha-DPG achieves results along a Pareto coverage-precision frontier, with low alpha values achieving the best coverage among all considered methods.

Main Findings

• RLVR is implicitly mode-seeking. The paper reproduces the result that the gradient of KL-Control's objective is proportional to the Reverse KL divergence to a target distribution that favors correct answers. For any fixed beta greater than 0, RLVR optimizes Reverse KL to that target, which concentrates mass on high-reward regions and sacrifices diversity. As beta approaches 0 from above, the target converges to the ideal filtered distribution.
• The ideal target distribution filters incorrect answers. The proposed target is p_x(y) proportional to pi_base(y|x) times v(y,x), where v(y,x) is a binary verifier. This distribution always outputs correct solutions and preserves the relative probabilities of correct answers from the base model. It is the single distribution that satisfies correctness and is closest to the base model in D_KL.
• Alpha-DPG unifies prior methods. The alpha-divergence family interpolates between Forward KL as alpha approaches 0 and Reverse KL as alpha approaches 1, with the squared Hellinger distance at alpha equal to 0.5. This unifies RLVR, KL-DPG, and Rejection Sampling Fine-Tuning under one umbrella. The alpha parameter directly controls the precision-diversity trade-off.
• Lean theorem-proving results favor the Pareto frontier. On a Lean benchmark, alpha-DPG models lie on or near the Pareto frontier between precision measured by pass@1 and coverage measured by pass@256. Low alpha settings, such as alpha equal to 0.25, achieve the highest coverage while still improving precision over SFT. Increasing alpha improves precision, reaching parity with RL-based methods at large values, such as alpha greater than or equal to 0.995, while typically retaining higher coverage.
• Specific model comparisons. Alpha-DPG with alpha equal to 0.999 generally dominates GRPO and other pure RL-based techniques, except ReMax, which has better coverage but somewhat lower precision. Alpha-DPG with alpha equal to 0.25 dominates the base model and diversity-preserving baselines such as Pass@k training and GRPO with KL regularization, achieving the best pass@256 performance. Rw-Ulkly starts with higher pass@1, but alpha equal to 0.25 outperforms it for k greater than or equal to 4. The paper also reproduces earlier findings that GRPO starts with a much higher pass@1 score, but the base model overpasses it at pass@16, and no single model dominates all others across all values of k.
• Problem difficulty transitions reveal a trade-off. GRPO and alpha-DPG with alpha equal to 0.999 improve many medium or hard problems into easy ones, but also degrade many hard problems so they become unsolved. Alpha-DPG with alpha equal to 0.25 and GRPO with High-KL are more conservative: they improve sample efficiency on fewer problems, but harder problems remain solvable, with only three problems becoming unsolvable.
• Diversity correlates with coverage, not precision. Higher diversity in tactics and premises used in Lean proofs correlates with improved pass@256 performance, whereas it is anticorrelated with pass@1. Perplexity analysis shows generated sequences are already highly probable under the base model, with very similar perplexities across models. For one problem, GRPO collapsed and produced 16 identical sequences, which were also highly probable under the base model.

Methodology in Plain English

The researchers start from a pre-trained LLM and define a target distribution by taking the base model's probability for each answer and keeping only answers that a verifier marks as correct. This filtered distribution preserves the relative probabilities of all correct answers, so it does not artificially favor a narrow subset. Because sampling directly from this target is infeasible due to an unknown normalization constant, they estimate that constant using importance sampling and the acceptance rate of the verifier when sampling from the base model. They then train an autoregressive policy using Distributional Policy Gradient algorithms, specifically f-DPG with alpha-divergences, which lets them choose how mode-seeking or mass-covering the training objective is. The alpha parameter interpolates between Forward KL and Reverse KL. They clip the pseudo-reward to a maximum M equal to 10 and use a leave-one-out per-context average as the baseline. The partition function is computed online using sampled responses for each problem, so it adds no extra computational cost relative to baselines. Experiments use DeepSeek-Prover-V1.5-SFT, a 7B parameter model, on 10K solvable Lean problems extracted from the Lean Workbook dataset, with 200 problems kept unseen as a test set. Training uses Lean4 and Mathlib4 version lean4:v4.9.0, runs on a single node of 4xA100 with 28 CPUs for parallel proof assessment, uses 512 generated sequences per step for 200 iterations, approximately 3 epochs, a maximum response length of 1024 tokens, and evaluation with temperature T equal to 1 and nucleus sampling p equal to 0.99. Baselines include GRPO, Dr. GRPO, High-KL with beta equal to 0.1, Rw-Ulkly with beta equal to 0.25, Pass@k training, GPG, ReMax, RLOO, and Base-SFT.

Why This Matters

The paper reframes RLVR as a form of filtering and reweighting of the base model's existing distribution rather than a process that creates fundamentally new reasoning capabilities. This helps explain why RLVR models often become accurate but less diverse, and it offers a principled way to control the trade-off between precision and coverage by tuning alpha. The work also connects distribution matching, f-divergences, and formal theorem proving in a single framework, which can inform future LLM post-training and inference scaling strategies.

Real-world applications:

• Formal theorem proving with proof assistants such as Lean, Coq, and Isabelle, where multiple valid proofs may exist and diversity helps solve harder theorems.
• AI for mathematics and mathematical discovery, where exploring a wide range of candidate proofs or solutions can be more valuable than producing one high-probability answer.
• Code generation and problem-solving benchmarks that use pass@k, where coverage across many samples matters for test-time scaling.
• Training LLMs for any verifiable task that requires robust sampling and multiple correct solutions, such as formal verification or structured reasoning.

Industry relevance: The method targets LLM post-training pipelines that use RLVR, RLHF, or proof assistant feedback. It offers a practical way to reduce mode collapse, preserve diversity, and improve coverage at inference time without requiring an additional critic model. The alpha parameter gives engineers a direct knob for balancing precision and diversity, which is relevant for products that rely on sampling many candidates, such as coding assistants, math solvers, and formal verification tools.

Future Directions

• Apply alpha-DPG beyond Lean to other verifiable domains, such as code generation, informal mathematics, and other proof assistants like Coq or Isabelle.
• Study how to choose or adapt alpha automatically for different tasks, models, and sampling budgets rather than tuning it manually.
• Improve partition function estimation, including using more samples, and further analyze the formal properties of alpha-divergences when the target distribution has restricted support over verifiable outputs only.
• Investigate whether these methods discover genuinely new solutions or mainly reweight solutions already likely under the base model, and test combinations with other diversity-preserving baselines such as Pass@k training and Rw-Ulkly.

Target Audience

This paper benefits machine learning researchers working on LLM reasoning, reinforcement learning, and distribution matching, as well as practitioners who tune LLMs with verifiable rewards. It is also relevant to researchers in formal verification and automated theorem proving, especially those using Lean, Coq, or Isabelle, and to engineers building post-training pipelines where diversity and coverage are important. Advanced students with background in reinforcement learning and probabilistic modeling will gain the most from the technical details, while the high-level framing is accessible to a broader AI audience.

Authors’ abstract

Reinforcement Learning (RL) has become the de facto standard for tuning LLMs to solve tasks involving reasoning. However, growing evidence shows that models trained in such way often suffer from a significant loss in diversity. We argue that this arises because RL implicitly optimizes the "mode-seeking" or "zero-forcing" Reverse KL to a target distribution causing the model to concentrate mass on certain high-probability regions of the target while neglecting others. In this work, we instead begin from an explicit target distribution, obtained by filtering out incorrect answers while preserving the relative probabilities of correct ones. Starting from a pre-trained LLM, we approximate this target distribution using the $α$-divergence family, which unifies prior approaches and enables direct control of the precision-diversity trade-off by interpolating between mode-seeking and mass-covering divergences. On a Lean theorem-proving benchmark, our method achieves state-of-the-art performance along the coverage-precision Pareto frontier, outperforming all prior methods on the coverage axis.