[Paper Review] ODE$^2$VAE: Deep generative second order ODEs with Bayesian neural networks

Çağatay Yıldız, Markus Heinonen+1 more|arXiv (Cornell University)|May 27, 2019

Generative Adversarial Networks and Image Synthesis35 references62 citations

TL;DR

ODE2VAE learns latent second-order ODE dynamics with a Bayesian neural network, decomposing latent space into position and velocity to model high-dimensional sequential data and improve long-term forecasting and imputation.

ABSTRACT

We present Ordinary Differential Equation Variational Auto-Encoder (ODE$^2$VAE), a latent second order ODE model for high-dimensional sequential data. Leveraging the advances in deep generative models, ODE$^2$VAE can simultaneously learn the embedding of high dimensional trajectories and infer arbitrarily complex continuous-time latent dynamics. Our model explicitly decomposes the latent space into momentum and position components and solves a second order ODE system, which is in contrast to recurrent neural network (RNN) based time series models and recently proposed black-box ODE techniques. In order to account for uncertainty, we propose probabilistic latent ODE dynamics parameterized by deep Bayesian neural networks. We demonstrate our approach on motion capture, image rotation and bouncing balls datasets. We achieve state-of-the-art performance in long term motion prediction and imputation tasks.

Motivation & Objective

• Provide a probabilistic, continuous-time latent dynamics model for high-dimensional sequential data.
• Decompose latent space into position and velocity to capture second-order dynamics.
• Incorporate Bayesian neural networks to model uncertain dynamics and prevent overfitting.
• Optimize via variational inference to jointly learn initial latent states, dynamics, and reconstructions.
• Demonstrate state-of-the-art results on motion capture, rotating MNIST, and bouncing balls datasets.

Proposed method

• Introduce a second-order ODE in latent space with position s_t and velocity v_t.
• Parameterize the acceleration field f_W(s_t,v_t) with a Bayesian neural network over weights W.
• Use a variational auto-encoder framework with encoders for initial position s_0 and initial velocity v_0 and a decoder for observations x_t.
• Apply the instantaneous change of variable to derive log-density evolution for the latent ODE flow.
• Employ a penalized variational loss (ODE2VAE-KL) to balance reconstruction, dynamics and prior constraints.
• Solve the latent dynamics via numerical ODE integration (RK method) and train with reparameterization and Monte Carlo estimates.

Experimental results

Research questions

• RQ1 How can a second-order, continuous-time latent dynamics model be learned for high-dimensional sequential data?
• RQ2 Does decomposing latent space into position and velocity with a Bayesian acceleration field improve long-term forecasting and imputation?
• RQ3 Can variational inference with Bayesian neural networks regularize latent dynamics to prevent overfitting on complex sequences?
• RQ4 How does ODE2VAE compare to RNNs and black-box ODE approaches on motion capture and image sequence datasets?
• RQ5 What are the benefits of a KL-weighted regularization scheme (ODE2VAE-KL) for balancing latent and dynamic losses?

Key findings

• Achieves state-of-the-art long-term motion prediction and imputation on high-dimensional sequences.
• Demonstrates improved long-horizon forecasting on motion capture, rotating MNIST, and bouncing balls datasets.
• Bayesian second-order dynamics with latent position-velocity decomposition helps mitigate overfitting and improves generalization.
• Outperforms neural ODEs and RNN-based models on several sequential tasks, particularly in extrapolation and missing-data scenarios.
• The KL-regularized variant (ODE2VAE-KL) further stabilizes training and enhances latent-flow alignment with encoder outputs.