PDE-Controller

LLMs for Autoformalization and Reasoning of PDEs

Mauricio Soroco 1*, Jialin Song 1*, Mengzhou Xia 2, Kye Emond 3,4, Weiran Sun 3, Wuyang Chen 1
*Equal Contribution
1School of Computing Science, Simon Fraser University
2Department of Computer Science, Princeton University 3Department of Mathematics, Simon Fraser University 4Department of Physics, Simon Fraser University
arXiv Code Data

Heat Example

Wave Example

Introduction

We present PDE-Controller, a framework that enables large language models (LLMs) to control systems governed by partial differential equations (PDEs). Traditional LLMs have excelled in commonsense reasoning but fall short in rigorous logical reasoning. While recent AI-for-math has made strides in pure mathematics, areas of applied mathematics, particularly PDEs, remain underexplored despite their significant real-world applications. Our approach enables LLMs to transform informal natural language instructions into formal specifications, and then execute reasoning and planning steps to improve the utility of PDE control. We build a holistic solution comprising datasets (both human-written cases and 2 million synthetic samples), math-reasoning models, and novel evaluation metrics, all of which require significant effort. Our PDE-Controller significantly outperforms the latest open-source and GPT models in reasoning, autoformalization, and program synthesis, achieving up to a 62% improvement in utility gain for PDE control. By bridging the gap between language generation and PDE systems, we demonstrate the potential of LLMs in addressing complex scientific and engineering challenges.

Figure: Overview of our PDE-Controller framework. The Translator directly autoformalizes an informal PDE control problem (yellow) into formal specifications with STL (blue). The Controller proposes novel STL subgoals (purple). Each STL is synthesized into specialized Python programs by the Coder (green) and optimized externally (white). From the initial condition (i.), our PDE reasoning optimizes a subgoal (ii.) before the original problem, improving the utility at the end of control (iii.). We train the Controller with reinforcement learning from human feedback (RLHF).

Figure: Case study of LLM reasoning for PDE control on heat (top) and wave (bottom) problems.



The website is under construction. We will continue updating this project!


Citation 

          
@article{soroco2025pdecontroller,
  title={PDE-Controller: LLMs for Autoformalization and Reasoning of PDEs},
  author={Soroco, Mauricio and Song, Jialin and Xia, Mengzhou and Emond, Kye and Sun, Weiran and Chen, Wuyang},
  journal={arXiv preprint arXiv:2502.00963},
  year={2025}
}