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

TL;DR

This paper introduces segment-level supervision for training LLM-based theorem provers, improving proof success rates by better aligning training data granularity with proof structure.

Contribution

It proposes a novel segment-level supervision strategy that enhances proof success and inference efficiency over traditional step-level or whole-proof methods.

Findings

Segment-level supervision achieves higher proof success rates on multiple datasets.

Goal-aware rollouts improve existing provers' success rates and reduce inference costs.

Code and models are publicly available at the provided GitHub repository.

Abstract

Automated theorem proving with large language models 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.