Direct Optimization of Generators for Search in Automated Theorem Proving
cs.AI, stat.ML
Submitted: 2026-09-22
Updated: 2026-09-22
License: http://creativecommons.org/licenses/by/4.0/
The gist: 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.
Terminology
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.
Sources
- Llemma: An Open Language Model For Mathematics
- Machine learning and information theory concepts towards an AI Mathematician
- Evaluating Large Language Models Trained on Code
- Training Verifiers to Solve Math Word Problems
- LoRA: Low-Rank Adaptation of Large Language Models
- HyperTree Proof Search for Neural Theorem Proving
- Policy-Guided Heuristic Search with Guarantees
- Single-Agent Policy Tree Search With Guarantees
- Learning Formal Mathematics From Intrinsic Motivation
- Generative Language Modeling for Automated Theorem Proving
- Qwen2.5 Technical Report
- Scaling LLM Test-Time Compute Optimally can be More Effective than Scaling Model Parameters
- Lean Copilot: Large Language Models as Copilots for Theorem Proving in Lean
- Optimizing Language Models for Inference Time Objectives using Reinforcement Learning
- BFS-Prover: Scalable Best-First Tree Search for LLM-based Automatic Theorem Proving
Related papers
- MAVEN-T: Reinforced Heterogeneous Distillation for Real-Time Multi-Agent Trajectory Prediction
- Model Discovery Agent: LLM-assisted Bayesian experiment design for data-efficient discovery of mechanistic world models
- The Clinician's Veto: Navigating Trust, Liability, and Uncertainty in Autonomous AI Prescribing
- MindHelper: Closed-Loop Embodied Mental-State Reasoning for Precision Intervention
- Incumbent Advantage: Brand Bias and Cognitive Manipulation Dynamics in LLM Recommendation Systems
- VSAL: A Vision Solver with Adaptive Layouts for Graph Property Detection