# Direct Optimization of Generators for Search in Automated Theorem Proving

> Compute-aligned training losses that reweight per-tactic cross-entropy gradients consistently outperform standard cross-entropy across all six search strategies in Lean theorem proving, with gains scaling as test-time compute grows.

- **Source:** [arXiv](https://arxiv.org/abs/2609.25575)
- **Published:** 2026-09-26
- **Permalink:** https://picx.dev/p/7OKctG
- **Whiteboard:** https://picx.dev/p/7OKctG/image

## Summary

## Summary (Overview)

- **Core contribution**: This paper extends Compute-Aligned Training (CAT) to tree-search-based Automated Theorem Proving (ATP), deriving tractable loss functions that align LLM training with deployment in policy-guided search algorithms under finite compute budgets.
- **Two complementary objectives**: (1) search-aware CAT losses that model specific search strategies (Pass@N, Best-First Search), and (2) a search-agnostic Uniform-Allocation (UA) loss that accounts for compute budget without specifying the search rule. Both induce scalar weights on per-tactic cross-entropy gradients.
- **Theoretical characterization**: The paper analyzes approximation error from trace-supported training, showing how alternative proofs ("bypass channel") deflate gradient weights while off-trace exploration costs ("trap channel") inflate them, with conditions for vanishing error at large budgets.
- **Empirical results**: On a Lean benchmark (458 held-out theorems), both CAT and UA achieve higher proof-success rates than Cross-Entropy (CE) across all six evaluated search strategies, with the shared UA adapter performing strongly across strategies.
- **Scaling behavior**: Budget sweeps show larger gains over CE at 256 expansions than at 16, demonstrating that compute-aligned training scales with test-time compute.

## Introduction and Theoretical Foundation

### Background and Motivation

Automated Theorem Proving (ATP) is a critical frontier in AI with applications in software and hardware verification. Formal verifiers like Lean 4 allow generated proofs to be checked, but proof generation remains difficult due to sparse rewards and long horizons. Large Language Models (LLMs) predict proof steps, yet a single flawed step can invalidate an entire attempt, motivating the pairing of LLMs with structured search algorithms.

### The Alignment Problem

Standard Supervised Fine-Tuning (SFT) with Cross-Entropy (CE) maximizes the likelihood of demonstrated traces:

$$L_{CE}(\theta) = -\sum_{t=1}^{L} \log \pi_{\theta}(y_t^* \mid x_{t-1}) = -\sum_{t=1}^{L} \log p_t$$

However, recent work shows CE is suboptimal when deployment uses search. Existing corrections focus on flat strategies (candidate aggregation, filtering), but extending alignment to tree search is fundamentally more difficult because:
- Proof discovery depends on exploration and recovery through off-trace states
- Supervised demonstrations do not reveal what happens when search deviates from the trace
- The exact objective requires off-trace continuations that demonstrated proof traces do not supply

### Related Work

The paper situates itself within three research threads:
1. **Neural ATP**: Systems like Lean Copilot, Llemma, and ReProver, powered almost exclusively by SFT
2. **Test-time search algorithms**: From uninformed traversal (DFS, Pass@N) to policy-guided searches (BFS, Levin Search)
3. **Test-time aligned training**: Prior work (e.g., Direct Coverage Optimization) focused only on non-branching strategies like Pass@N where the objective has closed-form expression; this paper addresses the branching case

## Methodology

### Formal Setup

The environment is defined by a state space $\mathcal{X}$, action space $\mathcal{V}$, and transition function $\mathcal{T}: \mathcal{X} \times \mathcal{V} \to \mathcal{X} \cup \{x_{err}\}$. A demonstrated trace is:

$$\tau^* = (x_0, y_1^*, x_1, \ldots, x_{L-1}, y_L^*, x_L), \qquad x_L = x_{QED}$$

### Search as Configuration Dynamics

A search algorithm $\mathcal{A}$ is specified by two operations:
- **Selection rule** $S_{\mathcal{A}}(x \mid c)$: picks a node from the frontier
- **Update rule** $U_{\mathcal{A}}(c, x, z)$: updates the configuration from observations

The true success probability $J^{\mathcal{A}}(c)$ satisfies a Bellman-like recurrence:

$$J^{\mathcal{A}}(c) = \begin{cases} 1 & \text{if } c \text{ solved;} \\ 0 & \text{if } c \text{ dead (and unsolved),} \\ \sum_{x \in \mathcal{F}(c)} S_{\mathcal{A}}(x \mid c) \sum_{y \in \mathcal{Y}} \pi_{\theta}(y \mid x) J^{\mathcal{A}}(U_{\mathcal{A}}(c, x, \mathcal{T}(x, y))), & \text{otherwise} \end{cases}$$

### Trace-Supported Surrogate

Since the exact objective cannot be evaluated from supervised traces, the paper restricts search to demonstrated states. The offline transition collapses all deviations (valid off-trace and invalid tactics) into a single "trace miss" event:

$$\widehat{\mathcal{T}}_{\tau^*}(x_{t-1}, y) = \begin{cases} x_t, & y = y_t^*, \\ \bot_{\tau}, & y \neq y_t^*, \end{cases}$$

This induces a tractable recurrence for the surrogate success probability $\widehat{J}_{\mathcal{A}}(\hat{c}; p)$.

### Key Theoretical Result: CAT Gradient

**Proposition 1 (CAT Loss)**: The CAT gradient is a weighted sum of per-step cross-entropy gradients:

$$\nabla_{\theta} \widehat{L}_{CAT}^{\mathcal{A}}(\theta) = \sum_{t=1}^{L} \hat{w}_t^{\mathcal{A}} \nabla_{\theta} L_{CE}^{(t)}, \qquad \hat{w}_t^{\mathcal{A}} = \frac{p_t}{\widehat{\mathcal{J}}_{\tau^*}^{\mathcal{A}}} \frac{\partial \widehat{\mathcal{J}}_{\tau^*}^{\mathcal{A}}}{\partial p_t}$$

### Uniform-Allocation (UA) Objective

For an expansion budget $N \geq L$, UA assigns each demonstrated step a uniform local budget $\bar{K} = N/L$:

$$\widehat{\mathcal{J}}^{UA}(p; N) = \prod_{t=1}^{L} \left[ 1 - (1 - p_t)^{\bar{K}} \right]$$

with gradient weights:

$$\hat{w}_t^{UA} = \frac{\bar{K} p_t (1 - p_t)^{\bar{K} - 1}}{1 - (1 - p_t)^{\bar{K}}}$$

At $N = L$, UA recovers CE; additional local attempts reduce emphasis on tactics already likely to be sampled.

### Strategy-Specific Instantiations

**Pass@N**: The surrogate success probability is:

$$\widehat{\mathcal{J}}_{\tau^*}^{\text{Pass@}N}(\theta) = 1 - \left(1 - \pi_{\text{proof}}\right)^N$$

where $\pi_{\text{proof}} := \prod_{t=1}^{L} p_t$, with uniform weights:

$$\hat{w}_t^{\text{Pass@}N} = \frac{N \pi_{\text{proof}} (1 - \pi_{\text{proof}})^{N - 1}}{1 - (1 - \pi_{\text{proof}})^N}$$

**Best-First Search (BFS)**: Nodes scored by geometric mean of tactic probabilities:

$$s_{\text{BFS}}(x_k; \pi_{\theta}) = \left(\prod_{j=1}^{k} \pi_{\theta}(y_j \mid x_{j-1})\right)^{1/k}$$

The surrogate treats misses as retries at the same state, with success probability:

$$\widehat{\mathcal{J}}_{\tau^*}^{\text{BFS}}(\theta) = \mathbb{P}(S_L \leq N)$$

where $S_L = \sum_{t=1}^{L} T_t$ and $T_t \sim \text{Geom}(p_t)$.

## Empirical Validation / Results

### Experimental Setup

- **Environment**: Lean 4 via LeanDojo, trained on leandojobenchmark-4-random split of mathlib4
- **Dataset**: 458 held-out theorems with reference proofs of 2–5 steps
- **Model**: Qwen2.5-Math-7B-Instruct with LoRA fine-tuning
- **Baseline**: Epoch-matched CE control (shared CE warm-up, one additional epoch)

### Experiment 1: Training for Search Across Strategies

At N = 256 expansions, evaluating Pass@N, BFS, DFS variants, and MCTS:

**Table 1: Proofs found (%) at N = 256 on 458 held-out theorems**

| Training objective | Pass@N | BFS | DFS | RDFS | VDFS | MCTS |
|-------------------|--------|-----|-----|------|------|------|
| CE | 21.2 | 23.8 | 19.2 | 20.7 | 21.2 | 20.3 |
| Search-aware CAT | 26.6* | 25.1 | 20.7 | 23.6* | 23.8 | 22.5 |
| Shared UA | 25.1* | 26.9 | 23.6* | 25.3* | 24.5* | 25.1* |

*Note: * marks p < 0.05 against CE under the same strategy; bold marks the highest reported rate.*

**Key findings**:
- Both search-aware CAT and UA outperform CE under all six strategies
- UA attains the highest point estimate under the five branching strategies
- Search-aware CAT achieves the highest observed success under Pass@N (26.6% vs 25.1% for UA)
- Matching training and deployment gives higher point estimates (e.g., Pass@N/BFS adapter pair: 26.6% vs 23.7% under Pass@N)

### Experiment 2: Alignment Gain vs. Test-Time Budget

**Table 2: Paired search-aware CAT gain over CE**

| Strategy | Budget N | n | b/c | Gain (pp), 95% interval | p |
|----------|----------|---|-----|-------------------------|---|
| Pass@N | 16 | 458 | 10/1 | +2.0 [0.5, 3.4] | 0.012 |
| Pass@N | 64 | 458 | 28/8 | +4.4 [1.8, 6.9] | 0.001 |
| Pass@N | 256 | 408 | 29/12 | +4.2 [1.1, 7.2] | 0.012 |
| BFS | 16 | 458 | 28/8 | +4.4 [1.8, 6.9] | 0.001 |
| BFS | 64 | 458 | 34/13 | +4.6 [1.7, 7.5] | 0.003 |
| BFS | 256 | 407 | 37/9 | +6.9 [3.6, 10.1] | < 0.001 |

*Note: n is the paired subset size; b and c count CAT-only and CE-only successes.*

**Key findings**:
- Positive paired gains at every reported budget for both strategies
- Gains at N = 256 exceed those at N = 16 (Pass@N: +4.2 vs +2.0; BFS: +6.9 vs +4.4)
- Suggests CAT scales with test-time compute

## Theoretical and Practical Implications

### Bypass Channel (Alternative Proofs)

**Theorem 1 (Bypass Deflation)**: For Pass@N with N > 1, accounting for alternative proofs reduces the deployed weight relative to the trace-supported weight:

$$w_t^{\star} \leq \hat{w}_t^{\text{Pass@}N} < 1 = w_t^{\text{CE}}$$

This shows the surrogate overestimates the importance of the demonstrated tactic when alternative proofs exist.

### Trap Channel (Off-Trace Exploration Costs)

**Theorem 2 (Shared-Budget Bracketing)**: When each miss costs $1 \leq \kappa \leq \bar{\kappa}$ expansions before returning to the trace state:

$$w_t^{\text{trap}}(N) \in \left[ \hat{w}_t(N), \hat{w}_t\left(L + \frac{N - L}{\bar{\kappa}}\right) \right] \subseteq [\hat{w}_t(N), 1]$$

The surrogate assumes $\kappa = 1$ (a miss costs one expansion), while real excursions can consume more budget.

### Large-Budget Scaling

**Theorem 3 (Large-Budget Limits and Recoverability)**:
- **Pass@N**: Both deployed and surrogate weights tend to zero as N → ∞, so the search-aware weight error vanishes while CE's error tends to one
- **Shared-budget excursions with permanent absorption probability $q_t$**: The surrogate weight tends to zero while the deployed weight tends to:

$$\omega_t := \frac{q_t}{p_t + (1 - p_t) q_t}$$

When all excursions are recoverable ($q_s = 0$), both weights vanish. With permanent absorption, avoiding misses retains value even with unlimited compute.

### Key Insight

The paper proves (Proposition 4) that two environments can agree on every demonstrated transition, all $p_s$, the search rule, and the budget, yet differ in what a miss causes. Thus, no rule using only trace-level inputs can be exact in both environments—the trace-supported recurrence cannot resolve missing information about recovery.

## Conclusion

### Main Takeaways

1. **Compute-aligned training works**: Both search-aware CAT and search-agnostic UA achieve higher proof-success rates than CE across all six evaluated search strategies in Lean.

2. **Strategy-specific modeling helps**: Search-aware CAT achieves the highest observed success under Pass@N, showing that incorporating strategy-specific structure can provide additional benefits beyond generic compute awareness.

3. **Gains scale with budget**: Larger gains over CE at N = 256 than at N = 16 suggest CAT scales with test-time compute.

4. **Theoretical understanding**: The analysis characterizes how alternative proofs and off-trace exploration affect gradient weights, identifying conditions for vanishing approximation error at large budgets.

### Limitations and Future Work

- Experiments use one model and short reference proofs (2–5 steps)
- Training is restricted to offline SFT with trace-only objectives
- Future directions include:
  - Sampling one-step deviations to reveal alternative transitions
  - Exploring local trees around demonstrations for evidence about alternative proofs and recovery costs
  - Allocating training compute between supervision and search
  - Refining search-aware objectives beyond fixed allocation using observed off-trace data

---

_Markdown view of https://picx.dev/p/7OKctG, served by PicX — AI-generated visual whiteboard summaries of research papers._