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

TL;DR

This paper introduces a filtering-based training method for large language models that balances reasoning accuracy and diversity by controlling the divergence used during training, outperforming reinforcement learning approaches.

Contribution

It proposes a novel approach starting from an explicit target distribution and uses the $oldsymbol{ extalpha}$-divergence to balance precision and diversity in reasoning tasks.

Findings

Achieves state-of-the-art coverage-precision trade-off on theorem proving benchmark.

Outperforms prior methods in coverage while maintaining high precision.

Demonstrates the effectiveness of filtering and divergence control in LLM training.

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 $\alpha$-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.