LemmaHead: RAG Assisted Proof Generation Using Large Language Models
Tianbo Yang, Mingqi Yan, Hongyi Zhao, Tianshuo Yang
TL;DR
The paper tackles the difficulty of generating and validating formal mathematical proofs with large language models by introducing LemmaHead, a retrieval-augmented generation (RAG) knowledge base that injects authoritative textbook knowledge into Lean-proof generation. It explores three pipelines—basic RAG, Enhanced Query Generation (EQG), and Iterative Proof Augmentation (IPA)—to augment GPT-4 with structured mathematical context from a CHROMA-backed repository. Empirical results on the MiniF2F benchmark show that IPA-based RAG markedly improves proof correctness (Pass@1) compared with baselines and competitive models, while naive RAG can harm performance if retrieval quality is poor. The work highlights practical benefits for formal proof automation and outlines limitations around long proofs and inference cost, with future directions including testing additional LLMs and scaling IPA iterations.
Abstract
Developing the logic necessary to solve mathematical problems or write mathematical proofs is one of the more difficult objectives for large language models (LLMS). Currently, the most popular methods in literature consists of fine-tuning the model on written mathematical content such as academic publications and textbooks, so that the model can learn to emulate the style of mathematical writing. In this project, we explore the effectiveness of using retrieval augmented generation (RAG) to address gaps in the mathematical reasoning of LLMs. We develop LemmaHead, a RAG knowledge base that supplements queries to the model with relevant mathematical context, with particular focus on context from published textbooks. To measure our model's performance in mathematical reasoning, our testing paradigm focuses on the task of automated theorem proving via generating proofs to a given mathematical claim in the Lean formal language.