Pub-AI: AI in Science

JEPA · 2026-09-30

Predicting the Next State Is Not Enough: JEPA Representations for Lean Theorem Proving

Aarnav Choudhary

arXiv:2609.32908PDF

Auto-summarized: this summary was generated by a language model from the paper’s text and has not been reviewed by an editor. Check the paper before relying on it. Benchmark numbers appear only after a human has verified them.

TL;DR

The authors study whether a JEPA-style latent proof-transition objective from one-step Lean transitions can rank kernel-validated successor states for search. On a same-theorem ranking diagnostic, JEPA has higher Top-1 than matched InfoNCE (50.18% vs 31.55%). But in fixed-budget kernel-checked search, JEPA solves fewer theorems (282.3 vs proposer ordering 308) and requires more tactic checks, indicating one-step transition ranking is insufficient long-horizon search value under this setup.

Why it matters

The paper provides a controlled evaluation separating representation quality (local transition ranking) from end-to-end proof-search performance (kernel-checked completions). It shows that stronger one-step transition ranking does not necessarily improve branch ordering for long-horizon proof search when successor states come from a fixed pretrained tactic proposer and scoring is kernel-gated.

Method

  • Train a JEPA-style latent transition model on recorded one-step Lean transitions (s,a,s′), using a context encoder, predictor, and an EMA target encoder; include a low-weight contrastive auxiliary and a matched direct-contrastive InfoNCE baseline.
  • Use a pinned pretrained ByT5 tactic generator to propose tactics; Lean executes/validates tactics, and the JEPA model scores only kernel-validated, nonterminal successors to order beam expansion (tactics are not generated by the learned model).
  • Evaluate both representation diagnostics (cross-theorem and same-theorem ranking Top-1/MRR, etc.) and a fixed-budget Lean-kernel search endpoint under identical protocol across three seeds, with kernel-checked proof completion as the primary measure.

Limitation

Our results are limited to one small encoder, one tactic generator, and one random split. The SFT rows contain independent one-step transitions rather than verified trajectories, precluding trajectory objectives and a learned remaining-step control; the design also does not distinguish objective mismatch from out-of-distribution scoring.

Abstract (from arXiv)

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.

Related papers