DiffPDE: Masked Diffusion Language Models as PDE Solver

2026-08-31Artificial Intelligence

Artificial Intelligence
AI summary

The authors propose a new method called DiffPDE to fix errors in code that solves partial differential equations (PDEs). Unlike traditional models that fix code from start to end, their approach focuses only on the parts with mistakes, making the process faster and more efficient. They also introduce a learning method, ID-GRPO, to handle problems that need multiple fixes in sequence. Their tests show that DiffPDE is both accurate and quicker than older methods.

Partial Differential EquationsPDE SolversAutoregressive ModelsDiffusion ModelsCode RepairReinforcement LearningDebuggingIterative Methods
Authors
Wenxuan Guo, Yuyang Hong, Lubin Fan, Zhaojin Fu, Lin Chen, Kun Ding, Shiming Xiang
Abstract
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.