Large language models improved for better automated theorem search
Direct Optimization of Generators for Search in Automated Theorem Proving
Artificial Intelligence
Summary
Automated theorem proving uses computer programs to find mathematical proofs. Large language models (LLMs) help by suggesting proof steps, but usually guide a search process rather than generate proofs in one try. The paper studies improved training methods that better align the LLM’s learning with how proofs are actually searched, including during complex tree searches involving many trial paths. The authors introduce new loss functions that account for the search process and prove these lead to better success rates in real tests with a theorem prover called Lean.
What this means in practice
- •For software engineers: Enhance automated theorem proving systems to find correct proofs more often by using search-aware trained language models.
- •For formal verification teams: Improve proof automation tools used in verifying hardware and software correctness by integrating models trained with search-aligned loss functions.
Authors
Adam Ousherovitch, Ambuj Tewari
Abstract
Fine-tuned Large Language Models (LLMs) significantly advance Automated Theorem Proving (ATP), but are often deployed as guiding policies within tree search rather than for single-attempt generation. Recent work shows cross entropy is suboptimal for an LLM used in flat search strategies such as aggregation or filtering and that work has developed new loss functions to correct this misalignment. Extending this alignment to tree search is more challenging: proof discovery depends on exploration and recovery through off-trace states that supervised demonstrations do not reveal. We extend Compute-Aligned Training (CAT) to this setting through an abstraction of policy-guided search, deriving tractable, trace-supported losses. Alongside these search-aware losses, we introduce a search-agnostic uniform-allocation (UA) loss that accounts for the budget without specifying the specific search. Both induce scalar weights on per-tactic cross-entropy gradients. We characterize how off-trace behavior affects the search-aware weights, including conditions for vanishing approximation error at large budgets. On a Lean benchmark, both approaches achieve higher observed proof-success rates than cross-entropy across six search strategies, with strong results from a single shared UA adapter. Budget sweeps show larger gains over cross-entropy at 16 than at 256 expansions, implying CAT scales with test time compute.