Hierarchical Attention Generates Better Proofs

TL;DR

This work tackles the difficulty of neural theorem provers capturing the inherent hierarchy in formal proofs. It introduces Hierarchical Attention, a regularization that enforces a five-level structure from $T_0$ (context) to $T_4$ (goal) and governs attention flow with a flow loss that respects level relations. Across miniF2F and ProofNet, the method yields higher proof-pass rates and shorter proofs, with notable reductions in proof length and improved attention patterns confirmed by ablations and visualization. The approach demonstrates that incorporating explicit hierarchical structure into LLM-based proof generation can bridge neural reasoning with formal mathematical rigor, offering practical gains and reusable code for reproducibility.

Abstract

Large language models (LLMs) have shown promise in formal theorem proving, but their token-level processing often fails to capture the inherent hierarchical nature of mathematical proofs. We introduce \textbf{Hierarchical Attention}, a regularization method that aligns LLMs' attention mechanisms with mathematical reasoning structures. Our approach establishes a five-level hierarchy from foundational elements to high-level concepts, ensuring structured information flow in proof generation. Experiments demonstrate that our method improves proof success rates by 2.05\% on miniF2F and 1.69\% on ProofNet while reducing proof complexity by 23.81\% and 16.50\% respectively. The code is available at https://github.com/Car-pe/HAGBP.

Figures (3)

• Figure 1: Overview of our hierarchical attention framework. Left: The five-level hierarchy from inner (context) to outer (goal) layer, illustrating the natural information flow in mathematical reasoning. Right: A concrete example showing how different components in a theorem proving state are assigned to hierarchical levels, with guided and unrestricted flow (solid arrows) representing allowed attention paths and limited flow (dashed arrows) representing restricted attention paths.
• Figure 2: Performance comparison between our method and baseline at $K = 64$. Left: Pass rate comparison across miniF2F (test/valid) and ProofNet (test/valid) datasets. Best-first search (BFS) consistently outperforms single-pass sampling (SPS), with our method further enhancing BFS performance. Solid bars represent our method while transparent bars represent the baseline. Right: Proof complexity ratio ($R_{avg}$), where values below 1.0 (dashed line) indicate more concise proofs. Our method with BFS achieves consistent complexity reductions across all datasets.
• Figure 3: Attention distribution analysis in different layers. Left: Hierarchy-constrained layers (where $\alpha_l \neq 0$). Right: Unconstrained layers (where $\alpha_l = 0$). This visualization is derived from averaging attention patterns across all evaluation samples on the LeanDojo Benchmark 4 test set. The x-axis represents different hierarchical levels, while the y-axis shows the percentage of attention scores, combining both cases where the level's tokens serve as source ($t_i$) and target ($t_j$). Blue and green bars represent the baseline and our method respectively, with different transparency levels indicating different attention flow types based on the relationship between source level($t_i$) and target level($t_j$).