GD-VAEs: Geometric Dynamic Variational Autoencoders for Learning Nonlinear Dynamics and Dimension Reductions

TL;DR

This work introduces Geometric Dynamic Variational Autoencoders (GD-VAEs), which extend variational autoencoders to learn nonlinear dynamical systems with latent spaces that have general geometries and topologies. By embedding manifold latent spaces extrinsically into Euclidean spaces and using projection maps with gradients computed via implicit differentiation, GD-VAEs enforce topology-aware representations and enable stable multi-step predictions. The authors derive a flexible training framework with ELBO-based losses and regularization terms, and demonstrate the approach on Burgers’ equation, constrained mechanical systems (torus and Klein bottle latents), and 2D Brusselator reaction–diffusion dynamics, showing improved stability, interpretability, and data efficiency compared with traditional methods. The results highlight how incorporating geometric and topological priors into latent spaces can yield parsimonious, robust, and interpretable models for high-dimensional dynamical systems, with broad potential for physics-informed learning and reduced-order modeling.

Abstract

We develop data-driven methods incorporating geometric and topological information to learn parsimonious representations of nonlinear dynamics from observations. The approaches learn nonlinear state-space models of the dynamics for general manifold latent spaces using training strategies related to Variational Autoencoders (VAEs). Our methods are referred to as Geometric Dynamic (GD) Variational Autoencoders (GD-VAEs). We learn encoders and decoders for the system states and evolution based on deep neural network architectures that include general Multilayer Perceptrons (MLPs), Convolutional Neural Networks (CNNs), and other architectures. Motivated by problems arising in parameterized PDEs and physics, we investigate the performance of our methods on tasks for learning reduced dimensional representations of the nonlinear Burgers Equations, Constrained Mechanical Systems, and spatial fields of Reaction-Diffusion Systems. GD-VAEs provide methods that can be used to obtain representations in manifold latent spaces for diverse learning tasks involving dynamics.

Figures (15)

• Figure 1: Learning Nonlinear Dynamics with Encoders and Decoders. Data-driven methods are developed for learning models to predict from $u(x,t)$ the non-linear evolution to future states $u(x,t + \tau)$ for PDEs and other dynamical systems. Probabilistic autoencoders are utilized that learn representations $\mathbf{z}$ of $u(x,t)$ in reduced dimensional latent spaces with prescribed geometric and topological properties. The trainable models make predictions using learnable maps that (i) encode an input $u(x,t) \in \mathcal{U}$ as $\mathbf{z}(t)$ in latent space (left), (ii) evolve the representation $\mathbf{z}(t) \rightarrow \mathbf{z}(t + \tau)$(middle), (iii) decode the representation $\mathbf{z}(t + \tau)$ to predict $\hat{u}(x,t + \tau)$(right).
• Figure 2: Variational Autoencoder (VAE) Framework. Representations of the non-linear dynamics are learned using probabilistic encoders and decoders trained using the VAE framework KingmaWellingVAE2014. Our deep neural network (DNN) models have the VAE architecture shown at the (top). This includes learnable probabilistic mappings, in contrast to more conventional neural network autoencoders shown on the (bottom-right). GD-VAEs includes a probabilistic encoder with distribution $\mathfrak{q}_{\theta_e}$ and a probabilistic decoder with distribution $\mathfrak{p}_{\theta_d}$, shown on the (bottom-left). These DNNs are trained (i) to serve as feature extractors to represent functions $u(x,t)$ and their evolution in a reduced dimensional latent space as $\mathbf{z}(t) \rightarrow \mathbf{z}(t + \tau)$, and (ii) to serve as approximators that can construct predictions $u(x,t + \tau)$ using the features $\mathbf{z}(t+ \tau)$.
• Figure 3: Learnable Mappings for Manifold Latent Spaces. We develop methods for using manifold latent space representations having general geometries and topologies through projection maps. For inter-operability with widely used trainable mappings, such as neural networks, we use a strategy of mapping inputs first to a point $\mathbf{w} = \tilde{\mathcal{E}}_\phi(\mathbf{x})\in\mathbb{R}^N$ for the embedding space which is then projected to a point in the manifold $\mathbf{z} = \Lambda(\mathbf{w}) \in \mathcal{M}\subset \mathbb{R}^N$. This provides trainable mappings for general manifold latent spaces. In practice, we can represent the manifold and compute projections based on general point cloud representations, analytic descriptions, product spaces, or other approaches.
• Figure 4: Burgers' Equation: Prediction of Dynamics. We consider responses for initial conditions $\mathcal{U}_1 = \{ u\, | \, u(x,t;\alpha) = \alpha \sin(2\pi x) + (1-\alpha) \cos^3(2 \pi x)\}$. Predictions are made for the evolution $u$ over the time-scale $\tau$ satisfying equation \ref{['eqn:burgers']} with initial conditions in $\mathcal{U}_1$. We find our nonlinear VAE methods are able to learn with $2$ latent dimensions the dynamics with errors $< 1\%$. Methods such as DMD DMD_Schmid_2010DMD_Theory_and_App_Kutz_2014 with $3$ modes which are only able to use a single linear space to approximate the initial conditions and prediction encounter challenges in approximating the nonlinear evolution. We find our linear VAE method with $2$ modes provides some improvements by allowing for using different linear spaces for representing the input and output functions. Results are summarized in Table \ref{['table:Burgers_compare_VAE_DMD_etc']}.
• Figure 5: Burgers' Equation: Latent Space Representations and Extrapolation Predictions. We show the latent space representation $z$ of the dynamics for the input functions $u(\cdot,t;\alpha) \in \mathcal{U}_1$. VAE organizes for $u$ the learned representations $z(\alpha,t)$ in parameter $\alpha$(blue-green) into circular arcs that are concentric in the time parameter $t$, (yellow-orange)(left). The reconstruction regularization with $\gamma$ aligns subsequent time-steps of the dynamics in latent space facilitating multi-step predictions. The learned VAE model exhibits a level of extrapolation to predict dynamics even for some inputs $u \not \in \mathcal{U}_1$ beyond the training data (right).