README_en.md

English | 中文

BGE-M3 Mathlib-RAG

BGE-M3 Mathlib-RAG is a Mathlib4 API retrieval model fine-tuned from BAAI/bge-m3. It is designed for Lean 4 formal proof scenarios, where the model retrieves relevant lemmas, theorems, and definitions from 325,959 Mathlib4 declarations based on the current proof goal and mathematical keywords.

The complete skill implementation is available on GitHub:

https://github.com/Tangtaizong-BUAA/Lean-formulizer

Model Use

This model serves as the retrieval core of the proof-formalizer skill pipeline. During mathematical proof formalization, it can be used to retrieve relevant Mathlib4 declarations by combining the Lean goal type with human-provided mathematical keywords.

A typical workflow is:

1. Extract the current ⊢ goal type from the Lean environment.
2. Combine the goal type with human-provided mathematical keywords.
3. Encode the query into a 1024-dimensional dense vector using this model.
4. Search a pre-built FAISS index over Mathlib4 declarations to retrieve the top-k most relevant results.

The model supports mixed English, Chinese, and Lean/code-style queries. The training data includes Chinese mathematical descriptions, English keywords, Lean declaration names, and Lean type signatures.

Usage

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("YuxuanGong/lean-RAG")

# Query: Lean goal + mathematical keywords
query = "x % y < y  nat mod lt"

# Candidate passages: Mathlib4 declarations from a FAISS index
passages = [
    "Keywords: nat mod lt theorem | Name: Nat.mod_lt | Kind: theorem | ...",
    "Keywords: nat mod add theorem | Name: Nat.ModEq.add | Kind: theorem | ...",
    # ... 325,959 declarations
]

query_emb = model.encode(query)
passage_embs = model.encode(passages)

similarities = model.similarity(query_emb, passage_embs)

Integration with the proof-formalizer Skill

In the proof-formalizer skill, this model is used by mathlib_lookup.sh in Route 4:

DENSE_QUERY="${GOAL_TYPE} ${KEYWORDS_TEXT}"

The dense query is then passed to search.py, which encodes the query with this model and performs FAISS-based nearest-neighbor search over Mathlib4 declarations.

Model Architecture

This model is based on BGE-M3, which uses XLM-RoBERTa as its backbone.

Parameter	Value
Base model	BAAI/bge-m3
Backbone	XLM-RoBERTa
Hidden size	1024
Transformer layers	24
Attention heads	16
Intermediate size	4096
Vocabulary size	250,002
Maximum position embeddings	8,194
Output dimension	1024
Pooling strategy	CLS token
Similarity function	Cosine similarity

SentenceTransformer(
  (0): Transformer(XLMRobertaModel, task='feature-extraction')
  (1): Pooling(cls pooling, 1024d)
  (2): Normalize()
)

Training Data

The training data was constructed from 325,959 Mathlib4 declarations.

Split / Field	Value
Training set	323,977 triplets
Validation set	500 examples
Format	BGE-style triplet format
Each sample	1 query, 1 positive passage, 7 negative passages
Query	Theorem name + type signature, with mixed English/Chinese descriptions
Positive passage	Structured declaration description: `Keywords	Name	Kind	Namespace	Doc	Type`
Negative passages	Random negatives + hard negatives

Data Example

Role	Content
Query	theorem length_take_of_le_length | (s.take n).length = n
Positive	`Keywords: le find iff	Name: Nat le_find_iff	Kind: lemma	Namespace: Nat	Type: ...`
Negative	`Keywords: length take le' list	Name: List length_take_le'	Kind: theorem	Namespace: List`

Training Details

Hyperparameter	Value
Loss function	MultipleNegativesRankingLoss, cosine similarity, scale = 20.0
Epochs	2
Training steps	6,750
Per-device batch size	24
Gradient accumulation	4
Effective batch size	96
Learning rate	1e-5
LR scheduler	Cosine
Warmup	5%
Optimizer	AdamW, fused
Precision	FP16
Random seed	42
GPU	RTX PRO 6000 Blackwell 98GB
Training time	Approximately 3.7 hours

Training Loss

Epoch	Step	Loss
0.61	2,050	0.146
0.89	3,000	0.167
1.20	4,050	0.145
1.48	5,000	0.131
1.78	6,000	0.136
2.00	6,750	0.140

Evaluation Results

Direct Retrieval

This setting evaluates whether the model can retrieve the corresponding Mathlib declaration given a declaration name and docstring.

Metric	Baseline BGE-M3	Fine-tuned model
Recall@1	-	29.4%
Recall@5	-	58.6%
Recall@10	-	66.2%
MRR	-	41.2%

Skill Mode F: Lemma Retrieval

This setting simulates a realistic proof-assistant scenario: given a lemma conclusion or current proof goal, retrieve the lemma that can prove it.

Metric	Value
Number of queries	626
Recall@1	26.7%
Recall@5	61.5%
Recall@10	71.9%
Recall@20	80.5%
MRR	42.1%

Skill Mode E: Initial Proof Retrieval

This setting simulates a harder scenario: given only the initial theorem goal, retrieve useful proof lemmas required to begin the proof.

Metric	Value
Number of queries	493
Recall@10	13.8%
Recall@20	18.3%
MRR	4.7%

Known Limitations

• Low recall in Mode E: Initial proof retrieval remains challenging. The current MRR is only 4.7%.
• Mathlib version compatibility: The model was trained on a specific Mathlib4 snapshot. Some APIs may change across Mathlib versions.
• Maximum sequence length: Although BGE-M3 supports up to 8192 tokens, the model was trained with truncation at 400 tokens to fit declaration lengths.
• Language preference: English queries generally perform better than Chinese queries.
• Retrieval-only model: This model retrieves relevant declarations but does not generate Lean proofs by itself.

Dependencies

pip install "sentence-transformers>=5.4.1"

Training Framework Versions

Component	Version
Python	3.12.3
Sentence Transformers	5.4.1
Transformers	5.5.4
PyTorch	2.8.0+cu128
Accelerate	1.13.0

Citation

If you use this model, please cite:

@misc{gong2025mathlib-rag,
  title={Mathlib-RAG: Fine-tuned BGE-M3 for Lean 4 Mathlib API Retrieval},
  author={Gong, Yuxuan},
  year={2025},
  publisher={Hugging Face},
  howpublished={\url{https://huggingface.co/YuxuanGong/lean-RAG}}
}

Base model citation:

@article{chen2024bge,
  title={BGE M3-Embedding: Multi-Lingual, Multi-Functionality, Multi-Granularity Text Embeddings Through Self-Knowledge Distillation},
  author={Chen, Jianlv and Xiao, Shitao and Zhang, Peitian and others},
  journal={arXiv preprint arXiv:2402.03216},
  year={2024}
}

License

This model is released under the Apache 2.0 license.