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

Amirhossein Mollaali; Bongseok Kim; Christian Moya; Guang Lin

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.

PDE system	Observation	Model	Forward ( $u_{T}$ )	Inverse ( $u_{0}$ )
Diffusion	Full (100%)	pADAM	$0.82$	$0.45$
	Full (100%)	Single-PDE	$0.55$	$0.68$
	Sparse (30%)	pADAM	$0.88$	$0.86$
	Sparse (30%)	Single-PDE	$1.00$	$2.31$
Advection	Full (100%)	pADAM	$1.03$	$1.25$
	Full (100%)	Single-PDE	$1.02$	$1.26$

Table 1: Comparative performance against PDE-specific baselines. Relative

L_{2}

error (%) for forward prediction (

u_{T}

) and inverse reconstruction (

u_{0}

) across fixed-parameter PDE systems. We compare pADAM (trained on a diverse diffusion, advection, and advection–diffusion library,

N=1,000

) against PDE-specific diffusion baselines (

N=500

per family) following DiffusionPDE Huang et al. (2024) . Results are averaged over 50 test instances under full (100%) and sparse (30%) spatial observations. The lowest errors in each row are shown in bold.

a — Scalar-field PDEs
PDE system	Observation	Forward ( $u_{T}$ )	Inverse ( $u_{0}$ )	Inverse ( $\phi$ )
Diffusion	Full (100%)	1.37	1.11	4.13
Diffusion	Sparse (30%)	2.03	2.40	5.60
Advection	Full (100%)	1.93	1.96	1.51
Advection	Sparse (30%)	1.93	1.97	2.24
Advection–diff.	Full (100%)	1.96	2.28	9.26

Table 2: Structural scaling across the multi-physics spectrum. Relative

L_{2}

errors (%) for pADAM evaluated on a library of six distinct physical regimes. a, Performance on scalar-field PDEs (diffusion, advection, advection–diffusion, and Allen–Cahn) for forward prediction (

u_{T}

), initial-state reconstruction (

u_{0}

), and parameter discovery (

\phi

). b, Performance on vector-field PDEs (Burgers’ and Navier–Stokes) where the model learns component-wise joint distributions to provide forward (

u_{T},v_{T}

) and inverse (

u_{0},v_{0}

) state estimation. All tasks are performed under both full (100%) and sparse (30%) spatial observations, with results reported as the mean over 50 test instances.

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})

.

a — Task-agnostic inference across a heterogeneous PDE library with different parametric dimensionality
PDE system	Observation	Forward ( $u_{T}$ )	Inverse ( $u_{0}$ )	avg. Inverse ( $\phi$ )
Diffusion	Full (100%)	0.98	1.18	2.89
Diffusion	Sparse (30%)	1.82	2.07	4.61
Advection	Full (100%)	0.85	1.03	6.57
Advection	Sparse (30%)	1.64	1.32	6.61
Advection–diff.	Full (100%)	1.49	1.72	5.99

Table 3: Parametric scaling and partial-parameter inference performance. a, Task-agnostic inference across a heterogeneous PDE library with different parametric dimensionality. Relative

L_{2}

errors (%) for state estimation (

u_{0}

and

u_{T}

) and parameter recovery (

\phi

) across three physical regimes with varying parametric depth: Diffusion (

\phi=[\nu]

), Advection (

\phi=[a_{x},a_{y}]

), and Advection–diffusion (

\phi=[\nu,a_{x},a_{y}]

). b, Impact of partial parameter constraints on parameter inference in the advection system. Comparison of joint inference (

a_{x},a_{y}

) against conditional settings where one velocity component (

a_{x}

or

a_{y}

) is provided at inference time. All tasks are performed under both full (100%) and sparse (30%) spatial observations, with results representing the mean over 50 test instances.

Figure 3: Probabilistic model selection across three PDE classes. Identification of governing laws from two snapshots

(u_{0},u_{T})

with only 30% of the spatial field observed. a–c, Left panels show sparse field snapshots; right panels compare the ground-truth PDEs with representative sampled PDEs and the associated 95% predictive intervals derived from the ensemble.

PDE group	IC parameter	Sampling rule
Scalar PDEs	Centroid $(x_{c},y_{c})$	$\mathcal{U}[0.2,0.8]^{2}$
	Gaussian Width $w_{0}$	$\mathcal{U}[0.025,0.075]$
Burgers’	Component Centroids $(c_{x,i},c_{y,i})$	$\mathcal{U}[0.2,0.8]^{2}$
	Component Widths $w_{1},w_{2}$	$\mathcal{U}[0.025,0.075]$
Navier–Stokes	Amplitude Factor $a$	$a\sim\mathcal{U}(0.5,1.5]$
	Trigonometric Basis $\phi,\psi$	$(\phi,\psi)\in\{\sin,\cos\}^{2}$

Extended Data Table 1 : Experimental configurations and parameter manifolds. a, Initial condition (IC) protocol standardized across all experimental regimes. ICs for scalar and Burgers’ systems follow localized Gaussian forms adapted to their respective domains (Eqs. 12 , 24 ), while Navier–Stokes follows the solenoidal trigonometric form (Eq. 27 ). b, Systematic variability of governing coefficients

\phi

across thematic investigations. These manifolds define the continuous physical regimes analyzed in the corresponding Results sections.

Extended Data Fig. 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.

Extended Data Fig. 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.

Extended Data Fig. 3 : Zero-shot extrapolation performance on unseen PDE (advection–diffusion–reaction dynamics). a, Quantification of the operator shift (

\Delta_{\mathrm{op}}

) between the advection–diffusion (AD) training prior and the unseen advection–diffusion–reaction (ADR) dynamics for two reaction rates

k

. This shift is defined as the relative

L_{2}

discrepancy between the terminal states (

u_{T}

) of the two systems under identical initial conditions; as

k

increases, the physical divergence between the two PDE systems grows. b, Relative

L_{2}

error for the joint reconstruction of full-field initial (

u_{0}

) and terminal (

u_{T}

) states conditioned on sparse spatial observations of endpoints. The model is conditioned on the closest known operator class (AD) and steered via observation-guided sampling. Error in

u_{T}

increases with

k

due to the accumulated influence of the unseen reaction term, while

u_{0}

remains stable. As observation sparsity increases, reconstruction accuracy degrades gracefully, demonstrating the robustness of the pADAM prior under physical misspecification. All results are averaged over 20 independent test instances.

Extended Data Fig. 4 : Empirical coverage of Bayesian posterior ensembles across different observation regimes. Mean predictive-interval coverage probability (PICP) of nominal 95% intervals (dashed grey line) as a function of ensemble size for forward and inverse tasks. Particularly under ill-posed conditions, such as inverse problems or prediction under sparse observations, where only 30% of the field is observed, the coverage of raw Bayesian intervals saturates well below the 95% target across all operator families. This systematic under-coverage underscores the necessity of conformal calibration for providing rigorous reliability guarantees in data-sparse or ill-posed regimes. Results are averaged over 20 test instances.

Method	Forward ( $u_{T}$ )	Inverse ( $u_{0}$ )
Ensemble only	58.33	36.31
Ensemble + Conformal	98.42	99.83

Extended Data Table 2 : Conformal calibration resolves the systematic under-coverage of Bayesian ensembles. Mean empirical coverage (%) of nominal 95% prediction intervals for the advection–diffusion system under sparse (30%) spatial observations (ensemble size

M=6

; 50 calibration instances). While raw Bayesian ensembles fail to meet the nominal target, the integration of conformal calibration into the pADAM framework effectively recovers statistical validity, ensuring reliable uncertainty quantification for physical inference. Results are averaged over 50 test instances.

Extended Data Fig. 5 : Effect of conformal calibration on prediction intervals. a, Standard ensemble-based prediction intervals show undercoverage (59.20%) for forward prediction of

u_{T}

in the advection–diffusion system under sparse (30%) observations. b, Conformal-calibrated intervals achieving 100.00% empirical coverage. c–j, One-dimensional slices at representative

y

-locations; shaded bands denote predictive intervals and blue curves indicate the true solution. Conformal calibration adaptively expands the intervals in regions of high epistemic uncertainty to recover the nominal 95% coverage. Critically, this post-hoc procedure preserves the shape of the underlying generative predictions while improving coverage reliability without model retraining.

PDE System	Observation	Forward ( $u_{T}$ )	Inverse ( $u_{0}$ )	Inverse ( $\phi$ )
Diffusion	Full (100%)	0.69	0.89	1.38
	Sparse (30%)	1.64	1.96	3.48
	Sparse (10%)	2.09	3.05	8.26
Advection	Full (100%)	1.91	1.13	0.72
	Sparse (30%)	2.11	1.47	1.48
	Sparse (10%)	2.13	2.55	2.69

Extended Data Table 3: Task-agnostic performance across the continuous physics manifold. Relative

L_{2}

errors (%) for forward prediction (

u_{T}

), initial-state reconstruction (

u_{0}

), and physical parameter discovery (

\phi

). pADAM was trained across three PDE families, each with a single variable coefficient: diffusion (

\phi=\nu

), advection (

\phi=a_{x}

), and advection–diffusion (

\phi=\nu

). Performance is evaluated under full (100%) and sparse spatial observations (30% and 10%), with results averaged over 50 test instances.

Extended Data Fig. 6 : Qualitative inference performance of pADAM under continuous physics manifold and structural scaling. a, Parameter discovery within the scalar physics manifold (diffusion, advection, and advection–diffusion). For the advection system, the physical coefficient

\phi=a_{x}

is inferred from paired states

(u_{0},u_{T})

under 10% spatial observations. b, c, Inference for the Burgers’ system under structural scaling to the 6-PDE library. b, Forward prediction of the Burgers’ velocity component

u_{T}

conditioned on full spatial observations of the initial states

(u_{0},v_{0})

. c, Inverse reconstruction of the Burgers’ initial velocity component

v_{0}

conditioned on 30% spatial observations of the terminal component

v_{T}

and the initial component

u_{0}

.

Extended Data Fig. 7 : Qualitative inference performance of pADAM under parametric scaling and heterogeneous parametric dimensionality. a, Initial-state reconstruction in the advection system with a variable parameter vector

\phi=[a_{x},a_{y}]

. The reconstructed initial state

u_{0}

is inferred by conditioning on the known parameter vector and full observation of the terminal state

u_{T}

. b, Parameter discovery in the advection–diffusion system with a three-dimensional parameter vector

\phi=[\nu,a_{x},a_{y}]

. The physical coefficients are inferred from paired states

(u_{0},u_{T})

under 30% spatial observations.

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

Abstract

Figures & tables

Physics-Informed Generative Modeling

Flow Learners for PDEs: Toward a Physics-to-Physics Paradigm for Scientific Computing

Physics-Informed Distillation of Diffusion Models for PDE-Constrained Generation

Learning to Think in Physics: Breaking Shortcut Learning in Scientific Diffusion via Representation Alignment

Abstract

Physics-Informed Generative Modeling

Explore similar work

Flow Learners for PDEs: Toward a Physics-to-Physics Paradigm for Scientific Computing

Physics-Informed Distillation of Diffusion Models for PDE-Constrained Generation

Learning to Think in Physics: Breaking Shortcut Learning in Scientific Diffusion via Representation Alignment