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

LLM-based automated theorem proving in Lean 4 commonly uses 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 and propose segment-level supervision, which extracts locally coherent proof segments to train 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 and 32.13%, 21.73%, and 17.20% on FATE-M, respectively, consistently outperforming both step-level and whole-proof baselines. Goal-aware rollout further improves existing step-level provers while reducing inference costs: on miniF2F, 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%; on FATE-M, the corresponding success rates increase from 45.60% to 50.67% and from 33.47% to 36.13%, showing that appropriate supervision granularity better aligns model learning with proof structure and search. Code and models are available at https://anonymous.4open.science/r/SEG-ATP.