PaperlayerView exact revision on arXiv

Source-linked AI summary

A Disentangled Recognition and Nonlinear Dynamics Model for Unsupervised Learning

Marco Fraccaro, Simon Kamronn, Ulrich Paquet, Ole Winther

arXiv:1710.05741v2stat.MLcs.LG

TL;DR

Existing video models often reason by generating full frames autoregressively rather than modeling object dynamics in a latent space. This paper introduces Kalman variational auto-encoders, which disentangle object representations from dynamics and outperform competing methods on generation and missing-frame imputation.

  • Problem

    Current temporal models often generate full sensory frames autoregressively, while inference for imputing unobserved frames requires using information from across the sequence.

  • Method

    Kalman variational auto-encoders disentangle object representations from dynamics by combining a recognition network with a recurrently adapted linear Gaussian state-space model that supports exact smoothing.

  • Results

    The KVAE outperforms competing methods in generative and missing-data imputation tasks, with smoothing improving inference and imputation.

  • Takeaways & Limitations

    Temporal reasoning and potentially planning can be performed in the disentangled latent space without generating high-dimensional frames at every time step.

Abstract

from arXiv · show

This paper takes a step towards temporal reasoning in a dynamically changing video, not in the pixel space that constitutes its frames, but in a latent space that describes the non-linear dynamics of the objects in its world. We introduce the Kalman variational auto-encoder, a framework for unsupervised learning of sequential data that disentangles two latent representations: an object's representation, coming from a recognition model, and a latent state describing its dynamics. As a result, the evolution of the world can be imagined and missing data imputed, both without the need to generate high dimensional frames at each time step. The model is trained end-to-end on videos of a variety of simulated physical systems, and outperforms competing methods in generative and missing data imputation tasks.

1 Introduction

The paper targets temporal reasoning in changing videos by separating object recognition from latent dynamics. It introduces Kalman variational auto-encoders, which model dynamics in latent space and outperform multiple methods on generation and missing-data imputation.

  • 1 Introduction: Humans use high-dimensional sensory representations to make temporal predictions, but end-to-end artificial models have captured recognition and Newtonian dynamics only weakly.The introduction motivates modeling a moving object's identity separately from intuition about its trajectory.
  • 1 Introduction: KVAEs disentangle an object representation from its dynamics by combining a variational auto-encoder with a linear Gaussian state-space model over latent encodings.At each time step, the variational auto-encoder compresses visual stimulus x_t into latent encoding a_t, whose learned manifold is modeled dynamically.
  • 1 Introduction: KVAEs outperform an array of methods in generative and missing-data imputation tasks.Their smoothing capability supports both tasks without requiring the paper's stated latent-space separation to be abandoned.

2 Background

The background introduces LGSSMs as first-order Markov models for sequential vectors and VAEs as deep generative models with stochastic latent encodings. LGSSMs additionally offer exact Kalman filtering and smoothing, naturally supporting missing-data handling.

  • Linear Gaussian state space models: LGSSMs model sequences of vectors through first-order Markov processes over latent states, optionally controlled by external inputs.Their transition and emission distributions are Gaussian.
  • Linear Gaussian state space models: The matrices At, Bt, and Ct represent state transition, control, and emission parameters, while Q and R are process- and measurement-noise covariance matrices.The starting state is Gaussian with mean zero and covariance Σ.
  • Linear Gaussian state space models: Exact filtered and smoothed latent-state posteriors can be computed with Kalman filter and smoother algorithms, providing a natural way to handle missing data.The relevant posteriors condition on observed sequences and, where applicable, external inputs.
  • Variational auto-encoders: VAEs define deep generative models by introducing latent encodings and approximating analytically intractable posteriors with variational distributions parameterized by φ.The model uses a likelihood pθ(xt|at) and typically Gaussian prior p(at).

3 Kalman Variational Auto-Encoders

The Kalman variational auto-encoder disentangles visual recognition from temporal dynamics by encoding frames into a low-dimensional latent representation and modeling its evolution with a Kalman-based state-space model. Its variational formulation enables exact smoothing, principled missing-data handling, and nonlinear dynamics through time-varying parameters.

  • Model architecture: KVAE encodes each video frame into a low-dimensional representation a_t that serves as a pseudo-observation for a linear Gaussian state-space model over dynamics z_t.The recognition network maps x_t to the mean and diagonal covariance of a Gaussian distribution over a_t.
  • Model architecture: The joint KVAE model factorizes into frame generation, latent representation dynamics, and latent state priors, allowing temporal reasoning and long-term prediction without autoregressively generating high-dimensional frames.The model uses pθ(x|a) pγ(a|z) pγ(z|u), with exact filtered and smoothed state distributions available from the LGSSM.
  • Variational learning: KVAE learns by replacing the intractable posterior with a variational distribution that combines neural recognition qφ(a_t|x_t) with the exact Kalman-smoothed conditional posterior pγ(z|a,u).Parameters θ, φ, and γ are jointly updated by maximizing the evidence lower bound.
  • Nonlinear dynamics: The model handles nonlinear dynamics by making the state-space parameters γ_t depend on the encoded history a_0:t−1 while preserving linear dependence between consecutive latent states.This allows events such as a ball-wall collision to alter subsequent dynamics while retaining tractable Kalman smoothing.
  • Nonlinear dynamics: An LSTM dynamics-parameter network produces time-varying weights that softly interpolate among K globally learned transition, control, and emission matrices, yielding a mixture of operating modes.The matrices A^(k), B^(k), and C^(k) represent different dynamics, with each mode emphasized according to its corresponding weight.

4 Missing data imputation

The KVAE imputes missing video frames by combining observed data with information from both past and future through LGSSM smoothing. Because its time-varying LGSSM parameters depend on missing encoded states, the method recursively estimates them during a forward pass before applying backward smoothing.

  • Motivation: Autoregressive models cannot use final frames to infer unobserved middle frames, motivating bidirectional inference for missing-data imputation.The observed sequence may include initial movement and final positions while intermediate frames are missing.
  • Smoothing-based imputation: The KVAE samples missing latent states from p(a_un, a_obs, z|x_obs, u) and generates unobserved frames from a_un, including non-contiguous missing times.Its LGSSM uses information from both past and future through smoothing.
  • Smoothing-based imputation: Ancestral sampling starts from x_obs, approximates p(a_obs|x_obs) with q_φ(a_obs|x_obs), applies Kalman smoothing with missing data, and samples p_γ(a_un|z).This procedure assumes γ is fully known when computing p_γ(z|a_obs, u).
  • Parameter estimation: Because each γ_t depends on previous encoded states including missing a_un, the KVAE recursively estimates γ during a forward pass using sampled observed states and predictive estimates for missing states.When x_1:t−1 is known but x_t is missing, the method samples a_1:t−1 from q_φ(a_1:t−1|x_1:t−1) and uses it to compute γ_1:t.
  • Parameter estimation: After the forward pass estimates γ and the filtered posterior for z, the Kalman smoother’s backward pass computes a smoothed posterior that is not exact because γ was estimated forward.At each missing step, γ_t+1 is estimated using the predictive estimate of the missing encoded state; observed states instead come from the VAE.

5 Experiments

The experiments evaluate KVAE on bouncing-ball, pendulum, and varied physical environments, testing dynamics learning, generation, and missing-data imputation. KVAE learns environment-specific dynamics, benefits from smoothing, and tolerates substantial missing training data.

  • Experimental setup: KVAE experiments use bouncing-ball videos to test separate recognition and dynamics learning, missing-data imputation, and generalization across environments.The study also compares KVAE with DVBFs on a pendulum example and trains all models end-to-end with stochastic gradients.
  • Bouncing-ball experiment: The bouncing-ball benchmark contains 5000 sequences of 20 time steps, with 32x32 binary frames and elastic wall collisions.Initial position and velocity are randomly sampled, and no forces act on the ball.
  • Missing-data imputation: KVAE smoothing improves missing-frame imputation by incorporating information from future observed frames, not only preceding observations.Imputation is evaluated after the first 4 frames are observed and the next 16 are randomly dropped with varying probabilities, using the average fraction of incorrectly guessed pixels.
  • Missing-data imputation: KVAE can be trained with up to 40% missing data with minor performance losses because of its smoothing capabilities.The paper also provides links to imputation and long-term generation results in the supplementary material and online.
  • Dynamics interpretation: The dynamics parameter network correctly selects transitions for constant central motion and velocity reversals near the box’s horizontal and vertical walls.In the illustrated MLP mixture network, different learned dynamics are selected according to the previous latent encoding.
  • Generalization: Across irregular-polygon bouncing, gravity-affected bouncing, and Pong-like environments, KVAE learns the dynamics and generates realistic-looking trajectories.These experiments test adaptation to environments with different properties.

6 Related work

Related work extends VAEs and state-space models to sequential data, while the KVAE preserves linear state transitions conditioned on latent encodings to enable exact posterior inference. Other lines of research model action-conditioned environments or combine physical world models with deep learning.

  • Probabilistic sequential models: Temporal VAE extensions use deep neural networks to parameterize transition and emission distributions in deep state-space models.These models define variational approximations to the posterior distribution of latent states at each time step.
  • Probabilistic sequential models: Unlike nonlinear transition models, the KVAE preserves linear dependencies between consecutive latent states by conditioning dynamics on latent encodings.This structure permits exact smoothed posterior computation given the encodings, supporting missing data imputation.
  • Action-conditioned dynamics: Reinforcement-learning research models environment evolution under actions, including latent frame representations and LSTM-parameterized action-conditional dynamics.These approaches, like the KVAE, can make long-term predictions without generating high-dimensional images at every time step.
  • Physical models and deep learning: Other recent work combines physical models of the world with deep learning to learn object dynamics in more complex but low-dimensional environments.The present work instead focuses on generative modeling for high-dimensional videos of simple physical systems.

7 Conclusion

The KVAE enables end-to-end unsupervised learning of high-dimensional videos while disentangling object representations from latent states describing dynamics. Exact smoothed-state inference supports improved inference and missing-data imputation, with temporal reasoning performed in latent space as a demonstrated proof of concept.

  • The KVAE disentangles an object’s latent representation a_t from a latent state z_t describing its dynamics and learns end-to-end from raw video.
  • Exact conditional smoothed posteriors over LGSSM states generally improve inference and missing-data imputation.
  • Disentangling representations enables temporal reasoning, and potentially planning, in latent space.
  • The model’s proof of concept focuses on videos of static worlds containing a few moving objects.

A Experimental details

Experiments used simulated Pymunk videos and a convolutional encoder-decoder, with optimization and staged training choices designed to prioritize learning temporal dynamics. LGSSM dimensions and mixture complexity were adjusted across physical-system experiments.

  • Data generation: 5,000 training and 1,000 testing videos were generated with the Pymunk physics engine.
  • Encoder/Decoder architecture for the KVAE: The image-based KVAE used three-layer convolutional encoder and Sub-Pixel decoder networks, each with 32 units per layer, 3x3 kernels, stride 2, and ReLU activations.MLPs were also tested in the pendulum experiment.
  • Optimization: ADAM optimization began at learning rate 0.007, decayed by 0.85 every 20 epochs, and reached convergence in roughly 80 epochs.Each epoch took 55 seconds on an NVIDIA Titan X.
  • Training tricks for end-to-end learning: Scaling the VAE reconstruction term by 0.3 helped training focus on temporal dynamics rather than reconstruction.
  • Training tricks for end-to-end learning: Training first optimized VAE parameters and globally learned matrices before jointly learning all parameters, helping establish embeddings, prior scale, and use of multiple dynamics.
  • Choice of hyperparameters for the LGSSM: Most experiments used a_t ∈ R2, z_t ∈ R4, and K = 3; gravity used z_t ∈ R5, while polygon experiments used K = 7.The larger gravity state accommodated a bias from external gravity, and polygons required more complex dynamics.

B Videos

The video demonstrations initialize models with four frames and sample sequences under filtering, smoothing, or long-generation settings, with masking schemes controlling which observations are available.

  • Generation setup: Video generation initializes all models with 4 frames before sampling.Filtering and smoothing variants may observe part of the sequence depending on the masking scheme.
  • Masking conditions: Filtering and smoothing videos use random masks with an 80% masking probability, except consecutive videos, which observe only the first and last 4 frames.The masking setup follows the schemes illustrated in figures 3a and 3b; only KVAE models have smoothing videos.
  • Visual conventions: Most videos show ground truth as a black ball and model output as a red ball.Videos marked long_generation omit the true sequence.
  • Availability: The videos are available through Google Drive and the KVAE website.The cited website is sites.google.com/view/kvae.

C Training with missing data.

The KVAE can be trained with missing data by using the smoothed posterior and masking missing observations in the ELBO. Although performance slightly degrades with 30% and 40% missing data, smoothing remains more accurate than filtering with all training data available.

  • Training with missing data: The smoothed posterior enables KVAE training with missing data by modifying the ELBO to mask missing data-point contributions.The masking is applied in the joint probability distribution and variational approximation.
  • Training with missing data: The mask variable I_t equals 0 for a missing data point and 1 otherwise.
  • Training with missing data: Performance slightly degrades with 30% and 40% missing data, yet smoothing remains more accurate than filtering with all training data available.

D Dynamics parameter network architecture · E Imputation in all environments

The dynamics parameter network architecture strongly affects modeling performance, with LSTMs outperforming the alternatives tested on bouncing-ball data. Figure 8 reports imputation results across all environments.

  • D Dynamics parameter network architecture: The α-network governs the nonlinear dynamics and significantly affects the model’s capabilities.
  • D Dynamics parameter network architecture: Three α-network architectures were considered: a two-hidden-layer MLP, an LSTM-based recurrent network, and a FIFO MLP accessing five time steps.
  • D Dynamics parameter network architecture: α can also be modeled as an approximate discrete random variable with the Concrete distribution, yielding an approximation to a switching Kalman filter.
  • D Dynamics parameter network architecture: On bouncing-ball data, all alternative architectures performed worse than the LSTM used in the other experiments.
  • D Dynamics parameter network architecture: LSTMs were believed to better model discretization errors from collisions and the 32x32 rendering of physics-engine trajectories.
  • E Imputation in all environments: Figure 8 presents imputation results for all environments.
Loading 1710.05741v2…