# Direct Optimization of Generators for Search in Automated Theorem Proving

arXiv: https://arxiv.org/abs/2609.25575 · HTML: https://mlsanity.com/paper/2609.25575/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 26 of its 37 references that are in the map.
- First posted on arXiv 22 Sep 2026.

## Builds on (most cited first, 26 of 26)

- [LoRA: Low-Rank Adaptation of Large Language Models](https://mlsanity.com/paper/2106.09685/) (2021)
- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [Evaluating Large Language Models Trained on Code](https://mlsanity.com/paper/2107.03374/) (2021)
- [Self-Consistency Improves Chain of Thought Reasoning in Language Models](https://mlsanity.com/paper/2203.11171/) (2022)
- [Qwen2.5 Technical Report](https://mlsanity.com/paper/2412.15115/) (2024)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Scaling LLM Test-Time Compute Optimally can be More Effective than Scaling Model Parameters](https://mlsanity.com/paper/2408.03314/) (2024)
- [LeanDojo: Theorem Proving with Retrieval-Augmented Language Models](https://mlsanity.com/paper/2306.15626/) (2023)
- [MiniF2F: a cross-system benchmark for formal Olympiad-level mathematics](https://mlsanity.com/paper/2109.00110/) (2021)
- [Generative Language Modeling for Automated Theorem Proving](https://mlsanity.com/paper/2009.03393/) (2020)
- [Llemma: An Open Language Model For Mathematics](https://mlsanity.com/paper/2310.10631/) (2023)
- [Draft, Sketch, and Prove: Guiding Formal Theorem Provers with Informal Proofs](https://mlsanity.com/paper/2210.12283/) (2022)
- [Autoformalization with Large Language Models](https://mlsanity.com/paper/2205.12615/) (2022)
- [HyperTree Proof Search for Neural Theorem Proving](https://mlsanity.com/paper/2205.11491/) (2022)
- [Formal Mathematics Statement Curriculum Learning](https://mlsanity.com/paper/2202.01344/) (2022)
- [Proof Artifact Co-training for Theorem Proving with Language Models](https://mlsanity.com/paper/2102.06203/) (2021)
- [Inference-Aware Fine-Tuning for Best-of-N Sampling in Large Language Models](https://mlsanity.com/paper/2412.15287/) (2024)
- [BFS-Prover: Scalable Best-First Tree Search for LLM-based Automatic Theorem Proving](https://mlsanity.com/paper/2502.03438/) (2025)
- [Formal Mathematical Reasoning: A New Frontier in AI](https://mlsanity.com/paper/2412.16075/) (2024)
- [Lean Copilot: Large Language Models as Copilots for Theorem Proving in Lean](https://mlsanity.com/paper/2404.12534/) (2024)
- [Optimizing Language Models for Inference Time Objectives using Reinforcement Learning](https://mlsanity.com/paper/2503.19595/) (2025)
- [Learning Formal Mathematics From Intrinsic Motivation](https://mlsanity.com/paper/2407.00695/) (2024)
- [InfAlign: Inference-aware language model alignment](https://mlsanity.com/paper/2412.19792/) (2024)
- [Rethinking Fine-Tuning when Scaling Test-Time Compute: Limiting Confidence Improves Mathematical Reasoning](https://mlsanity.com/paper/2502.07154/) (2025)
- [Single-Agent Policy Tree Search With Guarantees](https://mlsanity.com/paper/1811.10928/) (2018)
- [Policy-Guided Heuristic Search with Guarantees](https://mlsanity.com/paper/2103.11505/) (2021)

## Abstract (arXiv)

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.