Smooth and Sparse Latent Dynamics in Operator Learning with Jerk Regularization

TL;DR

The paper addresses the challenge of learning efficient, long-term spatiotemporal forecasts for systems governed by PDEs without requiring full governing equations. It introduces a spatiotemporal continuous operator learning framework that couples an autoencoder with jerk regularization, a neural ODE latent dynamics model, and a conditional INR decoder, enabling continuous space-time inference. On a 2D Navier–Stokes test case, jerk regularization improves accuracy, accelerates training convergence, and reveals a small set of intrinsic latent coordinates, while supporting high-resolution predictions from low-resolution training data. This approach offers a path to rapid, high-fidelity simulations and insight into the latent structure of complex dynamical systems.

Abstract

Spatiotemporal modeling is critical for understanding complex systems across various scientific and engineering disciplines, but governing equations are often not fully known or computationally intractable due to inherent system complexity. Data-driven reduced-order models (ROMs) offer a promising approach for fast and accurate spatiotemporal forecasting by computing solutions in a compressed latent space. However, these models often neglect temporal correlations between consecutive snapshots when constructing the latent space, leading to suboptimal compression, jagged latent trajectories, and limited extrapolation ability over time. To address these issues, this paper introduces a continuous operator learning framework that incorporates jerk regularization into the learning of the compressed latent space. This jerk regularization promotes smoothness and sparsity of latent space dynamics, which not only yields enhanced accuracy and convergence speed but also helps identify intrinsic latent space coordinates. Consisting of an implicit neural representation (INR)-based autoencoder and a neural ODE latent dynamics model, the framework allows for inference at any desired spatial or temporal resolution. The effectiveness of this framework is demonstrated through a two-dimensional unsteady flow problem governed by the Navier-Stokes equations, highlighting its potential to expedite high-fidelity simulations in various scientific and engineering applications.

Figures (5)

• Figure 1: Schematic of the Spatiotemporal Continuous Operator Learning Model.(a) Inference Stage: The encoder transforms the initial system state $\boldsymbol{u}(0)$ into a compressed latent representation $\boldsymbol{z}(0)$. This initial latent vector $\boldsymbol{z}(0)$ is then evolved using a neural ODE to obtain the corresponding latent state $\boldsymbol{z}(t)$ at any subsequent time $t$. A conditional INR decoder then reconstructs the continuous system state $u(\boldsymbol{x},t)$ from $\boldsymbol{z}(t)$ with an arbitrarily high spatial resolution by sampling at any desired spatial location. (c) Training Stage I: The autoencoder is trained to learn a nonlinear dimensionality reduction mapping between original and latent spaces yielding smooth and sparse latent space trajectories. To that effect, a reconstruction loss forces the output state to match the corresponding input at each time, while a jerk regularization loss enforces smoothness and sparsity by penalizing the third-order time derivative of the latent vector, approximated for a given time using four sequential latent vectors. (c) Training Stage II: The neural ODE is trained to learn a continuous-time model of the dynamics within the latent space. Once trained, it generates latent vectors at any desired time step, starting from the initial latent state $\boldsymbol{z}(0)$, allowing for predictions that extend beyond the temporal scope of the training data and with variable time granularity.
• Figure 2: Comparison of ground truth vorticity fields and model predictions for the Navier-Stokes equations example.(a) Time interpolation and extrapolation comparisons in the test set. (b) Error propagation inside (green area) and outside (pink area) the training time domain. (c) The proposed method can conduct arbitrary high-resolution forecasting in space and time.
• Figure 3: Jerk regularization promotes smoothness and sparsity in latent space. The average jerk is calculated by taking the mean value of the jerk measurements across all snapshots in a trajectory.
• Figure 4: Comparison of the stage I loss history with and without jerk regularization for a latent space dimension of 10.(a) MSE loss comparison for training and test sets. (b) Jerk loss comparison for training and test sets.
• Figure 5: Comparison of different jerk coefficients for autoencoder in the test set.