YuxuanGong/lean-RAG

BGE-M3 Mathlib-RAG

README_en:https://huggingface.co/YuxuanGong/lean-RAG/blob/main/README_en.md

基于 BAAI/bge-m3 微调的 Mathlib4 API 检索模型，专为 Lean 4 形式化证明场景设计。该模型能够根据当前证明目标（goal type）和数学关键词，从 Mathlib4 的 32.6 万条声明中检索最相关的引理/定理/定义。

在我的github主页可访问完整skill：https://github.com/Tangtaizong-BUAA/Lean-formulizer

模型用途

该模型是 proof-formalizer 技能链路的检索核心。在形式化一个数学证明时：

1. 从 Lean 环境中提取当前 ⊢ 目标类型（goal conclusion）
2. 结合人类提供的数学关键词
3. 用本模型编码为 1024 维向量
4. 在预构建的 FAISS 索引中检索 top-k 最相关的 Mathlib4 声明

支持中英文混合查询，训练数据包含中文数学描述和 Lean 代码。

使用方式

from sentence_transformers import SentenceTransformer

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

# 查询：Lean goal + 关键词
query = "x % y < y  nat mod lt"
# 候选：Mathlib4 声明（从 declarations.jsonl 构建的 FAISS 索引）
passages = [
    "Keywords: nat mod lt theorem | Name: Nat.mod_lt | Kind: theorem | ...",
    "Keywords: nat mod add theorem | Name: Nat.ModEq.add | Kind: theorem | ...",
    # ... 32.6 万条
]

query_emb = model.encode(query)
passage_embs = model.encode(passages)
similarities = model.similarity(query_emb, passage_embs)

与 proof-formalizer Skill 集成

# Skill 中的 mathlib_lookup.sh (Route 4) 使用本模型:
DENSE_QUERY="${GOAL_TYPE} ${KEYWORDS_TEXT}"
# 然后调用 search.py 用本模型编码查询 -> FAISS 搜索

模型架构

基于 BGE-M3 (XLM-RoBERTa) 架构：

参数	值
基座模型	BAAI/bge-m3
隐藏层维度	1024
Transformer 层数	24
注意力头数	16
中间层维度	4096
词表大小	250,002
最大位置编码	8,194
输出维度	1024
池化方式	CLS token
相似度函数	Cosine Similarity

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

训练数据

从 Mathlib4 的 325,959 条声明中构造三元组训练数据：

• 训练集: 323,977 条 (query, positive, negatives)
• 验证集: 500 条
• 格式: BGE 三元组格式 — 每条样本包含 1 个 query、1 个 positive passage、7 个 negative passages
• Query: 定理名 + 类型签名（中英混合）
• Positive: 对应引理的完整结构化描述（Keywords | Name | Kind | Namespace | Doc | Type）
• Negatives: 随机采样 + 困难负样本

数据示例

角色	内容
Query	theorem length_take_of_le_length | (s.take n).length = n
Positive	`Keywords: le find iff
Negative	`Keywords: length take le' list

训练细节

超参数	值
损失函数	MultipleNegativesRankingLoss (cosine, scale=20.0)
训练轮数	2 epochs
训练步数	6,750
Per-device Batch Size	24
梯度累积	4 (有效 batch size = 96)
学习率	1e-5
LR 调度	cosine
Warmup	5%
优化器	AdamW (fused)
精度	FP16
随机种子	42
GPU	RTX PRO 6000 Blackwell 98GB
训练时间	~3.7 小时

训练 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

评测结果

直接检索（Direct Retrieval）

给定声明名 + docstring，检索对应的 Mathlib 声明：

指标	Baseline (BGE-M3)	Fine-tuned (本模型)
Recall@1	-	29.4%
Recall@5	-	58.6%
Recall@10	-	66.2%
MRR	-	41.2%

Skill Mode F：引理检索（Lemma Conclusion → Find Lemma）

模拟真实使用场景 — "我有这个 goal，找到能证它的引理"：

指标	值
查询数	626
Recall@1	26.7%
Recall@5	61.5%
Recall@10	71.9%
Recall@20	80.5%
MRR	42.1%

Skill Mode E：证明初始检索（Initial Goal → Proof Lemmas）

模拟 "开始证这个定理，需要哪些引理"（更难任务）：

指标	值
查询数	493
Recall@10	13.8%
Recall@20	18.3%
MRR	4.7%

已知限制

• Mode E 召回率较低：从零开始的证明检索仍具有挑战性，MRR 仅 4.7%
• 跨版本 Mathlib 兼容性：模型训练于特定 Mathlib4 版本，API 可能随版本变化
• 最大序列长度：虽然 BGE-M3 支持 8192 tokens，训练时截断至 400 tokens（适配 declaration 长度）
• 中英文偏好：英文查询效果通常优于中文

环境依赖

pip install sentence-transformers>=5.4.1

框架版本（训练时）

• Python: 3.12.3
• Sentence Transformers: 5.4.1
• Transformers: 5.5.4
• PyTorch: 2.8.0+cu128
• Accelerate: 1.13.0

引用

如果你使用了这个模型，请引用：

@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}}
}

基座模型引用：

@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}
}

许可

Apache 2.0