Out-of-Distribution Generalization for Neural Physics Solvers

TL;DR

NOVA targets the core limitation of neural physics solvers: poor generalization under distributional shifts in PDE parameters, geometries, and initial conditions. By coupling physics-guided neural architecture search with a physics-informed zero-shot adaptation of the final layer, NOVA identifies inductive biases aligned with governing equations and enables rapid, data-free extrapolation. Across nonlinear heat, diffusion-reaction, and Navier–Stokes problems, NOVA achieves superior out-of-distribution accuracy, stabilizes long-time rollouts, and supports data-light generative design via guided diffusion. This co-design of architecture and physics regularization offers a robust path toward generalizable, efficient neural physics solvers with broad AI-for-Science impact, including potential extensions to multi-physics and three-dimensional problems.

Abstract

Neural physics solvers are increasingly used in scientific discovery, given their potential for rapid in silico insights into physical, materials, or biological systems and their long-time evolution. However, poor generalization beyond their training support limits exploration of novel designs and long-time horizon predictions. We introduce NOVA, a route to generalizable neural physics solvers that can provide rapid, accurate solutions to scenarios even under distributional shifts in partial differential equation parameters, geometries and initial conditions. By learning physics-aligned representations from an initial sparse set of scenarios, NOVA consistently achieves 1-2 orders of magnitude lower out-of-distribution errors than data-driven baselines across complex, nonlinear problems including heat transfer, diffusion-reaction and fluid flow. We further showcase NOVA's dual impact on stabilizing long-time dynamical rollouts and improving generative design through application to the simulation of nonlinear Turing systems and fluidic chip optimization. Unlike neural physics solvers that are constrained to retrieval and/or emulation within an a priori space, NOVA enables reliable extrapolation beyond known regimes, a key capability given the need for exploration of novel hypothesis spaces in scientific discovery

Figures (6)

• Figure 1: Illustration of NOVA. a. Concept of NOVA. In data-scarce settings, NOVA (a physics-guided neural physics solver discovery framework) identifies the optimal neural physics solver by leveraging underlying physics. The selected, more generalizable neural physics solver achieves accurate predictions on test tasks through physics-informed zero-shot adaptation, and provides robust predictions even when the test scenario diverges from the training distribution. b. Algorithmic description of NOVA. During the training and validation process, NOVA automates the design of network architectures (e.g., network operators and connections) on a distribution of PDE scenarios (training data distribution). For each candidate architecture, lightweight training is applied to all kernel parameters except the final layer for effective feature extraction. Subsequently, physics-informed zero-shot adaptation is used for the final layer, ensuring compliance with physical laws during validation. By incorporating physics at this stage, NOVA delivers fast and accurate predictions while being robust to the similarity of the new problem to the (potentially scarce) training distribution. c. Outcome of NOVA. Compared to purely data-driven neural architecture search (NAS) approaches, NOVA demonstrates significantly better generalization performance and is consequently more effective at learning from small datasets. d. NOVA-guided generative design. NOVA can function as a fast, data-light physics-compliant regressor for guided denoising in diffusion models, enabling the generation of designs for fluidic chips with desirable performance characteristics.
• Figure 2: Performance of NOVA for the 2D diffusion-reaction equations. a. Prediction results of NOVA on a representative case in $\mathbfcal{T}_{\rm test\_id}$ show improvement after 5, 50, and 200 architectures are evaluated during NAS. b-d. Predictions from NOVA, FNO and U-Net on (b) $\mathbfcal{T}_{\rm test\_ood1}$, (c) $\mathbfcal{T}_{\rm test\_ood2}$, and (d) $\mathbfcal{T}_{\rm test\_ood3}$. e. Mean RMSE on $\mathbfcal{T}_{\rm test\_id}$ and $\mathbfcal{T}_{\rm test\_ood}$ for NOVA, FNO, and U-Net indicates FNO and U-Net performance degrades significantly more than NOVA for out-of-distribution predictions. f. Mean RMSE at each time step across tasks in $\mathbfcal{T}_{\rm test\_id}$ from NOVA, FNO, and U-Net. Shaded area indicates the standard deviation. g. Training loss histogram after five epochs on the training dataset $\mathbfcal{T}_{\rm training\_NAS}$ for all 200 NOVA architectures evaluated during NAS. Also shown are the top 20 architectures, selected based on their validation performance on the validation dataset $\mathbfcal{T}_{\rm validation\_NAS}$ after physics-informed zero-shot adaptation. h. Training loss histogram after twenty epochs on the training dataset $\mathbfcal{T}_{\rm training\_NAS}$ for all 200 NAS-U-Net architectures. The top 20 architectures, identified based on their validation performance on the validation dataset $\mathbfcal{T}_{\rm validation\_NAS}$ during NAS, are also included. For a, b, c, and d, each subfigure for $u$ and $v$ displays predictions at $t = 0.5$, $t = 2.5$, and $t = 5$. RMSE values in brackets indicate an average over the entire simulation time. The pre-trained FNO and U-Net models are obtained from takamoto2022pdebench.
• Figure 3: Performance of NOVA for the 2D Navier-Stokes equations. a-b. Representative prediction results from NOVA and NAS-U-Net on (a) $\mathbfcal{T}_{\rm test\_id}$ and (b) $\mathbfcal{T}_{\rm test\_ood}$. The three images in each subfigure represent the predictions of $u$, $v$, and $p$ while the RMSE in brackets indicates the corresponding errors. c. RMSE for NOVA and NAS-U-Net on $\mathbfcal{T}_{\rm test\_id}$ when trained with 15, 25, 35, 45, and 55 training samples. d. Relative improvement in $\Delta P$ after guided generation with NOVA. Relative improvement is calculated as the difference between the $\Delta P$ for a design obtained by unconditional diffusion and NOVA-guided diffusion, normalized by the former's $\Delta P$. e. CFD-derived velocity and pressure fields for a sample fluidic chip design generated through diffusion with and without NOVA guidance. f. Mean RMSE on $\mathbfcal{T}_{\rm test\_id}$ and $\mathbfcal{T}_{\rm test\_ood}$ for NOVA and NAS-U-Net trained on 45 geometries. The $\Delta P$ values in d and e are obtained using a CFD solver on generated geometries.
• Figure 4: Generalization performance of NOVA for the 2D nonlinear heat equation. a-d. Predictions from NOVA and NAS-U-Net on a representative case from (a) $\mathbfcal{T}_{\rm test\_id}$, (b) $\mathbfcal{T}_{{\rm test\_ood}\_F1}$, (c) $\mathbfcal{T}_{{\rm test\_ood}\_F2}$, and (d) $\mathbfcal{T}_{{\rm test\_ood}\_F3}$. Each subfigure displays predictions at $t = 0.05$, $t = 0.5$, and $t = 1$. RMSE in brackets is the average over the entire simulation time. e. Mean RMSE from NOVA and NAS-U-Net on $\mathbfcal{T}_{\rm test\_id}$ and $\mathbfcal{T}_{{\rm test\_ood}\_F3}$. f. Mean RMSE at each time step for tasks in $\mathbfcal{T}_{\rm test\_id}$ for NOVA and NAS-U-Net. Shaded area indicates the standard deviation.
• Figure 5: CNN architecture diagram. The kernel parameters before the CNN last layer $\boldsymbol{\omega}$ are learned across multiple tasks for more generalizable feature representations. The final layer parameters $\boldsymbol{\theta}$ are learned to project the extracted features for each new PDE scenario. In this work, an autoregressive rollout is employed to handle transient problems.