| --- |
| license: apache-2.0 |
| base_model: |
| - Qwen/Qwen3-4B |
| pipeline_tag: text-generation |
| library_name: transformers |
| datasets: |
| - Pythagoras-LM/SFT_Dataset |
| --- |
| <div align="center"> |
| <img src="https://github.com/Pythagoras-LM/Pythagoras-LM.github.io/blob/main/figures/pythagoras-without-background.png?raw=true" width="240" alt="Pythagoras-Prover logo"><br> |
| <h1>Pythagoras-Prover</h1> |
| </div> |
|
|
| <div align="center"> |
| |
| [](https://pythagoras-lm.github.io/) |
| [](https://arxiv.org/abs/XXXX.XXXXX) |
| [](https://github.com/Pythagoras-LM/Pythagoras-Prover) |
| [](https://opensource.org/licenses/Apache-2.0) |
|
|
| </div> |
|
|
|
|
| ## 1. Introduction |
|
|
| We introduce **Pythagoras-Prover**, a compute-efficient family of open-source large language models for formal theorem proving in Lean 4. The family comprises two autoregressive provers at 4B and 32B parameters, together with **Pythagoras-Prover-Diffusion**, the first diffusion-based theorem prover, which iteratively refines Lean proofs at inference time. All three models are artefacts of a single methodological approach: a scalable, Lean-verified synthetic data pipeline. At its centre is **Augmented Lean Formalisation** (ALF), a structured mutation scheme that expands a verified seed corpus into formal variants without per-instance Lean compilation, then re-uses them as a self-distillation signal during training. This design lets careful data construction stand in for raw scale, closing much of the gap between small open provers and their largest counterparts — without relying on inference-time self-correction. |
|
|
| <div align="center"> |
| <img src="https://github.com/Pythagoras-LM/Pythagoras-LM.github.io/blob/main/figures/prover_fig1_v2_hi.png?raw=true" alt="Pythagoras-Prover benchmark overview" width="90%"> |
| </div> |
|
|
|
|
| ## 2. Model Summary |
|
|
| --- |
|
|
| **A Lean-Verified Synthetic Data Pipeline** |
|
|
| - Natural-language problems from general-math and competition sources are autoformalised into Lean and gated on the type-checker using predominantly sub-30B open models, with an auto-informalisation and alignment step discarding faithful-but-wrong formalisations to yield a verified seed corpus partitioned into easy, medium, and hard tiers. |
| - A rubric-guided distillation stage re-prompts on each rejected instance to target the specific Lean type-checker error responsible for its failure, lifting autoformalisation success and roughly doubling the verified training set. |
|
|
| --- |
|
|
| **Model Training** |
|
|
| - LoRA-only supervised fine-tuning of Qwen3-4B and Qwen3-32B under an 8K context, paired with a dynamic proof-reasoning filter and a difficulty-ordered easy→medium→hard curriculum, followed by reinforcement learning with a Lean-compilation reward and a final continued-SFT stage on the ALF corpus. |
|
|
| --- |
|
|
| **Augmented Lean Formalisation** |
| - ALF emits one structured variant per category — simplification, generalisation, lemma proposal, proof-step decomposition, and reformulation — for every seed statement, replacing per-instance Lean verification with a cheap statement-alignment check and expanding the seed corpus into roughly 2M formal variants. |
| - The post-RL prover proves the mutations, and these self-distilled proofs form a corpus that trains both the autoregressive and diffusion provers from a single recipe. |
|
|
| --- |
|
|
| **The Smallest Efficient Open-Source Lean Theorem Prover** |
|
|
| - We train Pythagoras-Prover-4B, one of the smallest and most compute-efficient open-source Lean theorem provers to date, reaching 86.07% on MiniF2F-Test at Pass@32 and surpassing the prior state of the art at a fraction of its parameter count. |
| - Pythagoras-Prover-Diffusion adapts a block-diffusion formulation with a tactic-based masking objective aligned to the discrete reasoning steps of Lean — to our knowledge the first demonstration that a diffusion language model can verifiably solve Lean theorems at non-trivial rates. |
|
|
| --- |
|
|
| The resulting models set a new bar for compute-efficient formal proving. **Pythagoras-Prover-32B** achieves state-of-the-art performance among open-source provers, reaching **93.03%** on MiniF2F-Test and solving **93 of 672** problems on PutnamBench, while **Pythagoras-Prover-4B** outperforms DeepSeek-Prover-V2-671B on MiniF2F-Test despite being roughly **167× smaller** — with no self-correction or test-time reinforcement learning. We additionally release **MiniF2F-ALF**, an ALF-mutated companion benchmark on which every evaluated prover degrades. |
|
|
| ## 3. Benchmark Performance |
|
|
| We evaluate Pythagoras-Prover on three Lean 4 benchmarks — MiniF2F-Test, PutnamBench, and the MiniF2F-ALF benchmark we introduce — under a single unified protocol (Lean 4.9.0-rc1, a 30,000-token generation limit, and a verbatim-statement pass criterion). Across all three, Pythagoras-Prover matches or exceeds open-source provers an order of magnitude larger, and does so **without** inference-time self-correction or test-time reinforcement learning. |
|
|
|
|
| <div align="center"> |
| <table style="margin: 0 auto;"> |
| <thead> |
| <tr> |
| <th>Method</th> |
| <th>#Params</th> |
| <th>Pass@32</th> |
| <th>Pass@1024</th> |
| <th>Best (N)</th> |
| </tr> |
| </thead> |
| <tbody> |
| <tr><td>Goedel-Prover-SFT</td><td>7B</td><td>57.6</td><td>–</td><td>62.7 (3200)</td></tr> |
| <tr><td>STP</td><td>7B</td><td>–</td><td>–</td><td>67.6 (25600)</td></tr> |
| <tr><td>Kimina-Prover-Preview-72B</td><td>72B</td><td>68.85</td><td>–</td><td>80.74 (8192)</td></tr> |
| <tr><td>DeepSeek-Prover-V2-7B</td><td>7B</td><td>75.6</td><td>–</td><td>82.0 (8192)</td></tr> |
| <tr><td>DeepSeek-Prover-V2-671B</td><td>671B</td><td>82.4</td><td>–</td><td>88.9 (8192)</td></tr> |
| <tr><td>Kimina-Prover-8B-Distill</td><td>8B</td><td>77.86</td><td>–</td><td>–</td></tr> |
| <tr><td>Kimina-Prover-70B</td><td>70B</td><td>84.0</td><td>87.7</td><td>92.2 (TTRL)</td></tr> |
| <tr><td>Goedel-Prover-V2-8B</td><td>8B</td><td>84.6</td><td>87.9</td><td>90.2 (8192)</td></tr> |
| <tr><td> + Self-Correction</td><td>8B</td><td>86.7</td><td>89.3</td><td>–</td></tr> |
| <tr><td>Goedel-Prover-V2-32B</td><td>32B</td><td>88.1</td><td>91.8</td><td>92.2 (8192)</td></tr> |
| <tr><td> + Self-Correction</td><td>32B</td><td>90.4</td><td>92.6</td><td>–</td></tr> |
| <tr><td><strong>Pythagoras-Prover-4B</strong></td><td>4B</td><td><strong>86.07</strong></td><td><strong>88.11</strong></td><td><strong>89.75 (2048)</strong></td></tr> |
| <tr><td><strong>Pythagoras-Prover-32B</strong></td><td>32B</td><td><strong>89.75</strong></td><td><strong>92.62</strong></td><td><strong>93.03 (2048)</strong></td></tr> |
| </tbody> |
| </table> |
| <!-- table caption --> |
| <caption align="bottom"><strong>Table 1</strong>: <em>Pythagoras-Prover-4B exceeds DeepSeek-Prover-V2-671B's pass@8192 result (88.9%) at pass@2048 — a quarter of the budget and ~167× fewer parameters. Pythagoras-Prover-32B sets the strongest reported MiniF2F-Test pass rate without self-correction or test-time RL.</em></caption> |
| </div> |
| |
|
|
| <br> |
|
|
| <div align="center"> |
| <table style="margin: 0 auto;"> |
| <thead> |
| <tr> |
| <th>#</th> |
| <th>Model</th> |
| <th>num-solved</th> |
| <th>compute</th> |
| </tr> |
| </thead> |
| <tbody> |
| <tr><td>1</td><td><strong>Pythagoras-Prover-32B</strong></td><td><strong>93</strong></td><td><strong>Pass@2048</strong></td></tr> |
| <tr><td>1</td><td><strong>Pythagoras-Prover-32B</strong></td><td><strong>59</strong></td><td><strong>Pass@64</strong></td></tr> |
| <tr><td>1</td><td><strong>Pythagoras-Prover-32B</strong></td><td><strong>48</strong></td><td><strong>Pass@32</strong></td></tr> |
| <tr><td>2</td><td>Goedel-Prover-V2-32B (self-correction mode)</td><td>86</td><td>Pass@184</td></tr> |
| <tr><td>2</td><td>Goedel-Prover-V2-32B (self-correction mode)</td><td>57</td><td>Pass@32</td></tr> |
| <tr><td>2</td><td>Goedel-Prover-V2-32B</td><td>43</td><td>Pass@32</td></tr> |
| <tr><td>3</td><td>DeepSeek-Prover-V2-671B</td><td>47</td><td>Pass@1024</td></tr> |
| <tr><td>3</td><td>DeepSeek-Prover-V2-671B</td><td>22</td><td>Pass@32</td></tr> |
| <tr><td>4</td><td>DSP+</td><td>23</td><td>Pass@128</td></tr> |
| <tr><td>5</td><td>Bourbaki</td><td>14</td><td>Pass@512</td></tr> |
| <tr><td>6</td><td>Kimina-Prover-7B-Distill</td><td>10</td><td>Pass@192</td></tr> |
| <tr><td>7</td><td>Self-play Theorem Prover</td><td>8</td><td>Pass@3200</td></tr> |
| <tr><td>8</td><td>Goedel-Prover-SFT</td><td>7</td><td>Pass@512</td></tr> |
| <tr><td>9</td><td>ABEL (closed-source)</td><td>7</td><td>Pass@596</td></tr> |
| </tbody> |
| </table> |
| <!-- table caption --> |
| <caption align="bottom"><strong>Table 2</strong>: <em>PutnamBench leaderboard (problems solved out of 657). Pythagoras-Prover-32B takes the top rank, solving 93 problems at Pass@2048 — 7 more than the previous best (Goedel-Prover-V2-32B, 86 at Pass@184 in self-correction mode) and nearly double DeepSeek-Prover-V2-671B's 47 at Pass@1024, despite being roughly 20× smaller. Seed-Prover (331 solved) is omitted from the ranked rows as it is closed-source with undisclosed test-time compute.</em></caption> |
| </div> |
| |
| <br> |
|
|
| <div align="center"> |
| <table style="margin: 0 auto;"> |
| <thead> |
| <tr> |
| <th>Model</th> |
| <th>Pass@32</th> |
| </tr> |
| </thead> |
| <tbody> |
| <tr><td>DeepSeek-Prover-V2-671B</td><td>79.71</td></tr> |
| <tr><td>Goedel-Prover-V2-8B</td><td>82.58</td></tr> |
| <tr><td>Goedel-Prover-V2-32B</td><td>83.61</td></tr> |
| <tr><td><strong>Pythagoras-Prover-4B</strong></td><td><strong>83.19</strong></td></tr> |
| <tr><td><strong>Pythagoras-Prover-32B</strong></td><td><strong>85.04</strong></td></tr> |
| </tbody> |
| </table> |
| <!-- table caption --> |
| <caption align="bottom"><strong>Table 3</strong>: <em>Performance of current state-of-the-art provers on MiniF2F-ALF (Pass@32, %). As MiniF2F-ALF is introduced in this work, all results are evaluated by us under a unified setup.</em></caption> |
| </div> |
| |
| <br> |
|
|
| <div align="center"> |
| <table style="margin: 0 auto;"> |
| <thead> |
| <tr> |
| <th>Model</th> |
| <th>Performance</th> |
| </tr> |
| </thead> |
| <tbody> |
| <tr><td>Pythagoras-Prover</td><td>86.1%</td></tr> |
| <tr><td>Pythagoras-Prover∗</td><td>74.6%</td></tr> |
| <tr><td><strong>Diffusion Pythagoras-Prover</strong></td><td><strong>63.3%</strong></td></tr> |
| </tbody> |
| </table> |
| <caption align="bottom"><strong>Table 4</strong>: <em>Diffusion Theorem Proving Performance (Pass@32). ∗ denotes the setting where training tokens are restricted to 4096 and evaluation is performed solely at 8192 tokens.</em></caption> |
| </div> |
| |
| ## 4. Model & Dataset Downloads |
|
|
| We release our Pythagoras-Prover models, the training dataset and the new MiniF2F-ALF benchmark for future future research. |
|
|
| <div align="center"> |
| |
| | Model | Download | |
| | -------- | -------- | |
| | Pythagoras-Prover-32B | Coming Soon | |
| | Pythagoras-Prover-4B | [🤗Download](https://huggingface.co/Pythagoras-LM/Pythagoras-Prover-4B) | |
| | Pythagoras-Prover-Diffusion-4B | [🤗Download](https://huggingface.co/Pythagoras-LM/Pythagoras-Prover-Diffusion-4B) | |
|
|
| </div> |
|
|
|
|
| <div align="center"> |
|
|
| | Dataset | Download | |
| | -------- | -------- | |
| | Pythagoras-Prover-SFT | [🤗Download](https://huggingface.co/datasets/Pythagoras-LM/SFT_Dataset) | |
| | Pythagoras-Prover-Distill-4B | Coming Soon | |
| | Pythagoras-Prover-Distill-32B | Coming Soon | |
| </div> |
|
|
| ## 5. Quick Start |
|
|
| This checkpoint is an A2D/BD3LM diffusion prover. It should be decoded with a diffusion sampler rather than `AutoModelForCausalLM.generate`. The snippet below is self-contained: it registers the A2D Qwen3 model class, implements the BD3LM block-diffusion sampler used for MiniF2F evaluation, loads the checkpoint with Transformers, and runs one Lean theorem-proving prompt. The decoding settings match the MiniF2F evaluation setting used for the reported Pass@32 result: bf16, SDPA attention, 8192 generated tokens, 8192 diffusion steps, and block size 8. |
|
|
| ````python |
| import copy |
| import math |
| import torch |
| import torch.nn.functional as F |
| import transformers |
| from transformers import AutoConfig, AutoModel, AutoModelForMaskedLM, AutoTokenizer |
| from transformers.cache_utils import DynamicCache |
| from transformers.modeling_attn_mask_utils import _prepare_4d_attention_mask |
| from transformers.modeling_outputs import BaseModelOutputWithPast |
| |
| |
| class A2DQwen3Config(transformers.Qwen3Config): |
| model_type = "a2d-qwen3" |
| |
| |
| class A2DQwen3Model(transformers.Qwen3Model): |
| def forward( |
| self, |
| input_ids=None, |
| attention_mask=None, |
| position_ids=None, |
| past_key_values=None, |
| inputs_embeds=None, |
| use_cache=None, |
| cache_position=None, |
| **kwargs, |
| ): |
| if (input_ids is None) == (inputs_embeds is None): |
| raise ValueError("You must specify exactly one of input_ids or inputs_embeds") |
| |
| if inputs_embeds is None: |
| inputs_embeds = self.embed_tokens(input_ids) |
| |
| if use_cache and past_key_values is None: |
| past_key_values = DynamicCache(config=self.config) |
| |
| if cache_position is None: |
| past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 |
| cache_position = torch.arange( |
| past_seen_tokens, |
| past_seen_tokens + inputs_embeds.shape[1], |
| device=inputs_embeds.device, |
| ) |
| |
| if position_ids is None: |
| position_ids = cache_position.unsqueeze(0) |
| |
| if not isinstance(causal_mask_mapping := attention_mask, dict): |
| if attention_mask is None: |
| attention_mask = torch.ones( |
| inputs_embeds.shape[:2], |
| device=inputs_embeds.device, |
| dtype=torch.long, |
| ) |
| |
| if not (isinstance(attention_mask, torch.Tensor) and attention_mask.ndim == 4): |
| attention_mask = _prepare_4d_attention_mask(attention_mask, self.dtype) |
| |
| causal_mask_mapping = {"full_attention": attention_mask} |
| if self.has_sliding_layers: |
| causal_mask_mapping["sliding_attention"] = attention_mask |
| |
| hidden_states = inputs_embeds |
| position_embeddings = self.rotary_emb(hidden_states, position_ids) |
| |
| for decoder_layer in self.layers[: self.config.num_hidden_layers]: |
| hidden_states = decoder_layer( |
| hidden_states, |
| attention_mask=causal_mask_mapping[decoder_layer.attention_type], |
| position_ids=position_ids, |
| past_key_values=past_key_values, |
| use_cache=use_cache, |
| cache_position=cache_position, |
| position_embeddings=position_embeddings, |
| **kwargs, |
| ) |
| |
| hidden_states = self.norm(hidden_states) |
| return BaseModelOutputWithPast( |
| last_hidden_state=hidden_states, |
| past_key_values=past_key_values if use_cache else None, |
| ) |
| |
| |
| class A2DQwen3LMHeadModel(transformers.Qwen3ForCausalLM): |
| config_class = A2DQwen3Config |
| |
| def __init__(self, config): |
| transformers.Qwen3PreTrainedModel.__init__(self, config) |
| self.model = A2DQwen3Model(config) |
| self.vocab_size = config.vocab_size |
| self.lm_head = torch.nn.Linear(config.hidden_size, config.vocab_size, bias=False) |
| self.post_init() |
| |
| |
| AutoConfig.register("a2d-qwen3", A2DQwen3Config) |
| AutoModel.register(A2DQwen3Config, A2DQwen3LMHeadModel) |
| AutoModelForMaskedLM.register(A2DQwen3Config, A2DQwen3LMHeadModel) |
| |
| |
| def prepare_for_sampling(x, block_size, pad_token_id): |
| batch_size, seq_len = x.shape |
| device = x.device |
| |
| valid = x != pad_token_id |
| pos_raw = torch.cumsum(valid.to(torch.long), dim=-1) |
| logical_pos = pos_raw - 1 |
| position_ids = torch.where(valid, logical_pos, torch.zeros_like(logical_pos)).to( |
| device=device, |
| dtype=torch.long, |
| ) |
| |
| pos = torch.arange(seq_len, device=device) |
| block_ids = torch.div(pos, block_size, rounding_mode="floor") |
| block_ids = block_ids.view(1, seq_len).expand(batch_size, -1) |
| block_ids = torch.where(valid, block_ids, torch.full_like(block_ids, -1)) |
| |
| bid_q = block_ids.view(batch_size, 1, seq_len, 1) |
| bid_k = block_ids.view(batch_size, 1, 1, seq_len) |
| attention_mask = (bid_k <= bid_q) & (bid_q >= 0) & (bid_k >= 0) |
| return attention_mask, position_ids |
| |
| |
| def add_gumbel_noise(logits, temperature): |
| if temperature == 0: |
| return logits |
| logits = logits.to(torch.float64) |
| noise = torch.rand_like(logits, dtype=torch.float64) |
| return logits.exp() / ((-torch.log(noise)) ** temperature) |
| |
| |
| def get_num_transfer_tokens(mask_index, steps): |
| mask_num = mask_index.sum(dim=1, keepdim=True) |
| num_transfer_tokens = torch.zeros( |
| mask_num.size(0), |
| steps, |
| device=mask_index.device, |
| dtype=torch.int64, |
| ) |
| for i in range(mask_num.size(0)): |
| for t, s, j in zip(range(steps, 0, -1), range(steps - 1, -1, -1), range(steps)): |
| reverse_transfer_prob = 1 - (s / t) |
| transfer = torch.round(mask_num[i, 0].to(torch.float64) * reverse_transfer_prob).to(torch.int64) |
| transfer = torch.minimum(transfer, mask_num[i, 0]) |
| num_transfer_tokens[i, j] = transfer |
| mask_num[i, 0] -= transfer |
| if mask_num[i, 0].item() == 0: |
| break |
| |
| rows = [] |
| max_len = 0 |
| for row in num_transfer_tokens: |
| nonzero = row[row > 0] |
| rows.append(nonzero) |
| max_len = max(max_len, nonzero.numel()) |
| |
| padded_rows = [] |
| for row in rows: |
| if row.numel() < max_len: |
| row = torch.cat([row, torch.zeros(max_len - row.numel(), dtype=row.dtype, device=row.device)]) |
| padded_rows.append(row) |
| return torch.stack(padded_rows, dim=0) |
| |
| |
| def diffusion_step_block(logits, x_block, mask_block, num_transfer_step, temperature): |
| batch_size, block_len, _ = logits.shape |
| logits_with_noise = add_gumbel_noise(logits, temperature=temperature) |
| x0 = torch.argmax(logits_with_noise, dim=-1) |
| |
| probs = F.softmax(logits, dim=-1) |
| x0_probs = torch.gather(probs, dim=-1, index=x0.unsqueeze(-1)).squeeze(-1) |
| |
| x0 = torch.where(mask_block, x0, x_block) |
| neg_inf = torch.full_like(x0_probs, -float("inf")) |
| confidence = torch.where(mask_block, x0_probs, neg_inf) |
| |
| transfer = torch.zeros_like(x0, dtype=torch.bool) |
| for batch_idx in range(batch_size): |
| k = int(num_transfer_step[batch_idx].item()) |
| if k <= 0: |
| continue |
| valid_count = (confidence[batch_idx] > -float("inf")).sum().item() |
| if valid_count == 0: |
| continue |
| _, selected = torch.topk(confidence[batch_idx], min(k, valid_count)) |
| transfer[batch_idx, selected] = True |
| |
| x_block_new = x_block.clone() |
| x_block_new[transfer] = x0[transfer] |
| return x_block_new |
| |
| |
| @torch.no_grad() |
| def generate_bd3lm( |
| model, |
| tokenizer, |
| input_ids, |
| max_new_tokens=8192, |
| steps=8192, |
| block_size=8, |
| temperature=1.0, |
| ): |
| if isinstance(input_ids, torch.Tensor): |
| input_ids = input_ids.tolist() |
| if input_ids and isinstance(input_ids[0], int): |
| input_ids = [input_ids] |
| |
| device = model.device |
| mask_id = tokenizer.mask_token_id |
| pad_id = tokenizer.pad_token_id |
| eos_id = tokenizer.eos_token_id |
| |
| inputs = [torch.tensor(ids, dtype=torch.long, device=device) for ids in input_ids] |
| prompt_lens = [ids.shape[0] for ids in inputs] |
| max_prompt_len = max(prompt_lens) |
| padded_prompt_len = math.ceil(max_prompt_len / block_size) * block_size |
| |
| x = torch.full( |
| (len(inputs), padded_prompt_len), |
| pad_id, |
| dtype=torch.long, |
| device=device, |
| ) |
| for batch_idx, ids in enumerate(inputs): |
| x[batch_idx, padded_prompt_len - ids.shape[0] : padded_prompt_len] = ids |
| |
| num_blocks = math.ceil(max_new_tokens / block_size) |
| steps_per_block = math.ceil(steps / num_blocks) |
| done = torch.zeros((len(inputs),), dtype=torch.bool, device=device) |
| generated = 0 |
| |
| for _ in range(num_blocks): |
| if done.all(): |
| break |
| |
| prefix_len = x.shape[1] |
| cur_block_len = min(block_size, max_new_tokens - generated) |
| if cur_block_len <= 0: |
| break |
| |
| prefix_attention_mask, prefix_position_ids = prepare_for_sampling( |
| x, |
| block_size=block_size, |
| pad_token_id=pad_id, |
| ) |
| prefix_outputs = model( |
| x, |
| attention_mask=prefix_attention_mask, |
| position_ids=prefix_position_ids, |
| use_cache=True, |
| ) |
| prefix_cache = prefix_outputs.past_key_values |
| |
| new_block = torch.full( |
| (x.shape[0], cur_block_len), |
| mask_id, |
| dtype=torch.long, |
| device=device, |
| ) |
| x = torch.cat([x, new_block], dim=1) |
| total_len = x.shape[1] |
| |
| block_mask_index = x[:, -cur_block_len:] == mask_id |
| num_transfer_tokens = get_num_transfer_tokens(block_mask_index, steps_per_block) |
| |
| full_attention_mask, full_position_ids = prepare_for_sampling( |
| x, |
| block_size=block_size, |
| pad_token_id=pad_id, |
| ) |
| block_attention_mask = full_attention_mask[:, :, prefix_len:total_len, :] |
| block_position_ids = full_position_ids[:, prefix_len:total_len] |
| |
| for step_idx in range(num_transfer_tokens.size(1)): |
| x_block = x[:, prefix_len:total_len] |
| mask_block = x_block == mask_id |
| if not mask_block.any(): |
| break |
| |
| logits = model( |
| x_block, |
| attention_mask=block_attention_mask, |
| position_ids=block_position_ids, |
| past_key_values=copy.deepcopy(prefix_cache), |
| use_cache=False, |
| ).logits |
| |
| x[:, prefix_len:total_len] = diffusion_step_block( |
| logits=logits, |
| x_block=x_block, |
| mask_block=mask_block, |
| num_transfer_step=num_transfer_tokens[:, step_idx], |
| temperature=temperature, |
| ) |
| |
| if eos_id is not None: |
| done = done | (x[:, prefix_len:total_len] == eos_id).any(dim=1) |
| generated += cur_block_len |
| |
| return x |
| |
| |
| def trim_completion(tokenizer, sequence_ids, prompt_ids): |
| sequence_ids = list(sequence_ids) |
| while sequence_ids and sequence_ids[0] == tokenizer.pad_token_id: |
| sequence_ids.pop(0) |
| |
| start = len(prompt_ids) |
| end = len(sequence_ids) |
| stop_ids = {tokenizer.eos_token_id} |
| if getattr(tokenizer, "eot_token_id", None) is not None: |
| stop_ids.add(tokenizer.eot_token_id) |
| |
| for idx in range(start, len(sequence_ids)): |
| if sequence_ids[idx] in stop_ids: |
| end = idx |
| break |
| |
| return tokenizer.decode(sequence_ids[start:end], skip_special_tokens=True) |
| |
| model_id = "Pythagoras-LM/Pythagoras-Prover-Diffusion-4B" |
| |
| formal_statement = """ |
| import Mathlib |
| import Aesop |
| |
| set_option maxHeartbeats 0 |
| |
| open BigOperators Real Nat Topology Rat |
| |
| /-- The volume of a cone is given by the formula $V = \frac{1}{3}Bh$, where $B$ is the area of the base and $h$ is the height. The area of the base of a cone is 30 square units, and its height is 6.5 units. What is the number of cubic units in its volume? Show that it is 65.-/ |
| theorem mathd_algebra_478 (b h v : ℝ) (h₀ : 0 < b ∧ 0 < h ∧ 0 < v) (h₁ : v = 1 / 3 * (b * h)) |
| (h₂ : b = 30) (h₃ : h = 13 / 2) : v = 65 := by |
| sorry |
| """.strip() |
| |
| prompt = """ |
| Complete the following Lean 4 code: |
| |
| ```lean4 |
| {}``` |
| |
| Before producing the Lean 4 code to formally prove the given theorem, provide a detailed proof plan outlining the main proof steps and strategies. |
| The plan should highlight key ideas, intermediate lemmas, and proof structures that will guide the construction of the final formal proof. |
| """.strip() |
| |
| if not torch.cuda.is_available(): |
| raise RuntimeError("CUDA is required for practical inference with this 4B diffusion checkpoint.") |
| |
| tokenizer = AutoTokenizer.from_pretrained(model_id) |
| if tokenizer.eos_token != "<|im_end|>": |
| tokenizer.eos_token = "<|im_end|>" |
| if tokenizer.pad_token_id is None: |
| tokenizer.pad_token = tokenizer.eos_token or tokenizer.bos_token |
| |
| model = AutoModelForMaskedLM.from_pretrained( |
| model_id, |
| dtype=torch.bfloat16, |
| device_map={"": 0}, |
| attn_implementation="sdpa", |
| ).eval() |
| |
| chat = [{"role": "user", "content": prompt.format(formal_statement)}] |
| input_ids = tokenizer.apply_chat_template( |
| chat, |
| tokenize=True, |
| add_generation_prompt=True, |
| enable_thinking=False, |
| ) |
| |
| outputs = generate_bd3lm( |
| model, |
| tokenizer, |
| input_ids, |
| max_new_tokens=8192, |
| steps=8192, |
| block_size=8, |
| temperature=1.0, |
| ) |
| |
| completion = trim_completion(tokenizer, outputs[0].tolist(), input_ids) |
| print(completion) |
| ```` |
|
|
| # Cite |
| ``` |
| @article{leang2026pythagoras, |
| title={Pythagoras-Prover: Advancing Efficient Formal Proving via Augmented Lean Formalisation}, |
| author={Leang, Joshua Ong Jun and Zhao, Zheng and Stoian, Mihaela Catalina and Xu, Qiyuan and Li, Haonan and Li, Wenda and Cohen, Shay B. and Giunchiglia, Eleonora}, |
| journal={arXiv preprint arXiv:2606.12594}, |
| year={2026} |
| } |
| ``` |
|
|
