Direct Optimization of Generators for Search in Automated Theorem Proving
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.
Related
- The Order Mass Identity of Rooted Forests, with a Direct Proof of the Extended Aldous-Broder Theorem
- From Corridor Selection to Earthwork: A Multi-Stage Framework for Automated Road Design via Steiner Trees and Convex Optimization
- Global Optimization Framework for Automated Low-Thrust Gravity-Assist Trajectory Design
- OptScientist: Automated Discovery of Optimization Laws
- Robust Direct Preference Optimization under Preference Label-Flip Attacks
BibTeX
@misc{ousherovitch2026direct,
title = {{Direct Optimization of Generators for Search in Automated Theorem Proving}},
author = {Adam Ousherovitch and Ambuj Tewari},
howpublished = {arXiv},
year = {2026},
month = sep,
doi = {10.48550/arxiv.2609.25575},
eprint = {2609.25575},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2609.25575},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - Direct Optimization of Generators for Search in Automated Theorem Proving
AU - Adam Ousherovitch
AU - Ambuj Tewari
PY - 2026
DA - 2026/09/22
PB - arXiv
DO - 10.48550/arxiv.2609.25575
UR - https://arxiv.org/abs/2609.25575
AB - 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.
ER -