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

Aarnav Choudhary

cs.LGSep 26, 2026

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

Authors: Aarnav Choudhary

Organizations: University of California, Los Angeles

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.

Figures & tables

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.

	Same theorem (16)		Cross theorem (64)
Objective	Top-1	MRR	Top-1	MRR
Untrained	0.0708	0.1985	0.0200	0.0825
JEPA	0.5018	0.6423	0.7837	0.8057
InfoNCE	0.3155	0.4754	0.8105	0.8422

Table 1: Test transition ranking, reported as three-seed means. Bold marks the best objective within each negative scope; full uncertainties and cosine margins appear in Appendix C .

Search order	Solved	Checks	Wall (s)
Proposer score	308/987	349.92	64.62
Progress	304/987	352.99	65.31
Untrained	$303.7{\pm}8.5$ /987	$351.03{\pm}2.86$	$65.99{\pm}0.30$
InfoNCE	$289.7{\pm}3.1$ /987	$359.89{\pm}0.50$	$66.02{\pm}0.03$
JEPA	$282.3{\pm}1.5$ /987	$363.45{\pm}1.02$	$66.61{\pm}0.08$

Table 2: Fixed-budget Lean-kernel search. Checkpoint rows report three-seed means and sample standard deviations. Bold marks the best value: higher for solved theorems and lower for checks and wall time. Paired comparisons use theorem-level outcomes.

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

Learn predict the recorded successor latent

\rightarrow

Audit remove split and retrieval shortcuts

\rightarrow

Ablate remove the low-weight contrastive auxiliary

\rightarrow

Test allocate a fixed Lean-kernel budget

Appendix

Figure 2: The evidence chain. Local ranking establishes representation quality; only completed, kernel-checked proofs establish search value.

state $s$ + tactic $a$ train-only byte-BPE $\longrightarrow$	context encoder $f_{\theta}$ 6 layers, width 512 $\longrightarrow$	predictor $q_{\phi}$ 3 layers, width 768 $\longrightarrow$	predicted latent $\hat{z}\in\mathbf{R}^{384}$
recorded Lean successor $s^{\prime}$ $\longrightarrow$	target encoder $f_{\xi}$ EMA, $m=0.996$ $\longrightarrow$	stop-gradient target $z^{\prime}\in\mathbf{R}^{384}$ $\longrightarrow$	$1-\cos(\hat{z},z^{\prime})$ $+0.1\mathcal{L}_{\rm var}+0.1\mathcal{L}_{\rm NCE}$
$f_{\theta}\;\xrightarrow{\text{EMA only}}\;f_{\xi}$
Fixed-budget inference
pinned ByT5 up to 64 tactics $\longrightarrow$	Lean execution invalid tactics discarded $\longrightarrow$	real successors or completed proof $\longrightarrow$	cosine branch score beam 8, depth 16

Appendix

Figure 3: Training and inference boundary. The learned model predicts and orders representations; Lean alone creates successors and certifies proof completion.

Setting	Value	Setting	Value
Tokenizer	byte-BPE, train only	Maximum length	1,024
Vocabulary / min. frequency	16,000 / 2	Encoder layers / heads	6 / 8
Hidden / feed-forward width	512 / 1,024	Projection / predictor width	384 / 768
Predictor layers / dropout	3 / 0.15	Trainable / EMA parameters	22.7M / 21.5M
Epochs / seeds	10 / 7, 17, 37	Physical / effective batch	32 / 512
Learning rate / weight decay	$10^{-5}$ / $10^{-4}$	Gradient clip / accumulation	1.0 / 16

Appendix

Table 3: Matched model and optimization settings. Direct InfoNCE changes the objective; untrained controls use the same architecture and tokenizer.

Stage	Objective / split	Negatives	Top-1	MRR	Margin
Initial large model	cosine, artifact split	16	0.0550	0.2027	$-0.0000$
Stabilized model	cosine + var. + NCE, artifact split	16	0.8944	0.9168	0.3054
Clean JEPA	same loss, theorem split	16	0.8862	0.9092	0.3620
Clean InfoNCE	direct contrastive, theorem split	16	0.9220	0.9548	0.2874

Appendix

Table 4: Development diagnostics before the official full-data matrix. Populations differ, so rows are provenance rather than a controlled comparison.

Split	Transitions	Theorem keys	Either side over 1,024 tokens
Train	250,814	59,557	0.38%
Validation	4,260	995	0.35%
Test	4,506	992	0.31%

Appendix

Table 5: Official Benchmark 4 transition split. Cross-split theorem overlap, empty fields, and source-pair mismatches are all zero.

Negative scope	Objective	Top-1	MRR	True cosine	Neg. cosine	Margin
Cross-theorem, 64	Untrained	$0.0200{\pm}0.0045$	$0.0825{\pm}0.0075$	$-0.0222{\pm}0.0142$	$-0.0236{\pm}0.0126$	$0.0013{\pm}0.0021$
Cross-theorem, 64	JEPA	$0.7837{\pm}0.0020$	$0.8057{\pm}0.0024$	$0.9350{\pm}0.0095$	$0.6318{\pm}0.0492$	$0.3032{\pm}0.0398$
Cross-theorem, 64	InfoNCE	$\mathbf{0.8105{\pm}0.0072}$	$\mathbf{0.8422{\pm}0.0056}$	$0.4281{\pm}0.0151$	$0.0870{\pm}0.0211$	$\mathbf{0.3411{\pm}0.0062}$
Same-theorem, 16	Untrained	$0.0708{\pm}0.0256$	$0.1985{\pm}0.0251$	$-0.0205{\pm}0.0138$	$-0.0095{\pm}0.0139$	$-0.0110{\pm}0.0056$
Same-theorem, 16	JEPA	$\mathbf{0.5018{\pm}0.0072}$	$\mathbf{0.6423{\pm}0.0063}$	$0.9403{\pm}0.0101$	$0.8178{\pm}0.0408$	$\mathbf{0.1225{\pm}0.0307}$
Same-theorem, 16	InfoNCE	$0.3155{\pm}0.0074$	$0.4754{\pm}0.0090$	$0.4730{\pm}0.0157$	$0.4282{\pm}0.0188$	$0.0448{\pm}0.0032$

Appendix

Table 6: Hardest official ranking scopes, reported as mean and sample standard deviation over seeds. Raw cosine levels are objective-dependent; rank and within-objective margins are the relevant diagnostics.

Objective	Cross Top-1@64	Cross MRR@64	Same Top-1@16	Same MRR@16
JEPA, $\lambda_{c}=0.1$	$\mathbf{0.7837{\pm}0.0020}$	$\mathbf{0.8057{\pm}0.0024}$	$\mathbf{0.5018{\pm}0.0072}$	$\mathbf{0.6423{\pm}0.0063}$
JEPA, $\lambda_{c}=0$	$0.3880{\pm}0.1347$	$0.4810{\pm}0.1289$	$0.3721{\pm}0.0237$	$0.5113{\pm}0.0269$

Appendix

Table 7: Requested loss ablation on the predefined representation diagnostics. Bold marks the stronger mean within each column.

Search order	Solved	Step success	Checks/thm.	Wall s/thm.	Timeouts
Proposer score	308/987	26.59%	349.92	64.62	83
Progress	304/987	28.26%	352.99	65.31	89
Untrained seed 7	294/987	25.90%	353.52	66.23	90
Untrained seed 17	307/987	26.40%	351.65	66.08	91
Untrained seed 37	310/987	26.16%	347.91	65.65	93
JEPA seed 7	284/987	25.16%	364.61	66.52	84

Appendix

Table 8: Complete fixed-protocol outcomes. Higher is better for solved and step success; lower is better elsewhere. Errors are retained as unsolved.

Comparator	JEPA $\Delta$ solves	$\Delta$ success (points)	95% interval	Wins/losses/ties
Proposer score	$-25.67$	$-2.600$	$[-3.715,-1.554]$	6/36/945
Progress	$-21.67$	$-2.195$	$[-3.445,-1.047]$	13/35/939
Untrained	$-21.33$	$-2.161$	$[-3.073,-1.317]$	12/49/926
InfoNCE	$-7.33$	$-0.743$	$[-1.317,-0.236]$	6/20/961

Appendix

Table 9: Theorem-clustered paired comparisons. Negative values favor the comparator; learned conditions are averaged within theorem over seeds.

Cross-theorem ranking InfoNCE 81.05% JEPA 78.37% InfoNCE wins

Same-theorem ranking JEPA 50.18% InfoNCE 31.55% JEPA wins

Kernel completion proposer 31.21% JEPA 28.61% proposer wins

Appendix

Figure 4: The proxy reversal. Either representation objective can win a local diagnostic while neither induces the best proof-search order.

Gate	Failure condition	Accepted evidence
Data	split overlap, empty state, source mismatch	259,580/259,580 exact pairs
Ranking	incomplete matrix or negative-pool drift	nine linked checkpoints
Kernel	missing key, budget drift, malformed error	11 conditions, 992 keys each
Pairing	unmatched theorem set or seed as sample	theorem-clustered comparison
Proof validity	sorry , metavariable, declaration failure	kernel-checked completion
Manuscript	placeholder, identifier, unembedded font, cell drift	exact result handoff

Appendix

Table 10: Fail-closed gates connecting imported data, checkpoints, search outcomes, and reported claims.

Lean Theorem Proving

Lean theorem proving uses the Lean proof assistant to construct machine-checked mathematical proofs, often with automated search and language models.

See all 123 papers

Oct 6, 2026

SCOPE: Certified Theorem Proving with a Language Model as the Policy Planner

In proof assistants such as Lean, a generated proof must pass machine compilation checks, so evaluation needs no human scoring. Direct generation fails on multi-step numeric propositions: a proof is valid only if every content integer is correct, so the pass rate is bounded by the k-th power of the per-integer accuracy. Controlled corruption across 2,617 reference proofs confirms this power law. SCOPE (State-Conditioned Operator Planning and Execution) enforces the natural division of labor: the model plans over an operator vocabulary, a symbolic engine executes the numerics, and a compiler renders the proof. On a 218-problem suite it certifies 191/218 (87.6%) with a 135M backbone; the 7B DeepSeek-Prover-V1.5-RL certifies 18/218 at 27.5 times the tokens and 37.5 times the wall-clock, and DeepSeek-Prover-V2-7B certifies zero on a bidirectional dual suite. Multi-step thinking costs 6.12 discrete decision actions per problem and produces no natural-language thinking text. Replacing the lagged engine state in the decision frame with the current one lifts the pass rate from 117/218 to 191/218, while up-weighting the chain-end loss hurts. On the public Lean-Workbook library, 2,132 of 3,536 gradeable admissible problems certify (60.29%) with zero regression on the main suite. All readings come from a version-frozen review with independent rechecks and reverse verification. Restricting free generation and keeping decision-time information visible is a more direct route than enlarging the model.

Hanchao Zhou, Jialei Li

Large Language Model-Guided Theorem Proving Automated Theorem Proving Lean Theorem Proving Efficient Language Model Reasoning

Ask PDF

Jan 22, 2026

Inference-Time Diversity in RL-Trained Lean Theorem Provers: A Diagnostic Study

RL-trained Lean theorem provers mode-collapse at inference time: on miniF2F-test with DeepSeek-Prover-V1.5-RL, doubling the i.i.d.\ sampling budget from

k{=}32

k{=}64

produces zero additional solved theorems (42/244 in both cases). A fixed schedule of 15 tactic skeletons breaks this plateau and recovers a

+45%

relative improvement at

k{=}16

(mean

Δ= +12.3 \pm 4.2

theorems across

n{=}3

seeds, sign preserved in every seed). A controlled diversity ablation rules out the prompt-diversity confound: tactic skeletons help, paraphrases match the baseline, and irrelevant Lean comments actively degrade. A leave-one-out formalization-difficulty stratification reveals a structural-content gradient across the three perturbations. The phenomenon is RL-specific: V1.5-Base proves zero theorems regardless of intervention, identifying RL as the stage that creates the proof capability which subsequently collapses; extending to two additional 7B Lean provers, RL-trained DeepSeek-Prover-V2-7B contributes

+3

frontier solves no i.i.d. baseline can reach despite a flat aggregate, while SFT-trained Goedel-Prover does not (

-10.0

\pm 4.4

theorems,

n{=}3

, sign preserved every seed). Inference-time structural diversity is a cheap, complementary axis for RL-trained provers, orthogonal to scaling model size or training compute.

Zachary Burton

Massachusetts Institute of Technology, Cambridge, MA, USA.

Inference-Time Search RL for Language Model Reasoning Lean Theorem Proving Large Language Model-Guided Theorem Proving

Ask PDF

May 26, 2026

MerLean-Prover: A Recursive Looping Harness for Lean 4 Theorem Proving

MerLean-Prover is an end-to-end Lean4 theorem prover that replaces sorry declarations with kernel-checkable proofs. It is built from three agent types (Planning, Check, and Lean) composed by a recursive outer loop whose unit of revision is the proof plan itself, and uses no fine-tuning, no custom RL objective, and no theorem-specific scaffolding. On FormalQualBench, a benchmark of 23 PhD-qualifying-exam theorems, MerLean-Prover solves 10/23, surpassing the strongest published open-source baseline (OpenGauss, 8/23). On Putnam2025, the same harness closes 12/12 with substantially lower total wall-clock than the next-best system that closes the full set. The harness also transfers to smaller models: Sonnet closes all four tested FormalQualBench problems, and Haiku closes the two short ones. These results suggest that harness design is a central factor in end-to-end Lean4 theorem proving, alongside raw model capability, and that a relatively simple harness can already be effective.

Jinzheng Li, Zeru Zhu, Yuanjie Ren

1Northeastern University · 2Stony Brook University · 3Massachusetts Institute of Technology

Large Language Model-Guided Theorem Proving Automated Theorem Proving Lean Theorem Proving

Ask PDF

Abstract

Lean Theorem Proving

Explore similar work

SCOPE: Certified Theorem Proving with a Language Model as the Policy Planner

Inference-Time Diversity in RL-Trained Lean Theorem Provers: A Diagnostic Study

MerLean-Prover: A Recursive Looping Harness for Lean 4 Theorem Proving