Rethinking Supervision Granularity: Segment-Level Learning for LLM-Based Theorem Proving

Abstract: Automated theorem proving with LLMs in Lean 4 is commonly approached through either step-level tactic prediction with tree search or whole-proof generation. These two paradigms represent opposite granularities for constructing supervised training data: the former provides dense local signals but may fragment coherent proof processes, while the latter preserves global structure but requires complex end-to-end generation. In this paper, we revisit supervision granularity as a training set construction problem over proof trajectories and propose segment-level supervision, a training data construction strategy that extracts locally coherent proof segments for training policy models. We further reuse the same strategy at inference time to trigger short rollouts for existing step-level models. When trained with segment-level supervision on STP, LeanWorkbook, and NuminaMath-LEAN, the resulting policy models achieve proof success rates of 64.84%, 60.90%, and 66.31% on miniF2F, respectively, consistently outperforming both step-level and whole-proof baselines. Goal-aware rollout further improves existing step-level provers while reducing inference costs. It increases the proof success rate of BFS-Prover-V2-7B from 68.77% to 70.74% and that of InternLM2.5-StepProver from 59.59% to 60.33%, showing that appropriate supervision granularity better aligns model learning with proof structure and search. Code and models are available at https://github.com/NJUDeepEngine/SEG-ATP.

• The paper introduces segment-level supervision that leverages open-goal changes to structure proof segments, enhancing local coherence in LLM training.
• It demonstrates that segment-level methods outperform both step-level and whole-proof approaches with substantial gains on benchmarks like miniF2F.
• The approach improves inference efficiency through goal-aware rollouts, reducing token usage and time costs while maintaining high proof success rates.

Segment-Level Supervision for LLM-Based Theorem Proving in Lean 4

Motivation and Problem Formulation

Automated theorem proving (ATP) with LLMs in Lean 4 is typically addressed via two regimes: step-level tactic prediction with tree search and whole-proof generation. The former delivers dense, local supervision by predicting the next tactic given the current proof state, but it fragments coherent proof processes, whereas the latter trains the model end-to-end to generate the complete proof script but suffers from long-horizon generation dependencies and expensive inference. This work critically reevaluates the granularity of supervision used for LLM-based ATP by proposing an intermediate regime: segment-level supervision, which organizes verified proof trajectories into locally coherent proof segments reflecting meaningful structural transitions, principally via changes in the number of open goals.

Figure 1: Illustration of step-level, segment-level, and whole-proof supervision targets on a NuminaMath-LEAN example, contrasting their respective inference and training units.

By leveraging goal structure as a lightweight signal for segment boundaries, this approach allows the same proof trajectory to be decomposed into training targets that are neither atomic nor globally holistic but instead preserve local proof coherence. This boundary-selection formulation unifies previous data construction paradigms and motivates the use of open-goal changes to mark transitions between contiguous proof segments.

Methodology

The authors formalize supervision granularity as a boundary-selection problem on verified proof-state trajectories. Given a proof trajectory, a boundary strategy selects a sequence of positions, and extracts the tactic subsequences between them as macro actions. Step-level supervision places boundaries at every proof state, resulting in single-tactic targets; whole-proof supervision only uses initial and terminal states, yielding the full tactic sequence as target. The segment-level regime places a boundary at every proof state where the open-goal count changes. Each segment thus captures a contiguous multi-tactic progression until the next structural transition.

The segment-level policy is trained with standard autoregressive MLE, with targets constructed as macro-tactic sequences starting from goal-changing proof states. The approach does not require new model architectures or synthesis of additional theorem–proof pairs; it merely reorganizes existing verified trajectories to align local learning units with goal structure.

The authors also introduce goal-aware rollout, where at inference time, step-level provers perform short rollouts in states immediately following a change in open-goal count, compressing search transitions and reducing fine-grained branching.

Empirical Evaluation

Extensive experiments were conducted using Qwen2.5-Math-7B on STP, LeanWorkbook, and NuminaMath-LEAN training datasets, with evaluation on both in-domain sets and the miniF2F benchmark. Three supervision granularities (step-level, segment-level, whole-proof) were compared under matched training and controlled evaluation protocols.

Segment-level supervision consistently outperformed both step-level and whole-proof baselines on miniF2F, achieving proof success rates of 64.84%, 60.90%, and 66.31% when trained on STP, LeanWorkbook, and NuminaMath-LEAN respectively. The gains are robust across in-domain and out-of-domain benchmarks, with segment-level yielding competitive or superior efficiency (token and time cost) in most settings.

Additionally, applying goal-aware rollout to BFS-Prover-V2-7B and InternLM2.5-StepProver without retraining improved best-first search success rates on miniF2F by 1.97% and 0.74%, respectively, while reducing average inference token and time cost substantially.

Figure 2: Training loss curves under different supervision granularities, highlighting optimization properties and convergence.

Training loss analysis revealed that segment-level supervision achieves steady loss reduction and eventual convergence comparably or superior to step-level and substantially faster than whole-proof, reflecting easier optimization and more coherent local conditioning.

Boundary and Ablation Analysis

Multiple boundary selection strategies were ablated—token-based, tactic-distance-based, and state-distance-based segmentation. Only open-goal-change boundaries achieved maximal performance (66.31% on miniF2F); alternative strategies caused notable accuracy drops, underscoring the importance of aligning macro action boundaries with goal structure rather than superficial textual or token granularity.

Goal-aware rollout horizon hyperparameters were also examined, demonstrating that modest horizons ($H=3,5,7$) suffice, with the optimal value depending on prover strength and test set composition.

Practical and Theoretical Implications

The demonstrated superiority of segment-level supervision highlights supervision granularity as a critical axis in ATP policy optimization. The findings indicate that existing formal proof corpora, when reorganized according to intrinsic proof structural signals (open goals), yield both better learning and search behavior. This approach serves as a valuable complement to prior data generation and expansion strategies.

Practically, segment-level policies and goal-aware rollout provide scalable, architecture-agnostic upgrades to LLM-based provers utilizing verified Lean proofs, offering improved accuracy and inference efficiency without requiring retraining or model modifications. The method is especially advantageous for datasets and domains with multi-step local proof progressions.

Theoretically, this work suggests that aligning supervision units with latent structural transitions facilitates both optimization and successful proof search, potentially unlocking advanced reasoning capabilities in LLMs for formal mathematics. Segment-level granularity bridges the gap between fine-atomic action prediction and intractable global autoregressive modeling, opening new avenues for hierarchical policy modeling and subgoal decomposition in deep theorem proving.

Anticipated further directions include: exploring alternative structural signals for segment boundaries, integrating dynamic segment-length adaptation, and combining segment-level learning with reinforcement learning and constrained search paradigms.

Conclusion

This paper provides a rigorous analysis and implementation of segment-level supervision for LLM-based ATP in Lean 4, demonstrating robust gains in proof success and inference efficiency across diverse datasets and provers. The results establish open-goal-change as an effective boundary signal, and segment-level granularity as a practical and theoretically justified regime for policy learning and search in formal theorem proving. These insights are expected to catalyze further advances in hierarchical reasoning, policy alignment, and efficient search for large-scale, LLM-driven automated mathematics.