{"slug": "high-throughput-lean-4-autoformalization-model-for-local-inference", "title": "High-Throughput Lean 4 Autoformalization Model for Local Inference", "summary": "Researchers developed a memory-efficient, single-GPU training and inference architecture for autoformalizing natural language mathematics into Lean 4 statement code, using the Qwen3-Coder-30B-A3B sparse Mixture-of-Experts model with whole-model NF4 quantization, custom parameter unfusing, GRPO with persistent Lean REPL compiler feedback, and dual-tier retrieval augmentation. The system fits the 30B-parameter model into 32 GB VRAM on an Nvidia RTX 5090ti, enabling high-throughput local inference for formal mathematics.", "body_md": "## High-Throughput Lean 4 Autoformalization Model for Local Inference\n\nA memory-efficient, single-GPU training and inference architecture for formalizing natural language mathematics into Lean 4 statement code. Incorporates whole-model NF4 quantization of a sparse Mixture-of-Experts (MoE) model, custom parameter unfusing, Group Relative Policy Optimization (GRPO) with persistent Lean REPL compiler feedback, and a dual-tier retrieval augmentation mechanism.\n\n### 1. Introduction\n\n#### 1.1 Autoformalization Bottleneck in Formal Mathematics\n\nAutoformalization, translating informal natural language mathematics into machine-checkable formal logic, is a primary bottleneck in formal verification. Modern interactive theorem provers, such as Lean 4, enforce rigid type-theoretic specifications. Minor syntax, namespace, or typeclass unification errors cause complete compilation failure. Manual translation is time-consuming, current autoformalization models are often expensive and error-prone.\n\n#### 1.2 Resource-Constrained Deployment\n\nState-of-the-art autoformalization models rely on dense multi-billion parameter LLMs requiring multi-GPU server clusters. Deploying fine-tuned models on consumer-grade single-GPU hardware presents severe memory limitations:\n\n1. Dense models (>14B parameters in 16-bit precision) exceed VRAM limits during training and generation.\n\n2. Standard 4-bit quantization libraries (such as `bitsandbytes`\n\n) fail on fused MoE weight tensors.\n\n3. Execution of sparse MoE models on single cards suffers from launch-bound per-expert CUDA kernels.\n\nEvaluation of autoformalization methods is non-trivial, as correct compilation does not imply successful translation. This makes training environment construction difficult. The lack of proper training data and an accurate evaluation engine are critical issues in the field.\n\nThis project aims to stretch the limits of a smaller transformer model at the domain of formal mathematics, test a set of fine tuning methods and assert their effectiveness, and eventually construct a multi-agent high-throughput system to rapidly iterate in problems.\n\n### 2. System Architecture\n\n#### 2.1 Base Model Selection\n\nThe base policy relies on `Qwen3-Coder-30B-A3B`\n\n, a sparse Mixture-of-Experts model containing 30B total parameters with ~3B active parameters per token across 48 transformer layers. Because Qwen-Coder is explicitly designed and pretrained for code generation and software synthesis, it has a better baseline comprehension of formal abstractions, strict type-theoretic semantics, and structured algorithmic logic required for interactive theorem proving compared to general-purpose language models. Furthermore, the sparse MoE architecture provides the expansive parameter capacity necessary to encode broad mathematical domain knowledge while constraining active forward-pass compute to ~3B parameters, this aids high throughput.\n\n#### 2.2 Memory Footprint & Quantization Engineering\n\nFor training and inference, a standard Nvidia RTX 5090ti was used.In standard `bfloat16`\n\nprecision, the 30B model requires ~60 GB VRAM, exceeding the 32 GB budget of a single GPU. To fit both the base model and optimization states into VRAM:\n\n1. Whole-Model NF4 Double-Quantization: The base model was loaded directly into host memory and quantized in-place using NormalFloat4 (NF4) with double quantization.\n\n2. GPU Allocation: Model parameters reside entirely on GPU VRAM (~16.7 GB post-load, peak ~26–30 GB during training with micro-batch size 2). CPU offloading was completely disabled to avoid PCI-e latency bottlenecks.\n\n#### 2.3 MoE Expert Unfusing Engineering\n\nThe standard implementation of `Qwen3Moe`\n\nstores expert weights in 3D fused tensors. Quantization libraries cannot process 3D fused parameters.\n\nTo enable whole-model 4-bit quantization and expert-specific adapters, a custome layer is introduced:\n\n1. Each layer’s fused expert block is decomposed into individual `_Expert`\n\nsubmodules containing standard `nn.Linear`\n\nlayers (`gate_proj`\n\n, `up_proj`\n\n, `down_proj`\n\n).\n\n2. Pretrained weights are copied from fused 3D tensors into the unfused `nn.Linear`\n\nmodules.\n\n3. Quantization is applied in-place via `bitsandbytes.nn.Linear4bit`\n\n.\n\n4. The forward pass is modified to execute expert routing over explicit linear layers without modifying the base model output semantics.\n\n``` python\nclass _Experts(nn.Module):\n    def __init__(self, config):\n        super().__init__()\n        self.num_experts = config.num_experts\n        self.act_fn = ACT2FN[config.hidden_act]\n        for j in range(self.num_experts):\n            self.add_module(str(j), _Expert(config.hidden_size, config.moe_intermediate_size))\n```\n\nPEFT (Parameter-Efficient Fine-Tuning) automatically converts `qwen3_moe`\n\nexpert target configurations into fused parameter names, missing custom unfused submodules. To bypass this, `model.config.model_type`\n\nis dynamically masked to `\"qwen3_moe_unfused\"`\n\nduring PEFT injection, permitting explicit target module matching.\n\n#### 2.4 LoRA vs. DoRA\n\nAdapter modules are attached across attention and expert layers:\n\n- Attention Modules (\n`q_proj`\n\n,`k_proj`\n\n,`v_proj`\n\n,`o_proj`\n\n): Rank (r = 64), (= 128). - Expert Modules (\n`gate_proj`\n\n,`up_proj`\n\n,`down_proj`\n\n): Rank (r = 8), (= 16). - Trainable Parameters: 468M parameters (~2.92% of total model parameters).\n- Adaptation Variants:\n- Standard QLoRA: Linear low-rank updates.\n- QDoRA (Weight-Decomposed Low-Rank Adaptation): Decomposes weight updates into directional vectors and magnitude scalars.\n\n```\n                       [ Informal Mathematical Statement ]\n                                        │\n                                        ▼\n             ┌─────────────────────────────────────────────────────┐\n             │ System B: Hybrid Retrieval Grounding (Premise RAG)  │\n             │  (Dense Slogan Embeddings + BM25 Lexical Rerank)    │\n             └──────────────────────────┬──────────────────────────┘\n                                        │ Augment Prompt Context\n                                        ▼\n             ┌──────────────────────────────────────────────────────┐\n             │ Qwen3-Coder-30B-A3B policy (NF4 Base + QLoRA Adapter)│\n             └──────────────────────────┬───────────────────────────┘\n                                        │ Sample Completion\n                                        ▼\n             ┌─────────────────────────────────────────────────────┐\n             │ Lean 4 REPL Worker Pool (Persistent Mathlib State)  │\n             └───────┬─────────────────────────────────────┬───────┘\n                     │ Type-Check Failure                  │ Type-Check Success\n                     ▼                                     ▼\n  ┌─────────────────────────────────────┐   ┌───────────────────────────────┐\n  │ System A: Compiler Repair Context   │   │ Formal Lean 4 Statement       │\n  │ (Extract Missing Identifiers/Types) │   │ (Type-Checked & Faithfulness  │\n  │ -> Re-prompt Policy (up to max_iter)│   │  Evaluated)                   │\n  └─────────────────────────────────────┘   └───────────────────────────────┘\n```\n\n*Figure 1: End-to-end architecture of the Lean 4 autoformalization framework. Informal mathematical statements pass through System B premise grounding prior to inference by the quantized MoE policy. Generated code is evaluated by a persistent Lean 4 REPL pool. Failed completions enter an agentic repair loop supported by System A identifier lookup.*\n\n### 3. Stage 1: Syntax Alignment\n\n#### 3.1 Dataset Preparation and Filtering\n\n- Corpus: Combined dataset from Herald and Lean-Workbook (~720,000 original informal-formal statement pairs).\n- Subsampling & Filtering: Subsampled to 40,000 clean pairs (mean sequence length: 198 tokens, p99: 470 tokens, max sequence length cap: 1024 tokens). Subsampling prevents overfitting to specific dataset phrasing while enforcing Lean 4 syntax formatting.\n- Prompt Structure: Standardized system prompt instructing the model to translate mathematical statements into Lean 4 theorem statements ending in\n`:= by sorry`\n\n.\n\n#### 3.2 Completion-Only Cross-Entropy Loss\n\nTo maximize parameter updates on syntax generation rather than prompt encoding, loss is computed exclusively on assistant completions:\n\n- Prompt tokens (system prompt + informal user input) are assigned a target label of\n`-100`\n\n. - Cross-entropy loss is evaluated only on output Lean 4 statement tokens:\n\n### 4. Stage 2: Reinforcement Learning with Lean Compiler Feedback (RLCF)\n\n#### 4.1 Group Relative Policy Optimization (GRPO)\n\nStage 1 produces syntactically valid code but struggles with deeper Lean 4 type-system semantics (such as typeclass instantiation and set coercion). Stage 2 uses GRPO to align output completions directly against Lean compiler responses.\n\nFor each informal statement prompt , the policy generates a group of completion outputs . The objective function is defined as:\n\nwhere the advantage is normalized within each group:\n\nand the KL penalty coefficient is set to .\n\n#### 4.2 Composite Gated Reward Function\n\nReward hacking is prevented by gating compiler feedback with a surface faithfulness metric relative to reference formalizations:\n\n- Well-formedness: Evaluated via heuristic AST structure checks (balanced delimiters, explicit theorem declaration keywords, presence of type signatures).\n- Compilation Check: Execution against Mathlib via Lean 4 REPL server.\n- Faithfulness Score: Sequence similarity ratio between canonicalized candidate and reference code (removing theorem identifiers and proof bodies).\n\n### 5. The A+B System\n\nAutoformalization failure often stems from hallucinated namespace prefixes or missing Mathlib declaration identifiers. The framework incorporates a dual-tier retrieval architecture operating without model weight modifications.\n\n#### 5.1 System B: Initial Premise Grounding\n\n- Function: Prepend relevant Mathlib theorem signatures and slogans to the initial informal prompt.\n- Mechanism: Hybrid dense-sparse retrieval:\n- Dense retrieval selects the top 100 declaration candidates using\n`SentenceTransformers`\n\nembeddings generated over Mathlib TheoremGraph slogans. - BM25 lexical reranking filters candidates down to the top (k=8) declarations, ensuring exact symbol match preservation.\n\n- Dense retrieval selects the top 100 declaration candidates using\n\n#### 5.2 System A: Compiler Error Context Repair\n\n- Function: Provide targeted context during repair iterations when compilation fails.\n- Mechanism: When the compiler returns an error matching\n`unknown identifier 'X'`\n\nor`unknown constant 'X'`\n\n, System A extracts symbol`X`\n\n, queries the Mathlib index using fuzzy substring matching, and injects candidate declarations (`Did you mean: ...`\n\n) directly into the repair prompt.\n\n```\n┌─────────────────────────────────────────────────────────────────────────────┐\n│                          System B: Premise Grounding                        │\n│                                                                             │\n│ Informal Query ──► SentenceTransformer ──► Dense Top-100 Candidates         │\n│                                                  │                          │\n│ Initial Prompt Context ◄── Mathlib Premises ◄── BM25 Lexical Reranking      │\n└─────────────────────────────────────────────────────────────────────────────┘\n\n┌─────────────────────────────────────────────────────────────────────────────┐\n│                      System A: Compiler Context Repair                      │\n│                                                                             │\n│ Compiler Error ──► Regex Ident Extractor ──► Candidate Lookup               │\n│ \"unknown identifier 'X'\"                          (Exact & Fuzzy Match)     │\n│                                                          │                  │\n│ Repair Prompt Context ◄── \"Did you mean: Y, Z\" ◄─────────┴──────────────────┘\n└─────────────────────────────────────────────────────────────────────────────┘\n```\n\n*Figure 2: Dual-tier retrieval pipeline. System B provides domain premise context prior to initial generation. System A dynamically inspects compiler error outputs during repair cycles to resolve identifier reference failures.*\n\n### 6. Agentic Compiler-Feedback Evaluation Loop\n\nDuring evaluation, formal statement generation is executed within a multi-turn agentic feedback loop:\n\n- Turn 1 Generation: The policy generates an initial theorem statement using System B prompt grounding.\n- Type-Check Verification: The statement is evaluated in\n`LeanREPLPool`\n\n. If it parses cleanly and type-checks without errors, it is marked as solved at iteration 1. - Iterative Repair Loop: If type-checking fails, compiler error logs (truncated to 800 characters) and System A candidate lookup suggestions are formatted into a repair message. The model generates a corrected statement. This cycle repeats up to\n`max-iters = 5`\n\n.\n\n### 7. Empirical Results\n\n#### 7.1 Quantitative Benchmark Results (ProofNet Test Set)\n\nPerformance across 6 model configurations evaluated on the ProofNet test set (n=100 for LoRA variants, n=50 for DoRA variants):\n\n| Metric | `qwen-lora-prosyntax` | `qwen-lora-proleanworkbook` | `qwen-lora-prominif2f` | `qwen-dora-postsyntax` | `qwen-dora-postleanworkbook` | `qwen-dora-postminif2f` |\n|---|---|---|---|---|---|---|\nAdapter Type | QLoRA | QLoRA | QLoRA | QDoRA | QDoRA | QDoRA |\nTraining Stage | Stage 1 (SFT) | Stage 2 (RLCF-WB) | Stage 2 (RLCF-F2F) | Stage 1 (SFT) | Stage 2 (RLCF-WB) | Stage 2 (RLCF-F2F) |\nSample Count ((n)) | 100 | 100 | 100 | 50 | 50 | 50 |\nWell-Formed Rate (%) | 96.0% | 99.0% | 99.0% | 100.0% | 98.0% | 98.0% |\nCompile@1 Rate (%) | 31.0% | 34.0% | 38.0% | 32.0% | 34.0% | 36.0% |\nCompile@2 Rate (%) | 39.0% | 53.0% | 53.0% | 44.0% | 52.0% | 54.0% |\nCompile@3 Rate (%) | 41.0% | 57.0% | 59.0% | 46.0% | 58.0% | 58.0% |\nCompile@4 Rate (%) | 42.0% | 59.0% | 60.0% | 48.0% | 58.0% | 62.0% |\nCompile@5 / Pass@5 (%) | 44.0% | 59.0% | 63.0% | 48.0% | 58.0% | 64.0% |\nMean Iterations Solved | 1.52 | 1.56 | 1.67 | 1.46 | 1.52 | 1.72 |\nStructural Faithfulness | 0.469 | 0.633 | 0.639 | 0.484 | 0.603 | 0.623 |\nGoal Exact Match (%) | 0.0% | 0.0% | 0.0% | 0.0% | 0.0% | 0.0% |\nGold Compiles Sanity (%) | 0.0% | 0.0% | 0.0% | 0.0% | 0.0% | 0.0% |\nThroughput (tok/s) | 6.5 | 7.0 | 6.8 | 3.0 | 3.0 | 3.2 |\n\n#### 7.2 Core Empirical Findings & Analysis\n\n- Effectiveness of Stage 2 RLCF: Reinforcement learning with compiler feedback increases statement compile@5 rates substantially over Stage 1 SFT baseline (QLoRA: 44.0% -> 63.0%; QDoRA: 48.0% -> 64.0%).\n- QLoRA vs. QDoRA Performance-Throughput Trade-off:\n- Accuracy: QDoRA achieves slightly higher overall compilation accuracy post-RLCF (64.0% vs. 63.0%) and higher initial SFT accuracy (48.0% vs. 44.0%).\n- Throughput: QLoRA achieves ~6.8–7.0 tokens/sec, whereas QDoRA achieves ~3.0–3.2 tokens/sec. The weight-decomposition calculation in DoRA adds execution overhead per forward pass on 4-bit MoE base weights.\n\n- Multi-Turn Repair Gains: Across all models, iterative compiler feedback improves overall pass rates significantly over single-turn generation (e.g.,\n`qwen-lora-prominif2f`\n\nimproves from compile@1 = 38.0% to compile@5 = 63.0%, a +25.0% net gain). - Error Distribution Dynamics: Unsolved problems transition from syntax/parse errors to missing identifier and typeclass resolution errors post-RLCF.\n\n#### 7.3 Tracking Model Evolution\n\nTo provide rigorous tracking as models evolve, subsequent research iterations should report the following 9 metrics:\n\n- Well-Formedness Rate (\n`well_formed`\n\n): Percentage of outputs parsing as valid Lean theorem signatures (balanced delimiters, formal structure). - Iterative Compile Rate (\n`compile@k`\n\n): Percentage of statements successfully type-checking against Mathlib within (k) feedback attempts ((k {1, 2, 3, 4, 5})). - Final Pass Rate (\n`pass@N`\n\n): Terminal compile success rate at maximum iteration cap (N). - Feedback Efficiency (\n`mean_iters_solved`\n\n): Average number of attempts required to solve successful problems. Lower values indicate higher quality initial attempts. - Structural Code Faithfulness (\n`faithfulness_code`\n\n): Sequence matching ratio between normalized generated statements and canonical references:\n\n- Semantic Goal Exact Match (\n`goal_exact_match`\n\n): Percentage of generated statement elaborated goals matching reference elaborated goal strings within the Lean kernel state. - Compiler Error Categorization (\n`error_breakdown`\n\n): Distribution of failure modes across unsolved instances (`unknown_ident`\n\n,`syntax`\n\n,`typeclass`\n\n,`type_mismatch`\n\n,`other`\n\n). - Generation Throughput (\n`tokens_per_sec`\n\n): Token generation speed during inference passes. - Hardware Resource Footprint: Peak VRAM allocation (GB) and GPU compute utilization (%).\n\n*Figure 3: Type-checking compile rate (**compile@k**) scaling across 5 feedback iterations. Stage 2 RLCF models consistently outperform Stage 1 SFT baselines, with multi-turn feedback adding 25–28% absolute compile accuracy.*\n\n*Figure 4: Generation throughput versus **compile@5** accuracy. QLoRA delivers ~2.2x higher generation speed with minimal accuracy drop compared to QDoRA.*\n\n*Figure 5: Shift in compiler failure modes across training stages. RLCF reduces missing identifier errors (*`unknown_ident`\n\n*) while syntax parsing errors dominate remaining unsolved instances.*\n\n#### 7.4 MoE Single-GPU Compute Bottlenecks\n\nWhile MoE models restrict active parameter computation to ~3B parameters per forward pass, running sparse MoE on a single GPU incurs fixed kernel launch overhead:\n\n- All 128 expert kernel branches execute sequentially or in small parallel launches per micro-batch.\n\n- Token packing (concatenating short sequences to length 1024) amortizes launch overhead, increasing throughput from <100 tok/s to **357 tok/s** during SFT training.\n\n#### 7.5 Sequence Packing vs. Completion Masking Trade-Off\n\n- Full Sequence Packing: Maximizes GPU compute efficiency but forces training on prompt tokens unless complex custom cross-entropy attention masking is implemented.\n- Completion-Only Masking: Prevents policy drift on prompt tokens but leaves padding overhead when sequences vary in length.\n\n**Conclusion**\n\nThe contribution is a purpose-built model-plus-agent architecture for Lean 4 autoformalization that runs on a single GPU. Sparse MoE inference (~3B active parameters), NF4 quantization, QLoRA adapters, and a persistent Lean REPL pool yield ~6.8–7.0 tok/s locally. That throughput supports multi-turn compiler repair and fast training/eval iteration without a cluster.\n\nSemantic exactness is not the strength of this stack (faithfulness ~0.63). Syntax and compilation are: **well-formedness 98–100%** on ProofNet, **compile@5 of 63–64%** after RLCF, and **100% well-formed / ~74% single-pass** compile on leak-free miniF2F. Systems A/B retrieval plus the agent loop convert modest compile@1 into those pass@5 rates. The result is a locally runnable system that is useful on undergraduate and competition-level statements, and cheap enough to iterate on.\n\n**References**\n\n[1] Qwen Team (2025). *Qwen3 Technical Report*. arXiv preprint arXiv:2505.09388. [https://arxiv.org/abs/2505.09388](https://arxiv.org/abs/2505.09388)\n\n[2] Tim Dettmers, Artidoro Pagnoni, Ari Holtzman, and Luke Zettlemoyer (2023). *QLoRA: Efficient Finetuning of Quantized LLMs*. arXiv preprint arXiv:2305.14314. [https://arxiv.org/abs/2305.14314](https://arxiv.org/abs/2305.14314)\n\n[3] Shih-Yang Liu, Chien-Yi Wang, Hongxu Yin, Pavlo Molchanov, Yu-Chiang Frank Wang, Kwang-Ting Cheng, and Min-Hung Chen (2024). *DoRA: Weight-Decomposed Low-Rank Adaptation*. arXiv preprint arXiv:2402.09353. [https://arxiv.org/abs/2402.09353](https://arxiv.org/abs/2402.09353)\n\n[4] Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Mingchuan Zhang, Y. K. Li, Y. Wu, and W. Liang (2024). *DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models*. arXiv preprint arXiv:2402.03300. [https://arxiv.org/abs/2402.03300](https://arxiv.org/abs/2402.03300)\n\n[5] Daya Guo, Dejian Yang, Haowei Zhang, Chaoyi Song, Ruoyu Zhang, Runxin Xu, et al. (2025). *DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning*. arXiv preprint arXiv:2501.12948. [https://arxiv.org/abs/2501.12948](https://arxiv.org/abs/2501.12948)\n\n[6] Guo Zheng, Stanislas Polu, Jesse Michael Han, Christian Szegedy, and Ilya Sutskever (2023). *ProofNet: Autoformalizing and Formally Proving Undergraduate-Level Mathematics Problems*. arXiv preprint arXiv:2302.12433. [https://arxiv.org/abs/2302.12433](https://arxiv.org/abs/2302.12433)\n\n[7] DeepSeek-AI (2024). *Lean-Workbook: A Large-Scale Dataset for Lean 4 Autoformalization*. Hugging Face Datasets. [https://huggingface.co/datasets/deepseek-ai/Lean-Workbook](https://huggingface.co/datasets/deepseek-ai/Lean-Workbook)\n\n[8] Alex J. Best (2024). *Herald: Natural Language to Lean 4 Autoformalization Dataset*. Hugging Face Datasets. [https://huggingface.co/datasets/alexjbest/herald](https://huggingface.co/datasets/alexjbest/herald)\n\n[9] Facebook Research (2021). *miniF2F: A Cross-System Benchmark for Formal Olympiad Mathematics*. Hugging Face Datasets. [https://huggingface.co/datasets/facebook/miniF2F](https://huggingface.co/datasets/facebook/miniF2F)\n\n[10] Hoskinson Center for Formal Mathematics (2023). *ProofNet: Undergraduate Mathematics Autoformalization Dataset*. Hugging Face Datasets. [https://huggingface.co/datasets/hoskinson-center/proofnet](https://huggingface.co/datasets/hoskinson-center/proofnet)", "url": "https://wpnews.pro/news/high-throughput-lean-4-autoformalization-model-for-local-inference", "canonical_source": "https://meshapplied.com/posts/lean4-autoformalization", "published_at": "2026-08-22 12:26:56+00:00", "updated_at": "2026-08-22 12:43:20.251347+00:00", "lang": "en", "topics": ["artificial-intelligence", "machine-learning", "large-language-models", "ai-research", "ai-infrastructure"], "entities": ["Qwen3-Coder-30B-A3B", "Lean 4", "Nvidia RTX 5090ti", "Group Relative Policy Optimization", "Mixture-of-Experts"], "alternates": {"html": "https://wpnews.pro/news/high-throughput-lean-4-autoformalization-model-for-local-inference", "markdown": "https://wpnews.pro/news/high-throughput-lean-4-autoformalization-model-for-local-inference.md", "text": "https://wpnews.pro/news/high-throughput-lean-4-autoformalization-model-for-local-inference.txt", "jsonld": "https://wpnews.pro/news/high-throughput-lean-4-autoformalization-model-for-local-inference.jsonld"}}