JEPA representations improve one-step tactic ranking in theorem proving
Predicting the Next State Is Not Enough: JEPA Representations for Lean Theorem Proving
Machine Learning
Summary
When computers try to prove math statements, they need to pick the right next move from many possibilities. The authors tested a new method called JEPA that predicts how promising the next steps might be, aiming to pick better moves. JEPA was better at ranking these moves correctly right after one step than a common method, but this didn’t always lead to solving more problems overall. This shows that simply picking the best next step isn’t enough to finish longer proofs in complex math tasks.
What this means in practice
- •For automated theorem proving teams: Optimize tactic selection in automated proof systems by improving one-step successor ranking accuracy with JEPA-style representations.
- •For formal methods engineers: Incorporate learned successor state representations to refine exploration ordering when verifying complex software and hardware properties.
Authors
Aarnav Choudhary
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.