Magnushammer: A Transformer-Based Approach to Premise Selection

TL;DR

Magnushammer employs contrastive learning with Transformer to improve premise retrieval, achieving 59.5% success on PISA with 4x fewer parameters.

cs.LG 🔴 Advanced 2023-03-08 34 views
Maciej Mikuła Szymon Tworkowski Szymon Antoniak Bartosz Piotrowski Albert Qiaochu Jiang Jin Peng Zhou Christian Szegedy Łukasz Kuciński Piotr Miłoś Yuhuai Wu
Automated Theorem Proving Transformer Contrastive Learning Premise Selection Mathematical Reasoning

Key Findings

Methodology

Magnushammer utilizes a hierarchical two-stage retrieval process: first, the SELECT module encodes proof states and premises into a shared embedding space, selecting the top 1024 premises based on cosine similarity via a contrastive InfoNCE loss. Next, the RERANK module performs contextualized re-ranking by jointly encoding proof state and premise pairs, producing relevance scores through a transformer-based interaction. The system is trained on a large dataset of over 4.4 million (proof state, premise) pairs, combining supervised signals from human and Sledgehammer-generated proofs. During inference, the model efficiently retrieves premises with a single pass, then refines relevance scores for precise selection.

Key Results

  • On the PISA benchmark, Magnushammer achieves a proof success rate of 59.5%, outperforming Sledgehammer’s 38.3%. On miniF2F, it reaches 34.0%, surpassing 20.9% baseline. When combined with the neural-symbolic prover Thor, the success rate on PISA improves from 57.0% to 71.0%, with a 4x reduction in model parameters. The approach demonstrates robustness across different compute budgets, especially in moderate resource scenarios.
  • The large-scale textual dataset and contrastive training enable the model to learn semantic relevance beyond superficial text similarity, leading to significant performance gains. The hierarchical retrieval balances speed and accuracy, making the system scalable to hundreds of thousands of premises. Ablation studies confirm the importance of the two-stage process and dataset scale.
  • In multi-step proof scenarios, integrating Magnushammer into Thor yields state-of-the-art results, confirming its effectiveness for complex reasoning tasks. The model maintains high performance with limited parameters, indicating high efficiency and potential for deployment in real-world proof assistants.

Significance

This work advances automated theorem proving by replacing symbolic, handcrafted premise selection with a data-driven, neural approach. The contrastive Transformer framework effectively leverages textual proof descriptions, reducing engineering overhead and broadening applicability across different proof systems. The large, publicly available dataset fosters further research, and the demonstrated success in formal proof benchmarks highlights the method's potential to accelerate formal mathematics and AI reasoning. The approach addresses longstanding challenges of scalability, generalization, and resource efficiency, paving the way for more autonomous and intelligent proof systems.

Technical Contribution

The paper introduces a novel contrastive learning framework for premise embedding using Transformer models, with a hierarchical retrieval architecture combining fast cosine similarity-based filtering and contextualized re-ranking. It constructs the largest open-source dataset for Isabelle premise selection, integrating human and Sledgehammer-generated proofs. The approach achieves high accuracy with significantly fewer parameters, demonstrating the effectiveness of textual representations and contrastive training in formal reasoning tasks. These innovations open avenues for scalable, data-efficient AI reasoning systems.

Novelty

This is the first study to integrate contrastive learning with Transformer-based premise retrieval in formal theorem proving. The hierarchical two-stage retrieval process, combining rapid embedding-based filtering with contextual re-ranking, surpasses previous heuristic and symbolic methods. Additionally, the creation and open release of the largest premise selection dataset for Isabelle uniquely support deep learning approaches in formal logic, marking a significant step forward in AI-assisted mathematics.

Limitations

  • Despite strong performance, the model struggles with highly complex or ambiguous proof states, especially when information is sparse or multiple reasoning paths exist. Transferability to other proof assistants or logics requires additional adaptation. Although training is data-efficient relative to model size, large-scale pretraining still demands substantial computational resources, limiting accessibility. The generalization to novel or less-represented proof scenarios remains to be validated, and further robustness improvements are needed.

Future Work

Future research will explore multi-modal inputs combining textual and structural proof information, aiming to improve understanding of complex proofs. Extending the approach to other proof systems like Coq or Lean, and developing more efficient models for real-time inference, are key directions. Incorporating active and continual learning strategies could enable the system to adapt dynamically to expanding mathematical knowledge bases. Additionally, integrating symbolic reasoning modules with neural retrievers may further enhance proof automation capabilities.

AI Executive Summary

Automated theorem proving faces a persistent challenge: how to efficiently and accurately select relevant premises from vast knowledge bases to support proof construction. Traditional symbolic methods, while precise, rely heavily on handcrafted features and domain-specific engineering, making them labor-intensive and less adaptable. Recent advances in deep learning, particularly Transformer models, have opened new avenues for addressing this bottleneck.

This paper introduces Magnushammer, a novel premise selection framework that leverages contrastive learning and hierarchical retrieval strategies. The core idea is to encode proof states and premises into a shared semantic space using a Transformer backbone pre-trained on large textual datasets. The first stage, SELECT, rapidly filters the top 1024 premises based on cosine similarity, enabling quick retrieval. The second stage, RERANK, performs a detailed contextualized scoring by jointly encoding proof state-premise pairs, refining the selection.

The approach is trained on a massive dataset of over 4.4 million (proof state, premise) pairs, combining human and Sledgehammer-generated proofs, which ensures diversity and robustness. Experimental results on the PISA and miniF2F benchmarks demonstrate that Magnushammer achieves 59.5% and 34.0% proof success rates respectively, outperforming the widely used Sledgehammer tool (38.3% and 20.9%). When integrated with the neural-symbolic prover Thor, the success rate on PISA jumps to 71.0%, with a fourfold reduction in model size.

These results highlight the potential of deep neural models to revolutionize formal reasoning by reducing reliance on handcrafted heuristics and domain knowledge. The large, open-source dataset further accelerates research in this area, providing a valuable resource for future developments. The hierarchical retrieval design balances speed and accuracy, making the system scalable to large proof libraries.

Despite its success, the model still faces challenges with highly complex or ambiguous proofs, and transferability across different proof systems remains an open question. Future work will focus on multi-modal inputs, broader proof system adaptation, and integration with symbolic reasoning modules. Overall, Magnushammer marks a significant step toward more autonomous, efficient, and scalable automated theorem proving systems.

Deep Dive

Abstract

This paper presents a novel approach to premise selection, a crucial reasoning task in automated theorem proving. Traditionally, symbolic methods that rely on extensive domain knowledge and engineering effort are applied to this task. In contrast, this work demonstrates that contrastive training with the transformer architecture can achieve higher-quality retrieval of relevant premises, without the engineering overhead. Our method, Magnushammer, outperforms the most advanced and widely used automation tool in interactive theorem proving called Sledgehammer. On the PISA and miniF2F benchmarks Magnushammer achieves $59.5\%$ (against $38.3\%$) and $34.0\%$ (against $20.9\%$) success rates, respectively. By combining \method with a language-model-based automated theorem prover, we further improve the state-of-the-art proof success rate from $57.0\%$ to $71.0\%$ on the PISA benchmark using $4$x fewer parameters. Moreover, we develop and open source a novel dataset for premise selection, containing textual representations of (proof state, relevant premise) pairs. To the best of our knowledge, this is the largest available premise selection dataset, and the first one for the Isabelle proof assistant.

cs.LG cs.AI cs.LO