LeanTree: Efficient Tree Methods in ML & Proof Search
- LeanTree is a term for efficiency-oriented tree methods that decompose large search or prediction spaces in ML reasoning, extreme classification, and Lean 4 proof search.
- Its LLM variant employs a guided tree search with a learned value network and adaptive exploration budgets to achieve near state-of-the-art accuracy at a fraction of token cost.
- In extreme classification and Lean 4 proof search, LeanTree leverages recursive partitioning and state factorization to optimize latency and enable parallel, white-box reasoning.
Searching arXiv for the cited LeanTree-related papers and closely related work. I’ll look up the arXiv records to ground the article in the relevant papers. Searching for "LeanTree" and the specified arXiv IDs. LeanTree is a name used for multiple distinct tree-based methods in contemporary machine learning and automated reasoning research. In the sources considered here, it denotes three different systems: a guided tree-search procedure for LLM reasoning introduced under the title "LiteSearch: Efficacious Tree Search for LLM" (Wang et al., 2024); a label-tree construction method for extreme multi-label classification (Baharav et al., 2021); and a Lean 4 toolkit and dataset for white-box proof search with factorized states (Kripner et al., 19 Jul 2025). The shared motif is not a single unified algorithm, but the use of tree structure to make large search or prediction spaces computationally tractable.
1. Distinct meanings of the term
The three uses of LeanTree differ in domain, object of search, and optimization target.
| Domain | Tree object | Primary trade-off |
|---|---|---|
| LLM mathematical reasoning | Partial chain-of-thought search tree | Accuracy versus tokens generated |
| Extreme multi-label classification | Label tree over a very large label set | Latency versus statistical performance |
| Lean 4 automated theorem proving | Factorized proof-state tree / AND-OR structure | Proof success versus state size, parallelism, and verifier interaction |
In the LLM setting, the root is the question , and each node is a partial chain of thought (Wang et al., 2024). In the extreme classification setting, each leaf corresponds to one label and internal nodes support binary decisions during beam search (Baharav et al., 2021). In the Lean 4 setting, a proof state is a finite sequence of open goals, and factorization is defined by the connected components of a dependency graph over metavariable sharing (Kripner et al., 19 Jul 2025).
This suggests that the term "LeanTree" is best understood as a recurrent naming convention for efficiency-oriented tree structuring, rather than as a single research lineage.
2. Guided tree search for LLM reasoning
In "LiteSearch: Efficacious Tree Search for LLM," LeanTree refers to a guided tree-search algorithm designed for multi-step mathematical reasoning (Wang et al., 2024). The motivation is explicit: greedy decoding ("System 1") is fast but error-prone on multi-step math problems, whereas off-the-shelf tree searches such as DFS, BFS, and MCTS can improve accuracy but typically consume more LLM compute, measured in tokens generated.
The core design has three parts. First, a learned value network estimates the "promise" of a partial reasoning node without requiring step-wise labels. Second, node selection combines estimated future success with current progress along the greedy trajectory. Third, each selected node receives a node-level exploration budget , interpreted as the maximum number of children to expand, with the budget shrinking as estimated value rises and as depth grows. The search alternates between selection and budgeted expansion until either a complete answer with value exceeding a threshold is found or the iteration budget is exhausted.
The node-selection rule combines a value term and a progress term : where balances exploitation versus fastest-progress. In practice, is described as or similarly for tokens.
Budget computation is depth- and value-dependent: 0 where 1 is a target cumulative success probability and 2 is a user-set upper bound. The calibrated value is
3
so that the first few steps are nudged toward the value 4 of the corresponding greedy-chain node. The paper states the intended effects directly: the deeper 5 is, the smaller 6, and the higher the calibrated value 7, the smaller 8.
Value-network training uses only final-answer supervision. If 9 is a sampled full chain and 0 indicates whether the final answer matches the gold answer, then every prefix 1 receives that delayed reward under the MSE loss
2
The model instantiation described in the source uses Llama3-8B’s last-token hidden state with a linear head 3.
Empirically, the paper reports near-MCTS/BFS performance at only 4 of the token cost. On GSM8K with Mixtral-8×7B, the reported figures include: Greedy 5, SoftVoting@10 6, LeanTree-Batch 7, and LeanTree-Incr 8. On TabMWP with Mixtral-8×7B, the reported figures are: Greedy 9, SoftVoting@10 0, LeanTree-Batch 1, and LeanTree-Incr 2. A further GSM8K comparison reports MCTS+Rerank at 3 accuracy and 4 tokens, SoftVoting@20 at 5 and 6, LeanTree-Batch at 7 and 8, and LeanTree-Incremental at 9 and 0.
The ablation claims are equally specific. Replacing the adaptive budget by a static 1 roughly triples the token cost with no accuracy gain, and removing the depth penalty or greedy calibration degrades either cost or performance. Sensitivity to 2 and 3 is described as smooth, permitting explicit speed-versus-accuracy tuning. Deployment guidance distinguishes Incremental expansions as the lowest-cost mode and Batch expansions as the highest-accuracy mode; recommends a single mixed-domain value network in multiple reasoning domains; advises calibration of the first few steps with greedy-chain values; introduces an optional fallback threshold 4 on 5; and notes that the value network can be deployed in parallel because its cost is 6 LLM decoding cost.
3. Label-tree construction for extreme multi-label classification
In "Enabling Efficiency-Precision Trade-offs for Label Trees in Extreme Classification," LeanTree denotes a method for constructing label trees in extreme multi-label classification (XMC) (Baharav et al., 2021). The formal setup uses training examples 7 with 8 and sparse label vectors 9, where 0 may be 1. The motivation is that a naive 1-vs-All approach incurs 2 inference cost per query, whereas a label tree can support 3 inference when balanced.
The tree semantics are operational. Each leaf 4 is assigned to one label. Each internal node trains a binary classifier to decide whether the set of true labels lies in the left subtree or the right subtree. Inference proceeds by beam search, extending only the top 5 partial paths at each level; if the tree is balanced, the expected number of node evaluations is 6.
The defining contribution is a multi-objective criterion that interpolates between latency and accuracy. For leaf depths 7 and label marginals 8, the expected-depth proxy is
9
For empirical loss 0, the combined objective is
1
where 2 trades off latency versus accuracy. The construction is top-down and recursive, with each split optimized by a fast alternating-maximization solver over label assignments 3, unit vectors 4 and 5, and a balance constraint 6. The source describes the corresponding surrogate as combining "similarity-driven" clustering and "frequency-driven" splitting.
The method is explicitly linked to information theory. Entropy 7 is cited as a lower bound on the expected code length of any prefix-free code, and Fano’s code is invoked as a top-down scheme that splits the symbol set into two subsets with total probability as close to 8 as possible, producing expected depth at most 9. As 0, the LeanTree surrogate reduces to a Fano-style frequency split at each node and thus nearly minimizes expected depth. At 1, the method becomes balanced spherical 2-means on label embeddings 2, matching the similarity-driven extreme.
The recursive construction algorithm is defined by alternating updates. Given a label subset 3, if 4, 5 becomes a leaf. Otherwise, node-specific weights 6 are computed, 7 and 8 are initialized, and the algorithm repeats until stationary: with 9 fixed, update 0 and 1, then normalize; with 2 fixed, compute
3
and solve
4
in 5 time by sorting 6. The full tree is built in 7 time with 8 memory for embeddings and 9 for frequencies.
The empirical record spans public benchmarks and proprietary e-commerce data. On Wiki-500K, LeanTree with 0 reduces expected depth@1 by 1 relative to Parabel at equal precision@1. On Amazon-670K and AmazonCat-13K, the reported Pareto curves show that small increases in expected depth yield large gains in precision, and vice versa. Replacing FastText hierarchical-softmax with LeanTree at 2 improves p@1 on Wiki-500K by 3 at only 4 increase in expected depth. On five e-commerce logs with 5 labels, LeanTree maintains p@1 while reducing expected depth@1 by 6.
Practical considerations are central to the paper. For 7, two heuristics for constructing the Fano-style frequency vector 8 are described: greedy removal and assigning each context to its most frequent label. Local optima are acknowledged, and 9 random restarts per node are reported to suffice in practice. The paper also notes that realizing expected-depth savings as wall-clock gains requires adaptive beam search that stops when 00 high-confidence leaves are found, rather than always traversing to fixed depth, and recommends Laplace smoothing 01 to guard against train/test distribution shift in frequencies.
4. White-box proof search with factorized states in Lean 4
In "LeanTree: Accelerating White-Box Proof Search with Factorized States in Lean 4," LeanTree denotes both a tool built in the Lean 4 language and a dataset of factorized intermediate proof states (Kripner et al., 19 Jul 2025). The paper frames the work against a standard distinction in automated theorem proving. In black-box ATP, the model receives a theorem statement 02 and must generate a complete proof 03 in one shot, with only final verification 04. In white-box ATP, proof construction is incremental: 05 and success occurs when 06 has no open goals.
The difficulties for white-box search are identified as state-size explosion, branching factor, and lack of tooling and data. In Lean, a proof state 07 is a list of open goals, and each goal accumulates local hypotheses from previous steps. Tactics such as cases or induction create multiple subgoals; without factorization, the prover operates on a monolithic list of goals and loses the opportunity to solve independent subgoals in parallel.
The system architecture has two layers. At the verifier level, LeanTree patches the official Lean REPL to report open goals and contexts after each tactic, factorize states, and incremental-verify only the new assignments. At the orchestration level, a thin Python wrapper manages a pool of Lean processes and exposes three API calls: load_file(path), apply_tactic(state_id, tactic_str), and factorize(state_id).
The formal representation is graph-theoretic. A proof state is
08
where each goal 09 is a metavariable 10 together with a local context 11. Two goals are coupled if they share at least one metavariable. This induces the dependency graph
12
Theorem 2.2 states that the connected components of 13 partition 14 into independent sub-states
15
and that no tactic that closes goals in 16 can affect goals in 17 for 18. The factorization algorithm therefore consists of building the dependency graph, finding connected components, and partitioning the goal list accordingly.
The associated dataset is large-scale. Its sources are Mathlib 4, contributing 19 factorized tactic proofs, and DeepSeek-Prover-V1, contributing 20 proofs. The total is
21
Each theorem record stores imports and, for each by-block, a proof tree whose nodes contain id, proof_size, proof_depth, state.goals, tactic_string, children, and tactic_depends_on. Verification is performed by replaying each tree in Lean REPL with incremental kernel-checking. The average branch count is reported as
22
with the accompanying observation that most proofs are linear but a nontrivial tail has 23.
The factorization analysis is quantitative. Original context complexity is defined as
24
while factorized complexity is
25
The reduction ratio satisfies
26
For search time, serial roll-out cost is 27, ideal parallel cost is 28, and the speedup is described by
29
The implementation also hashes factorized substates by their multisets of goals for state reuse and attaches Lean error codes to failing tactic edges to support pruning without discarding unaffected branches.
The reported experiments are preliminary but concrete. Using Llemma-7B on the Lean 4 version of MiniF2F with 30 test theorems, the procedure performs 31 independent linear rollouts of up to 32 tactics. The success rates are:
| Method | Success rate |
|---|---|
| Whole-proof generation | 33 |
| Black-box linear rollout (34) | 35 |
| White-box linear rollout (36) | 37 |
| Best-first search | 38 |
The paper states that providing factorized states more than triples the success rate compared to black-box rollout, and that white-box linear rollout outperforms full proof generation in one shot.
5. Recurrent technical patterns across the three LeanTree usages
Despite the absence of a single common algorithm, the three LeanTree variants share several structural patterns. Each introduces a decomposition of a large combinatorial object into locally manageable units. In LLM reasoning, the units are partial chains of thought with value-guided node selection and node-level exploration budgets (Wang et al., 2024). In extreme classification, the units are recursively partitioned label subsets with a tunable split criterion controlled by 39 (Baharav et al., 2021). In Lean 4 proof search, the units are independent sub-states derived from connected components of a dependency graph over goals (Kripner et al., 19 Jul 2025).
Each usage also defines a specific efficiency metric. The LiteSearch version measures compute in generated tokens and explicitly compares against greedy decoding, MCTS, BFS, and SoftVoting. The XMC version uses expected depth@40 and precision@41, and relates tree structure to prefix-code length and beam-search latency. The Lean 4 version defines context-size reduction, ideal parallel speedup, and theorem-proving success rate under verifier interaction. This suggests that "LeanTree" is repeatedly associated with controllable operating points rather than with a fixed optimization target.
A further commonality is that all three systems avoid purely exhaustive search. The LLM version shrinks the exploration budget for deeper or higher-valued nodes. The XMC version does not optimize solely for balanced depth or solely for similarity clustering, but interpolates between them. The Lean 4 version does not insist on monolithic proof-state expansion, but instead factorizes and reuses independent substates. A plausible implication is that the name has been attached, in multiple subfields, to methods that seek efficiency by selective expansion rather than brute-force enumeration.
6. Scope, misconceptions, and open directions
A common misconception would be to treat LeanTree as one method with a stable meaning across arXiv. The record here does not support that interpretation. The term refers, in different papers, to a lightweight value-guided tree search for LLM reasoning (Wang et al., 2024), a label-tree construction algorithm for XMC (Baharav et al., 2021), and an open-source Lean 4 toolkit with a factorized-state dataset for white-box proof search (Kripner et al., 19 Jul 2025).
Another misconception would be to view all three solely as accuracy-improvement methods. The explicit objectives are broader. The LiteSearch paper emphasizes competitive performance at significantly lower computational cost and reports 42 of the tokens relative to baseline methods. The XMC paper is organized around an efficiency-precision trade-off and states latency improvements such as 43 on Wiki-500K at equal accuracy and 44 on e-commerce logs while maintaining p@1. The Lean 4 paper foregrounds context reduction, parallel search across branches, efficient reuse of states, and feedback in case of errors, and describes its experimental findings as preliminary.
The open problems are domain-specific. In the LLM setting, deployment guidance already anticipates value-network misestimates via a fallback threshold 45 and a small SoftVoting run. In XMC, the source highlights local optima in alternating maximization, sensitivity to frequency shift, and the engineering difficulty of translating expected-depth gains into wall-clock savings. In Lean 4 proof search, the paper lists future work including retrieval-augmented models in the white-box loop, more sophisticated AND-OR search algorithms such as MCTS or best-first, support for more advanced Lean tactics such as calc and conv, and reinforcement learning for discovering new proof strategies.
Taken together, these works show that LeanTree is not a single canonical technique but a family of efficiency-oriented tree formulations that appear in distinct research programs. The stable theme is explicit control over combinatorial growth: through value-guided budget allocation, information-theoretic tree construction, or graph-based factorization of proof states.