Learning Where to Simulate: Generative Active Sampling for Online PDE Surrogate Training

2026-06-08

AI Summary

This work addresses the challenge that uniform sampling in training data generation for PDE surrogate models often fails to adequately cover regions of complex dynamics, leading to high prediction errors and large variance. To overcome this limitation, the authors propose an Online Generative Active Sampling (OGAS) framework, which introduces a conditional diffusion model as an active sampler for the first time. OGAS dynamically couples PDE solving with surrogate training, guiding data generation toward difficult configurations by leveraging the surrogate’s loss or uncertainty estimates. Experiments on 2D Kuramoto–Sivashinsky, Navier–Stokes, and Gray–Scott equations demonstrate that OGAS significantly reduces both the 99th percentile prediction error and overall error dispersion across parameter spaces up to 308 dimensions, while incurring negligible computational overhead. This approach effectively enhances tail performance and worst-case reliability of the surrogate models.

Abstract

Data-driven PDE surrogates are trained with data produced by numerical PDE solvers. However, when the surrogate's goal is to generalize across a wide range of PDE configurations (e.g., initial conditions and physical coefficients), generating a representative training set is non-trivial. Uniform sampling of configuration parameters often under-represents trajectories exhibiting challenging dynamics, leading to high prediction errors and large error variance in the trained surrogate. Online training, where data generation and surrogate training are coupled, offers a natural advantage by allowing solver parameters to be steered on-the-fly. To efficiently exploit this capability, we introduce Online Generative Active Sampling (OGAS), an active learning method that reactively learns the relationship between configuration parameters and surrogate performance to control the sampling distribution. OGAS trains a fast diffusion model in parallel to the surrogate to act as a conditional sampler, mapping a surrogate-derived difficulty signal (e.g., loss or uncertainty) to configuration parameters. By actively drawing target signals from a prior biased toward high difficulty, OGAS continuously steers data generation toward challenging regimes without delaying the training workflow. We evaluate OGAS across 2D PDEs with distinct challenging dynamics (Kuramoto-Sivashinsky, Navier-Stokes, Gray-Scott) and up to 308 parameters, using multiple surrogate architectures. Across all settings, OGAS consistently improves tail statistics, yielding substantial reductions in errors above the 99th percentile and overall error dispersion compared to uniform sampling. While prioritizing challenging trajectories introduces a trade-off with average error, OGAS effectively ensures worst-case reliability of trained surrogates with negligible wall-time overhead.

Problem

• PDE surrogate
• active sampling
• online training
• configuration parameters
• error variance

Innovation

• Online Generative Active Sampling
• PDE surrogate
• diffusion model
• active learning
• tail error reduction