pADAM: A Plug-and-Play All-in-One Diffusion Architecture for Multi-Physics Learning

Abstract

Generalizing across disparate physical laws remains a fundamental challenge for artificial intelligence in science. Existing deep-learning solvers are largely confined to single-equation settings, limiting transfer across physical regimes and inference tasks. Here we introduce pADAM, a unified generative framework that learns a shared probabilistic prior across heterogeneous partial differential equation families. Through a learned joint distribution of system states and, where applicable, physical parameters, pADAM supports forward prediction and inverse inference within a single architecture without retraining. Across benchmarks ranging from scalar diffusion to nonlinear Navier--Stokes equations, pADAM achieves accurate inference even under sparse observations. Combined with conformal prediction, it also provides reliable uncertainty quantification with coverage guarantees. In addition, pADAM performs probabilistic model selection from only two sparse snapshots, identifying governing laws through its learned generative representation. These results highlight the potential of generative multi-physics modeling for unified and uncertainty-aware scientific inference.

Figures (10)

• Figure 1: Schematic of the pADAM framework for unified multi-physics learning.a--c, The pADAM framework learns across disparate physical laws, illustrated here for scalar-field PDEs, by projecting heterogeneous equation families into a shared generative prior. A class-conditional diffusion model learns the joint distribution of system states and physical parameters (a, b), enabling the generation of diverse physical regimes from Gaussian noise (c, orange trajectories). d, Task-agnostic inference via Bayesian conditioning. By incorporating full or sparse observations through plug-and-play guidance (green trajectory), the shared pADAM prior supports forward prediction, initial condition reconstruction, parameter inference, and probabilistic model selection within a single framework. This unified manifold allows pADAM to navigate a range of inference tasks without task-specific retraining.
• Figure 1: Mechanistic analysis of internal attention patterns across physical operators. Representative attention maps are extracted from a decoder block at a fixed denoising step ($t=1000$) for diffusion, advection, and advection--diffusion regimes under identical initial conditions. Comparison of local attention weights reveals distinct structural signatures: the advection--diffusion maps (right) exhibit intermediate activations in spatial regions where features are prominent in pure advection (middle) but absent in pure diffusion (left). This graded attention pattern across the three operator classes suggests that the shared pADAM prior reuses and composes operator-specific internal representations to capture complex mixed dynamics.
• Figure 2: Forward and inverse inference for Navier--Stokes dynamics using pADAM.a, Forward prediction of the velocity component $v_T$ conditioned on 30% spatial observations of the initial states $(u_0, v_0)$. b, Inverse reconstruction of the initial velocity component $v_0$ conditioned on full observation of the terminal component $v_T$ and the initial component $u_0$. c, Forward prediction of the velocity component $u_T$ conditioned on 30% spatial observations of the initial states $(u_0, v_0)$.
• Figure 2: Quantitative assessment of representational similarity across PDE families. Pairwise cosine similarities of attention maps calculated for one encoder and one decoder block at two representative denoising steps ($t=500$ and $t=1000$) across five test samples. Across all configurations, similarities between the advection--diffusion (mixed) regime and the pure extremes (diffusion or advection) are consistently higher than the similarity between the two pure extremes. This hierarchy indicates that the model's latent space is organized according to the underlying mathematical composition of the physical laws, with the mixed operator acting as a representational bridge.
• Figure 3: Advection regime