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

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

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