Advancing Generalization in PINNs through Latent-Space Representations

TL;DR

PiDo introduces a latent-space physics-informed PDE solver that generalizes across initial conditions, PDE coefficients, and training horizons by encoding PDE solutions into a low-dimensional latent space and evolving these embeddings with coefficient-conditioned Neural ODEs. The method couples a grid-independent INR-based decoder with an encoder via auto-decoding to map initial data into latent trajectories and applies regularization in latent space to stabilize training and improve long-horizon extrapolation. Latent Dynamics Smoothing and Latent Dynamics Alignment address instability and embedding drift in the physics-informed loss, enabling robust generalization. Empirical results on 1D Burgers–KdV and 2D Navier–Stokes benchmarks show PiDo outperforms data-driven and other physics-informed baselines and that its learned representations transfer well to downstream tasks such as long-term integration and inverse problems, highlighting practical impact for flexible PDE modeling under varying conditions.

Abstract

Physics-informed neural networks (PINNs) have made significant strides in modeling dynamical systems governed by partial differential equations (PDEs). However, their generalization capabilities across varying scenarios remain limited. To overcome this limitation, we propose PIDO, a novel physics-informed neural PDE solver designed to generalize effectively across diverse PDE configurations, including varying initial conditions, PDE coefficients, and training time horizons. PIDO exploits the shared underlying structure of dynamical systems with different properties by projecting PDE solutions into a latent space using auto-decoding. It then learns the dynamics of these latent representations, conditioned on the PDE coefficients. Despite its promise, integrating latent dynamics models within a physics-informed framework poses challenges due to the optimization difficulties associated with physics-informed losses. To address these challenges, we introduce a novel approach that diagnoses and mitigates these issues within the latent space. This strategy employs straightforward yet effective regularization techniques, enhancing both the temporal extrapolation performance and the training stability of PIDO. We validate PIDO on a range of benchmarks, including 1D combined equations and 2D Navier-Stokes equations. Additionally, we demonstrate the transferability of its learned representations to downstream applications such as long-term integration and inverse problems.

Figures (9)

• Figure 1: Physics-informed neural PDE solvers. Colorful squares denote spatial/temporal coordinates ("$x$" and "$t$") and embeddings of PDE solutions ("$c$") at different times, while curves represent continuous trajectories of embeddings unrolled by the dynamics model, with colors indicating different values. The "$\times$" denotes Cartesian product. "ENC", "DEC" and "DYN" stand for the encoder, decoder and dynamics model, respectively. PiDo encodes PDE solutions into a compact latent space, enabling generalization to unseen initial conditions. It models and unrolls latent trajectories under varying PDE coefficients ($\boldsymbol{\alpha}$ and $\boldsymbol{\alpha'}$) using a dynamics model, leveraging its strength in temporal extrapolation. The entire framework is trained end-to-end with a physics-informed loss.
• Figure 2: Training instability arises from overly complex dynamics. We randomly sample 3 dimensions from the 128-dim embeddings. We visualize the loss and predictions of a trajectory at different time steps. Training steps ($t = \{0, 1, 2, 3, 4\}$) are highlighted in green, while the test step ($t = 0.5$) is shown in red.
• Figure 3: Time extrapolation degradation arises from latent embedding drift. We take the CE3 setting as an example. The black dash line separates training and testing horizon. For visualization, we randomly sample 3 dimensions from the 64-dim embeddings.
• Figure 4: Regularization. We apply smoothing regularization $R_{S}$ and alignment regularization $R_{A}$ to the predicted latent trajectory (the blue curve). $R_{S}$ prevents rapid local changes, while $R_{A}$ aligns the predicted latent embeddings with the anchor embeddings (the orange curve), which are obtained by auto-decoding the predicted solutions. This ensures that the embeddings remain within their typical range.
• Figure 5: Performance (y-axis) of various models on the NS2 test set with unseen Reynolds numbers (x-axis). $L_2$ relative error ($\%$) is reported.