Hierarchical Attention Generates Better Proofs

本文提出层次注意力正则化方法，通过引导大型语言模型的注意力机制与数学推理的五级层次结构对齐，在 miniF2F 和 ProofNet 基准上分别提升证明成功率 2.05% 和 1.69%，并显著降低证明复杂度。

Large Language Model, Reasoning, Hierarchical Attention, Pre-training, Fine-tuning, AI for Science

Jianlong Chen, Chao Li, Yang Yuan, Andrew C Yao

The Chinese University of Hong Kong, Shenzhen, Shanghai Qi Zhi Institute, IIIS, Tsinghua University

Generated by grok-3

Background Problem

近年来，人工智能与数学的交叉领域成为重要研究方向，尤其是在形式化定理证明方面。尽管大型语言模型（LLMs）在生成数学证明和解决广泛问题方面表现出色，但其基于令牌序列的处理方式往往无法捕捉数学证明固有的层次结构，导致在生成复杂证明时失败或生成冗长的证明。本文旨在解决这一关键问题：如何让 LLMs 更好地理解和利用数学推理中的结构化信息流，以提高证明的成功率和简洁性。

Method

本文提出了一种名为层次注意力（Hierarchical Attention）的正则化方法，旨在通过引导 LLMs 的注意力机制与数学推理的层次结构对齐来改进证明生成。具体方法如下：

核心思想：数学推理具有从基础元素到高级概念的自然层次结构，作者将其形式化为五级层次（从上下文到目标），并通过控制注意力分布来模拟这种信息流。
实现步骤：
1. 提取层次模式：将输入定理的令牌序列解析为五个层次（上下文、案例、类型、实例、目标），基于 Lean 语言的语义特征。
2. 注意力约束：通过注意力掩码（Attention Mask）控制信息流，允许高层次令牌访问同级或低层次信息，但限制低层次令牌访问高层次信息。
3. 损失函数设计：引入流动损失（Flow Loss） $\mathcal{L}_{flow} = \frac{1}{|T|} \sum_{l=1}^{L} \alpha_l \cdot \sum_{i,j} \text{ReLU}(att_l(t_i, t_j) \cdot (1 - M_{ij}))$ ，惩罚违反层次约束的注意力分布，并结合标准交叉熵损失进行训练，其中 $\alpha_l = 1 - l/L$ 为层级适应因子，在早期层施加更强约束。
4. 层级适应：通过层级适应因子在较深层逐渐放松约束，允许更灵活的注意力分布以适应复杂推理。
关键点：该方法不修改模型架构，仅通过训练时的注意力正则化实现层次结构的内化，旨在保持模型的通用性同时提升其对数学结构的理解。

Experiment

本文在多个定理证明基准数据集上进行了实验评估，具体设置和结果如下：

数据集与配置：使用 LeanDojo Benchmark 4 作为训练数据集，在 Pythia-2.8B 模型上进行微调，并在 miniF2F 和 ProofNet 数据集的测试和验证集上评估。实验采用两种搜索策略：最佳优先搜索（Best-First Search, BFS）和单次采样（Single-Pass Sampling, SPS），计算预算为 $K \times T$ （ $T=100$ ）。
实验设计合理性：实验设置涵盖了不同计算预算（ $K$ 从 1 到 64），并通过 pass@K 准确率和证明复杂度（平均复杂度比 $R_{avg}$ ）两个指标评估性能。选择 LLMSTEP 作为基线，便于结果复现。
结果分析：
1. 在 miniF2F 测试集上，BFS 策略下成功率从基线的 29.51% 提升至 31.56%（提升 2.05%），证明复杂度降低 23.81%（ $R_{avg}=0.76$ ）；SPS 策略下成功率从 23.36% 提升至 27.87%（提升 4.51%）。
2. 在 ProofNet 测试集上，BFS 策略下成功率从 13.56% 提升至 15.25%（提升 1.69%），证明复杂度降低 16.50%（ $R_{avg}=0.84$ ）。
3. 注意力分布可视化表明，方法成功实现了层次信息流约束，目标层有效整合低层信息（高达 84.5% 的向上注意力）。
评价与不足：虽然成功率和证明简洁性有所提升，但提升幅度较小，尤其在复杂证明任务中可能不够显著。实验未涉及更先进的模型（如 DeepSeek-Prover），可能低估了方法在高性能模型上的潜力。此外，层次结构定义依赖 Lean 语义，通用性存疑，实验未充分探讨动态推理场景下的表现。

Further Thoughts

本文提出的层次注意力机制为 LLMs 在形式化定理证明中的应用提供了一个有趣的方向，但其固定五级层次结构可能限制了方法对更复杂数学推理场景的适应性。未来可以探索动态层次结构或基于图结构的注意力机制，以捕捉数学证明中非线性的依赖关系。此外，考虑到不同证明语言（如 Coq、Isabelle）的语义差异，跨语言的通用性测试将是重要研究方向。另一个值得思考的点是，该方法是否可以与其他技术（如强化学习或蒙特卡洛树搜索）结合，进一步优化证明搜索效率？与近期工作如 DeepSeek-Prover 的对比也可能揭示层次注意力在高性能模型上的潜力或局限性。最后，层次注意力机制或许不仅适用于数学推理，也可以扩展到其他需要结构化推理的领域，如法律文本分析或代码生成，值得跨领域探索其适用性。