[NeurIPS] Latent Matters: Mastering Dynamics with Extended Kalman VAEs

TL;DR

Deep State-Space Models (DSSMs) often look like they are working (high ELBO) while failing to actually understand the "physics" of the data they process. This paper introduces the Extended Kalman VAE (EKVAE) and a Constrained Optimization (CO) framework. By replacing black-box RNNs with neural-linearized Kalman filters and replacing standard training with a Lagrangian approach, the authors achieve superior system identification and prediction from raw pixels.

Key Insight: To truly learn dynamics, the latent state must be Markovian. If your model uses an RNN, "physics" information like velocity gets trapped in the RNN's hidden state instead of the latent variable, leading to a "non-Markovian" mess that fails at long-term prediction.

The Hidden Failure of ELBO

In the world of Generative Models, we often use the Evidence Lower Bound (ELBO). For sequential data, the ELBO is composed of:

1. Distortion: How well can we reconstruct the image?
2. Rate: How well does our transition model match our posterior?

The authors point out a critical flaw: High ELBO $eq$ Good Dynamics. You can achieve a high ELBO by simply having a great decoder that ignores the transition model. Standard "annealing" (slowly increasing the weight of the Rate) is a fickle heuristic that often fails to find the right balance.

Figure 1: The CO framework in action. By balancing distortion and rate through a Lagrange multiplier $\lambda$, the model first learns to see (reconstruction) and then learns to move (dynamics).

Methodology: The Extended Kalman VAE (EKVAE)

The EKVAE departs from the "RNN-as-a-filter" trend. Instead, it resurrects classic Control Theory using two main pillars:

1. Neural Linearization

Traditional Extended Kalman Filters (EKF) require a Taylor expansion (Jacobian) of the transition function. In deep learning, the transition function is unknown. The EKVAE solves this by directly learning the Jacobian as a function of the current state: ${\mathbf{z}}_{t+1}\approx {\mathbf{F}}_{\psi }({\mathbf{z}}_t,{\mathbf{u}}_t){\mathbf{z}}_t+{\mathbf{B}}_{\psi }({\mathbf{z}}_t,{\mathbf{u}}_t){\mathbf{u}}_t$ This allows the model to compute closed-form Bayesian smoothing updates while remaining fully differentiable and non-linear.

2. Disentangling Static and Dynamic Features

By using a time-invariant auxiliary-variable model $p({\mathbf{a}}_t|{\mathbf{z}}_t)$, the EKVAE forces the latent space to partition itself.

• Static features (at): Captured through the encoder-decoder (e.g., the position of a pendulum).
• Dynamic features: Captured via the transition model (e.g., angular velocity).

Figure 2: The EKVAE Graphical Model. Red arrows denote inference; black arrows denote generation. Note the clean separation of auxiliary variables.

Experiments: Why RNNs Limit System Identification

The authors compared the EKVAE against heavyweights like the Deep Kalman Filter (DKF) and KVAE.

A startling discovery was made regarding Non-Markovian Latents. In models like the KVAE or RSSM, the transition depends on an internal RNN state ($h_t$). Consequently, the latent $z_t$ doesn't actually contain the velocity. It's stored in the RNN memory. When you try to "smooth" or predict long-term, the model falls apart because the latent variables alone are insufficient to describe the system state.

Table 1: On the Pendulum task, the EKVAE identify dynamics almost perfectly, while standard DKS fails significantly.

Deep Insight: Policy Learning Without Rewards

Because the EKVAE disentangles position and velocity, we can perform Model-Based RL without external rewards. By simply defining a target position in the "static" dimensions of the latent code, or a target speed in the "dynamic" dimensions, we can derive a reward function directly from the latent distance.

Figure 3: Robot policies learned purely from latent representations. Top: Swing-up. Middle/Bottom: Steady rotation at specific speeds.

Conclusion

"Latent Matters" is a call to return to the fundamentals of state-space modeling. By ensuring the latent space is truly Markovian and using a constrained optimization framework to balance the objective, we can build AI that doesn't just "parrot" pixels but understands the underlying physics of the world.