TopoAlign: A Framework for Aligning Code to Math via Topological Decomposition

AI summary

Overview

Research area: Natural Language Processing, specifically autoformalisation — translating informal mathematical statements into formal languages such as Lean 4 — and the use of code repositories as training data for mathematical LLMs.

Technical level: Advanced. The paper assumes familiarity with LLM pretraining/fine-tuning, formal proof assistants, dependency graphs, and autoformalisation metrics (typechecking, bidirectional equivalence, pass@k).

Scope: The paper introduces TopoAlign, a framework that decomposes Python code into docstrings, main functions, and dependency functions, reassembles them into structures that mirror Lean 4 statements, and trains three base models on this aligned data alongside formal math data on a task called code autoformalisation (CAF).

What This Paper Is About

Math LLMs can reason in natural language and in formal languages such as Lean 4, but they remain weak at autoformalisation because large, high-quality corpora pairing informal statements with formal ones are scarce. The authors observe that formal languages share a structural resemblance to programming languages, and that code is available at scale, but existing code-trained models do not transfer well to formal math because of structural and syntactic differences. TopoAlign addresses this mismatch by decomposing code into components that mirror the anatomy of formal mathematical statements and training models on the reassembled data.

Key Contributions

1. TopoAlign, a framework that structurally aligns code with formal mathematical languages to address the shortage of training corpora for Math LLMs, using a function-level topological dependency parser (breadth-first search over abstract syntax trees) rather than file-level dependency extraction.
2. Code Autoformalisation (CAF), a training task that uses the structurally aligned code dataset to emulate autoformalisation, treating docstrings as informal statements, main functions as formal statements, and dependency functions as supporting lemmata, thereby reducing dependence on annotated informal–formal statement pairs.
3. Release of a 300 million token large-scale semi-synthetic pre-training dataset (the paper reports the filtered corpus as 156,684 functions comprising 324.5 million tokens) of structurally aligned code for autoformalisation tasks, available on Hugging Face.
4. Ablation studies showing that a balanced ratio of aligned code data to formal mathematical statements yields optimal autoformalisation performance (α = 0.5 outperforming α = 0.25 and α = 0.75).

Main Findings

• DeepSeek-Math gains: TopoAlign improves DeepSeek-Math by 17.77% on BEq@10 and 68.82% on typecheck@10, and yields relative BEq improvements of 36.7% for DeepSeek-Math and 6.2% for Herald according to the introduction.
• Herald gains: Herald improves by 0.12% on BEq@10 and 1.09% on typecheck@10, despite TopoAlign introducing no new mathematical knowledge.
• Broad benchmark success: TopoAlign outperforms the benchmarked settings across most datasets and model families on the BE@10 metric, with the exception of Herald on ProofNet.
• Best average rank: TopoAlign achieves the best average rank across datasets for DeepSeek-Math (1.13), Qwen-3 (1.63), and Herald (1.69), compared with baseline, math-only, and code-only settings.
• Unaligned code hurts: Training Herald on unaligned code only degraded its BEq scores and typecheck accuracy, with generated outputs frequently syntactically invalid, indicating raw code data may actively degrade an already capable formaliser.
• Degenerate outputs from code-only training: Qwen-3 trained on raw code frequently generated degenerate outputs such as endless type definitions (for example Nat to Nat) that pass typechecking but are semantically meaningless.
• Error profile: In a qualitative analysis of 200 failure cases (50 per dataset) from Qwen-3 with TopoAlign, Type 1 errors (missing or incorrect conditions) accounted for approximately 70% of all failures and over 90% on ProofNet and Putnam, identifying deep problem comprehension as the main bottleneck.
• Balanced data mix is best: DeepSeek-Math trained with α = 0.5 reached 51.75 TC@1 and 14.47 BEq@1 on MiniF2F-valid, versus 44.74/11.84 at α = 0.25 and 52.63/10.96 at α = 0.75; α = 0.25 was consistently weakest.
• Base model matters: The DeepSeek-Math base model without autoformalisation pretraining failed to generate meaningful outputs, while Herald's CAF training improved BEq@1 on MiniF2F-valid from 23.07% to 25.08% and on MiniF2F-test from 22.70% to 25.64% in the unbiased pass@1 results.
• Research-level limits: On FATE-H, all models yielded less than BEq of 0.05, and continued pre-training with TopoAlign data did not improve BEq for Goedel-Formalizer, which was trained on Mathlib.

Methodology in Plain English

The authors start from an analogy: a programming function that solves a subproblem using helper functions looks structurally similar to a formal mathematical statement that is expressed using lemmata and theorems from a library. They extract Python repositories from the Stack v2 dataset and, instead of working file by file, build function-level dependency trees with a parser that uses abstract syntax trees and breadth-first search to trace calls, parent definitions, methods, recursive calls, and imports. Each tree yields a docstring (treated as the informal problem statement), the root function (treated as the formal statement), and the called dependency functions (treated as supporting lemmata). Repositories are filtered so that dependency trees have a maximum depth between 3 and 6 and a maximum sibling count between 3 and 10, excluding both overly complex codebases and simple scripts. Missing or low-quality documentation is filled in with Qwen-3 generating concise summaries of a function's implementation logic, and a 10-gram contamination analysis is used to check for leakage against the test set.

Training combines two data streams: the aligned code data under the CAF objective and formal mathematical statements from the Herald Statements corpus, with dependency functions for the math side extracted using the jixia library. A mixing parameter α controls the fraction of math samples and 1 − α the fraction of code samples, with the loss a weighted sum of the math loss and the CAF loss, each computed by next-token prediction as negative log-likelihood. For each training instance the model receives an input (a code docstring or informal math statement) plus dependencies and must generate the target (the main function or the formal statement). Four settings are compared — Baseline, Math, Code, and TopoAlign — with the same number of training samples across settings; TopoAlign uses 4,000 randomly selected samples from GitHub repositories and 4,000 from Herald Statements at α = 0.5. Evaluation is on MiniF2F (validation and test), ProofNet, and Putnam, using typechecking with the Lean 4 compiler (v4.21.0) and bidirectional equivalence (BEq), reported under pass@1 and pass@10.

Why This Matters

Impact on research: The work reframes code repositories as a scalable and structurally compatible data source for formal mathematics, rather than something merely useful for model initialisation. It offers a concrete alternative to synthetic informal–formal pair generation, which is bounded by the quality of the underlying generative models, and its rank-based results indicate the approach complements training on formal math data rather than replacing it.

Real-world applications:

• Building and maintaining formalised mathematical libraries, where statements must be translated into Lean 4 or similar systems before proofs can be checked.
• Automated verification pipelines that rely on proof assistants such as Isabelle or Lean 4 to enforce rule-based logical consistency.
• Educational tools that let students state problems informally and receive machine-checkable formal versions.
• Scaling formal methods work in settings where domain experts are scarce, by reducing the human annotation required to produce formal statements.

Industry relevance: The framework targets the pre- and mid-training phases of LLM development, using about 28% of Stack-V2 (Python) as its starting corpus and leaving room to extend to other programming languages. Companies with access to large internal codebases and pretraining infrastructure can apply the decomposition and alignment procedure without new mathematical content. The paper also notes that TopoAlign is orthogonal to fine-tuning-based approaches, meaning it can be layered onto existing training pipelines rather than competing with them.

Future Directions

• Scaling the aligned corpus beyond the 4,000 code and 4,000 math samples used in the main experiments, and extending the approach to programming languages beyond Python from Stack-V2.
• Applying the same decomposition and alignment idea to related tasks such as theorem proving, which the authors identify as a possible extension.
• Improving the handling of implicit natural language constraints, since Type 1 errors (missing or incorrect conditions) dominate failures, especially on ProofNet and Putnam where they average over 90% of failures.
• Investigating how to complement rather than substitute for high-quality formal data, given that TopoAlign produced less than BEq of 0.05 on FATE-H and did not improve BEq for Goedel-Formalizer.
• Determining whether retraining state-of-the-art fine-tuned formalisers with TopoAlign data under identical conditions would change the comparison, which the authors state is beyond the scope of this study.

Target Audience

Researchers and engineers working on LLM-based mathematical reasoning, autoformalisation, and neuro-symbolic systems that pair language models with proof assistants such as Lean 4 or Isabelle. It is also relevant to practitioners building training pipelines for Math LLMs, and to those interested in cross-domain transfer between code and formal mathematics. Readers unfamiliar with proof assistants, type systems, or LLM training objectives will need background reading first.

Authors’ abstract

Large Language Models (LLMs) excel at both informal and formal (e.g. Lean 4) mathematical reasoning but still struggle with autoformalisation, the task of transforming informal into formal mathematical statements. Yet, the performance of current Math LLMs is constrained by the scarcity of large-scale corpora, particularly those containing pairs of informal and formal statements. Interestingly, the formal languages used in autoformalisation share structural similarities with programming languages, and code data is available at scale. However, current models trained on code do not transfer effectively to formal math, due to structural and syntactic differences between them. To address this, we propose TopoAlign, a framework that unlocks widely available code repositories as training resources for Math LLMs. TopoAlign decomposes code into docstrings, main functions, and dependency functions, and reassembles these components into analogues that structurally mirror formal statements. We train three state-of-the-art models, DeepSeek-Math, Qwen-3 and Herald, and evaluate them on the MiniF2F, Putnam, and ProofNet benchmarks. TopoAlign provides substantial gains for DeepSeek-Math, improving performance by 17.77% on BEq@10 and 68.82% on typecheck@10, and also measurably improves Herald by 0.12% on BEq@10 and 1.09% on typecheck@10 despite introducing no new mathematical knowledge.