A JEPA model scores kernel-validated Lean successor states from LeanDojo transitions to reorder branches in beam-search proof search.
Abstract
Neural theorem provers must both propose tactics and decide which valid successor states to explore. We study whether one-step Lean transitions provide a self-supervised signal for branch ordering. A JEPA-style model predicts latent successor representations and scores only kernel-validated, nonterminal successors generated by a fixed pretrained ByT5 proposer. JEPA achieves higher Top-1 than matched InfoNCE on a same-theorem ranking diagnostic(50.18% versus 31.55%), but averages 282.3 of 987 solved theorems across three seeds versus 308 for proposer ordering, while requiring more tactic checks. In this setting, accurate one-step transition ranking is therefore insufficient as a long-horizon search value. The controlled evaluation separates representation from proposal quality and treats kernel-checked proof completion as the primary endpoint.
Problem
Neural Lean provers must both propose tactics and decide which valid successor states to explore. It is unclear whether one-step Lean transitions provide a useful self-supervised signal for ordering search branches.
Approach
A JEPA-style model with a context encoder, predictor, and EMA target encoder predicts latent successor representations from recorded (state, tactic, successor) transitions in LeanDojo Benchmark 4. A fixed pretrained ByT5 tactic generator proposes candidates that Lean validates; the JEPA model scores only kernel-validated, nonterminal successors to influence beam-search branch ordering without altering tactic generation. It is compared against a matched InfoNCE objective, untrained weights, a hand-designed progress heuristic, and proposer score under a fixed kernel-check budget.
Figure 1: The search pipeline: a pinned ByT5 model proposes up to K tactics, which are executed and validated by Lean; invalid tactics are discarded and completed proofs end the search. The JEPA model scores only valid nonterminal successor states, influencing branch expansion order without affecting tactic generation or correctness.
Results
JEPA achieves higher same-theorem Top-1 ranking (50.18%) than InfoNCE (31.55%), but solves fewer theorems in fixed-budget search (282.3 of 987 across three seeds) than proposer ordering (308) while using more kernel checks. Accurate one-step transition ranking is insufficient as a long-horizon search value; the study yields a controlled negative result.
