Wassersstein pipeline improves training data for faster solving of complex equations
WaSP: Wasserstein Data Synthesis Pipeline for Solver-Free Operator Learning
Machine LearningArtificial Intelligence
Summary
Solving complex math equations called partial differential equations (PDEs) can be very slow, so people train AI models to do it faster. However, training these models usually needs lots of example input and output pairs, which can be expensive to get. The authors developed a new method called WaSP that creates better training data step-by-step, making the AI models more accurate without needing to solve the equations often. Their method works well on several challenging PDEs and improves results significantly compared to previous approaches.
What this means in practice
- •For computational engineers: Train PDE-solving surrogate models more accurately without running expensive solvers, speeding up simulations in engineering design.
- •For weather modelers: Improve statistical predictions in chaotic systems like weather by using better data for training AI surrogates of dynamic equations.
Authors
Ryan Liu, Chuwei Wang, Ricardo Baptista, Andrew Stuart, Anima Anandkumar
Abstract
Neural operators offer fast surrogates for solving PDEs, but training them typically requires expensive input-solution pairs. Two solver-free alternatives include physics-informed learning and training on manufactured data. Physics-informed learning approximates solutions by minimizing PDE residuals. While the residual norm must be strong enough to control solution error, stronger norms increase mini-batch variance and hinder optimization. Manufactured data constructs input-solution pairs for learning by applying the PDE operator to artificial solutions. However, the resulting training distribution can be mutually singular with the target distribution, yielding poor accuracy on the target distribution. We introduce the Wasserstein Data Synthesis Pipeline (WaSP), which trains on manufactured input-solution pairs while gradually mitigating the distribution shift. WaSP synthesizes new pairs by applying the learned operator to a path of intermediate distributions converging to the target input distribution. Training proceeds in stages along this path. At each stage, the previous surrogate generates solutions, and the exact PDE operator derives corresponding inputs, yielding training pairs for the next surrogate. Theoretically, we prove a stepwise error-propagation bound for the distribution shift introduced during synthesis. Empirically, WaSP reduces solution error by $3.9\times$ to $9.5\times$ compared to existing solver-free methods on PDEs such as Allen-Cahn, Kuramoto-Sivashinsky, and Ginzburg-Landau, and better reproduces the two chaotic systems' statistical observables.