PaperlayerView exact revision on arXiv

Source-linked AI summary

DiffPDE: Masked Diffusion Language Models as PDE Solver

Wenxuan Guo, Yuyang Hong, Lubin Fan, Zhaojin Fu, Lin Chen, Kun Ding, Shiming Xiang

arXiv:2608.30532v1cs.AI

TL;DR

Autoregressive PDE-solver repair redundantly regenerates globally when bugs are localized, motivating a more targeted approach. DiffPDE uses discrete diffusion models for localized re-masking and infilling, and ID-GRPO supports multi-round debugging of coupled bugs. Across PDEBench’s five PDE families, it remains competitive with larger AR baselines, improves over same-scale AR models, and repairs much faster.

  • Problem

    Autoregressive solver repair regenerates full contexts sequentially despite localized bugs, while coupled failures require multiple repairs and produce sparse feedback.

  • Method

    DiffPDE uses discrete diffusion models for localized re-masking and infilling, combined with ID-GRPO for multi-round repair using intermediate rewards.

  • Results

    Across five PDEBench families, DiffPDE remains competitive with larger AR baselines, substantially improves over same-scale AR models, and delivers much faster localized repair.

  • Takeaways & Limitations

    Localized infilling supports high-quality solver synthesis and more efficient, robust debugging than same-scale autoregressive baselines.

  • Takeaways & Limitations

    The framework focuses on localized debugging; failures requiring global structural revisions remain beyond its scope.

Abstract

from arXiv · show

Existing approaches for synthesizing Partial Differential Equation (PDE) solvers predominantly rely on autoregressive models, yet their global left-to-right decoding incurs substantial redundancy when addressing inherently localized bugs. In this work, we challenge this inefficient paradigm and propose DiffPDE, a framework leveraging discrete diffusion language models for targeted code repair. By introducing a localized re-masking and infilling strategy, DiffPDE regenerates only erroneous regions while preserving correct context, naturally aligning generation with the sparse nature of PDE errors. Furthermore, to handle coupled bugs requiring sequential interventions, we present Iterative Debugging GRPO (ID-GRPO), a reinforcement learning scheme that enables multi-round debugging within single trajectories via intermediate rewards. Experiments on PDEBench show that DiffPDE achieves competitive accuracy, outperforms same-scale AR models, and significantly accelerates repair.

1 Introduction

PDE solver synthesis requires iterative debugging, but autoregressive repair regenerates globally even when errors are localized. DiffPDE uses localized diffusion infilling, while ID-GRPO supports multi-round repair for coupled bugs.

  • PDE solver implementation demands expertise in discretization, boundary conditions, and stability, while reliable synthesis also requires execution-driven iterative debugging.
  • AR repair sequentially regenerates code after localized errors, causing redundant computation over surrounding regions that are already correct.
  • DiffPDE re-masks only suspected erroneous regions and infills them while preserving surrounding correct code as conditional context.
  • ID-GRPO addresses coupled bugs and sparse feedback through multi-round debugging trajectories, task allocation, dynamic buffering, and intermediate rewards.
  • Experiments on benchmark PDE tasks show competitive accuracy, stronger robustness than same-scale AR baselines, and lower repair latency.

2 Related Work

PDE solving spans classical numerical methods, neural surrogates, and LLM-based code generation. Diffusion models are positioned as an alternative to autoregressive repair because they enable parallel refinement instead of sequential full-context processing.

  • Classical finite-difference, finite-element, and spectral solvers provide rigorous convergence guarantees but require substantial domain expertise for correct implementation.
  • Neural surrogates offer a data-driven alternative to classical PDE solvers, while LLMs increasingly support scientific reasoning and code generation.
  • Agentic scientific-code frameworks use iterative refinement and tree search to improve the robustness of synthesized programs.
  • Autoregressive LLMs incur latency by sequentially processing full error contexts, whereas diffusion models enable efficient parallel refinement for PDE solving.

3 Preliminaries

The paper formulates numerical PDE solving as code generation for differential equations with source and boundary or initial conditions. Continuous domains are discretized so numerical solvers produce spatiotemporal reference tensors, while diffusion LLMs progressively mask and recover code tokens.

  • 3.1 PDE Task Formulation: Numerical PDE solving is formulated as a code-generation task whose output is a numerical solver.
  • 3.1 PDE Task Formulation: A PDE uses differential operator L and boundary or initial-condition operator B to relate source term ϕ and constraint term ψ.
  • 3.1 PDE Task Formulation: Because closed-form solutions are usually unavailable, the domain is discretized into spatial grids and time steps, producing spatiotemporal tensors as reference solutions.
  • 3.2 Diffusion LLM Generation: A masked diffusion language model progressively corrupts a code sequence by replacing tokens with the special [MASK] token.
  • 3.2 Diffusion LLM Generation: Its reverse process predicts lower-noise sequences conditioned on the current noisy sequence xt.

4 Method

DiffPDE unifies from-scratch PDE solver synthesis with localized code repair using diffusion-based generation, while ID-GRPO trains multi-round debugging from execution feedback.

  • Code Generation and Localized Infilling: DiffPDE trains a diffusion language model for both complete solver synthesis and localized span infilling under execution feedback.SFT covers both from-scratch generation and localized repair, while RL further improves these capabilities.
  • Code Generation and Localized Infilling: At each reverse step, the dLLM predicts masked positions in parallel, with blocks generated sequentially and diffusion refinement performed within each block.This semi-autoregressive strategy balances generation quality and efficiency for long code sequences.
  • Code Generation and Localized Infilling: Localized infilling masks a contiguous erroneous span while conditioning on both left and right context, preserving surrounding PDE solver code.Bidirectional conditioning supports repairs involving tightly coupled numerical variables, tensor shapes, and iterative updates.
  • Supervised Fine-Tuning: The unified masked-diffusion objective uses high noise for from-scratch generation and low noise for localized infilling.Sampling the noise level uniformly trains recovery across varying corruption levels.
  • Iterative Debugging GRPO: ID-GRPO progressively shifts reinforcement learning from simple infilling toward error-driven debugging through task allocation and a dynamic buffer.Successful solvers are repurposed into synthetic infilling tasks, creating constructive and corrective challenges with adjustable sampling ratios.
  • Iterative Debugging GRPO: Multi-round rewards combine the best outcome across repair attempts with accumulated intermediate progress, providing learning signals even when all attempts fail.Intermediate rewards credit strict reductions in normalized error signatures and apply only to error-driven infilling tasks.

5 Experiments

Experiments on PDEBench evaluate DiffPDE across five PDE families using accuracy, executability, efficiency, ablations, and repair-behavior analyses. DiffPDE improves same-scale solver performance and reliability while localized infilling substantially reduces repair latency, though elliptic Darcy Flow remains challenging.

  • Datasets and evaluation: The evaluation covers five PDE families spanning linear transport, nonlinear conservation, reaction–diffusion, compressible flow, and steady-state elliptic flow.The benchmark includes Advection, Burgers, 1D Reaction–Diffusion, 1D Compressible Navier–Stokes, and Darcy Flow.
  • Solver accuracy: DiffPDE achieves 4.5 × 10−3 nRMSE on 1D Advection, outperforming GPT-5.2, and matches state-of-the-art performance on 1D Reaction-Diffusion at 3.0 × 10−2.These comparisons are reported against large-parameter LLMs.
  • Solver accuracy: On Burgers, DiffPDE reduces error to 9.7 × 10−4, compared with errors > 1.3 for Qwen2.5-Coder-7B and Qwen3-8B.The result illustrates the reported accuracy gain over same-scale autoregressive baselines on a complex PDE.
  • Generation success rate: DiffPDE reaches a 0.78 Pass@32 on 1D Compressible Navier-Stokes, exceeding Claude-Sonnet-4.6 at 0.44 and Deepseek-v3.2 at 0.50.Pass@32 measures the proportion of candidates that execute successfully within the debugging budget.
  • Limitations: DiffPDE’s Pass@32 is 0.09 on 2D Darcy Flow, where elliptic global coupling, permeability variation, and joint equation-boundary constraints remain challenging.The discussion identifies limited elliptic-PDE priors as an important bottleneck.
  • Efficiency: Localized infilling reduces debugging generation time to below 1 second, yielding a reported 30× to 100× acceleration over autoregressive repair attempts.Reported infilling times range from 0.65 s to 0.99 s, versus roughly 20 to 120 seconds for AR repair attempts.
  • Efficiency: DiffPDE lowers worst-case average end-to-end cost to 55.5 s versus 102.1 s for Qwen2.5-Coder-7B and 119.3 s for Qwen3-8B under k = 3 repairs.The end-to-end cost combines initial generation with three repair attempts.

6 Conclusion

DiffPDE is a diffusion-LLM framework for automated PDE solver generation and repair. Across five PDEBench families, it remains competitive with larger autoregressive baselines, improves over same-scale AR models, and enables faster localized repair.

  • DiffPDE combines supervised fine-tuning with ID-GRPO to jointly optimize PDE solver generation and debugging.The training pipeline uses dual-source infilling tasks and intermediate reward signals.
  • DiffPDE replaces full-regeneration debugging with bidirectional re-masking and localized infilling of erroneous code segments.Correct context is preserved while suspected error regions are regenerated.
  • Across five PDE families from PDEBench, DiffPDE remains competitive with larger AR baselines and substantially improves over same-scale AR models.
  • DiffPDE delivers much faster localized repair while maintaining competitive solver performance.

Limitations

The study is scoped to localized PDE solver debugging, while global structural revisions and broader evaluation settings remain open directions.

  • DiffPDE focuses on localized debugging where execution feedback identifies specific code spans for revision.
  • Failures requiring global structural revisions remain beyond the scope of the localized infilling framework.
  • Future work should evaluate DiffPDE on more diverse PDE regimes, realistic scientific programming environments, and richer interactive debugging workflows.

Ethical Statement

The work uses model-generated solver code and mitigates potential instability through execution-based verification and quantitative filtering. Experiments rely on PDEBench and automatically constructed solver trajectories from benchmark tasks.

  • Model-generated solver code may be unstable, so the study applies execution-based verification and quantitative filtering.
  • The experiments use PDEBench and automatically constructed solver trajectories derived from benchmark PDE tasks.
  • The reported data consist of synthetic scientific problems and code and are described as containing no personally identifying information or offensive content.

A Additional Method Details

DiffPDE localizes repair by deterministically mapping execution traceback information to a code span, while ID-GRPO supplies intermediate rewards for progressive error reduction during debugging.

  • Deterministic Error Localization: The error-localization procedure extracts the failing line number L from the execution error message.
  • Deterministic Error Localization: The procedure expands from L until reaching a blank line, comment-only line, or shallower-indented line.
  • Deterministic Error Localization: The expanded block is converted to a token span, masked, and infilled using preserved left and right context.
  • Intermediate-Reward Construction: Intermediate reward is granted only when the new normalized error-signature set is a strict subset of the previous set.
  • Intermediate-Reward Construction: New syntax errors or new non-syntax error signatures invalidate the intermediate reward under the default configuration.

B Benchmark and Data Details

The benchmark evaluates generated PDE solvers across five representative PDE families under standardized execution and array-based comparison. It specifies task-specific grids, temporal dimensions, boundary conditions, data splits, and solver interfaces.

  • Benchmark scope: PDEBench covers Advection, Burgers, Reaction–Diffusion, Compressible Navier–Stokes, and Darcy Flow across time-dependent and steady-state regimes.Generated solvers are executed in a standardized harness and compared with corresponding ground-truth solution trajectories or fields.
  • PDE families and conditions: The PDE task definitions use periodic conditions for the listed one-dimensional families, while Darcy Flow uses Dirichlet conditions.The benchmark includes 1D Advection, Burgers, Reaction–Diffusion, and Compressible Navier–Stokes, plus 2D Darcy Flow.
  • Solver interfaces: Solver interfaces require task-specific Python functions whose generated outputs are directly executed and uniformly scored across PDE families.Inputs and outputs vary by task, including batched initial conditions, time coordinates, multiple physical fields, or coefficient fields.
  • Data splits: Each selected PDEBench parameter setting contributes mutually exclusive SFT, RL, and test subsets of 100 instances.The settings are represented as HDF5 tensors under fixed discretizations and sampled with a fixed random seed.
  • Discretization: One-dimensional Advection, Burgers, and Reaction–Diffusion use N = 1024 and Nt = 200, while one-dimensional CNS uses N = 1024 and Nt = 100.Darcy Flow uses a 128 × 128 spatial grid.
  • Tensor shapes: Time-dependent outputs include the initial frame and T future steps, with shape R^B×(T+1)×N.Darcy Flow has no temporal dimension because it is steady-state.

C Training Environment and Hyperparameters

DiffPDE combines staged supervised training with reinforcement learning and uses distinct prompting for full generation versus localized diffusion repair. The task prompts specify executable Python solvers, PDE-specific interfaces, and required tensor shapes.

  • Training stages: DiffPDE is trained through generation SFT, infilling-aware SFT, and reinforcement learning over multi-PDE generation and infilling tasks.Training starts from the Dream-Coder-7B checkpoint.
  • Prompting: DiffPDE uses a debugging prompt for re-masking and span-level infilling, whereas the autoregressive baseline uses the same full-regeneration prompt structure as CodePDE.The comparison keeps the from-scratch and AR debugging prompt structure aligned while changing the diffusion repair prompt.
  • Prompt variants: The training and evaluation materials include separate system prompts for from-scratch generation and debugging prompts for autoregressive full regeneration.The from-scratch prompt requires a complete, efficient, accurate Python solver with concise comments.
  • Localized repair: The repair prompt instructs the model to fill only MASKed code spans while preserving the surrounding reasoning, structure, and intent.It explicitly prohibits rewriting the rest of the answer or generating a completely new solution.
  • Task specification: The representative advection task asks for batched Python prediction of a trajectory with shape [batch_size, T+1, N] under periodic boundaries.The prompt describes the 1D advection equation and notes that smaller internal time steps may improve simulation stability.

E Additional Case Study: Localized Repair on Burgers

A Burgers case study shows execution-guided localized repair of a runtime failure by re-masking the faulty diffusion-factor span and infilling a minimal correction. The repaired solver executes successfully with nRMSE 0.1.

  • Workflow implication: The case study connects execution feedback, error localization, re-masking, and span-level infilling into a repair workflow.This workflow is contrasted with expensive full regeneration.
  • Initial failure: The initial Burgers solver fails with UnboundLocalError because it uses dt to construct the diffusion factor before dt is assigned inside the time loop.The error is explicitly reported as local variable dt referenced before assignment.
  • Localized intervention: DiffPDE re-masks only the error-localized span around the diffusion-factor construction instead of regenerating the full solver.The remaining code is kept unchanged and the repair is represented as masked span infilling.
  • Repair mechanism: The repaired Burgers solver computes the diffusion factor using the time difference between adjacent output times.The case uses viscosity ν = 0.1 and an implicit Fourier-space diffusion treatment.
  • Outcome: nRMSE of 0.1 is achieved after repair, and the corrected solver executes successfully.The example demonstrates recovery from a concrete runtime failure through a minimal code edit.
Loading 2608.30532v1…