Connections between State-Space Models and VAEs

When we model sequences or time-dependent phenomena with Variational Autoencoders, such as with Recurrent VAEs (RVAEs) or VAEs incorporating attention mechanisms, we are implicitly or explicitly defining a system that evolves over time. This naturally brings us to the well-established domain of State-Space Models (SSMs), which have a long history in fields like control engineering, econometrics, and signal processing. Understanding the links between these two modeling paradigms can provide deeper insights into how sequential VAEs function and can guide the development of more principled architectures.

State-Space Models: A Quick Refresher

At their core, SSMs describe a system using a set of unobserved (latent) state variables $z_t$ that evolve over time, and a set of observed variables $x_t$ that depend on the current latent state. A typical discrete-time SSM is defined by two equations:

1. State Transition Equation: This describes how the latent state evolves from one time step to the next, often influenced by a control input $u_t$ (if present) and process noise $w_t$. Probabilistically, this is $p(z_t | z_{t-1}, u_t)$. $$z_t = f(z_{t-1}, u_t) + w_t$$
2. Observation Equation: This defines how the observed data $x_t$ is generated from the current latent state $z_t$, potentially influenced by $u_t$ and affected by measurement noise $v_t$. Probabilistically, this is $p(x_t | z_t, u_t)$. $$x_t = g(z_t, u_t) + v_t$$

The functions $f$ and $g$ define the dynamics and observation process, respectively. In classical linear-Gaussian SSMs, such as those handled by the Kalman filter, $f$ and $g$ are linear functions, and $w_t$ and $v_t$ are assumed to be Gaussian noise. These models allow for exact inference of the latent states $p(z_t | x_{1:T})$ (smoothing) or $p(z_t | x_{1:t})$ (filtering) using efficient algorithms.

Sequential VAEs as Non-Linear State-Space Models

Sequential VAEs can be viewed as a powerful generalization of SSMs, particularly non-linear SSMs. Let's draw the parallels:

• Latent States $z_t$: The sequence of latent variables $z_1, z_2, \ldots, z_T$ in a VAE directly corresponds to the hidden states in an SSM. These variables aim to capture the underlying dynamical structure of the data.
• Transition Model $p_\theta(z_t | z_{t-1})$: The prior distribution over the latent sequence in a VAE, often defined by a recurrent neural network (RNN) or a simpler Markovian dependency $p_\theta(z_t | z_{t-1})$, acts as the state transition model. This network learns the dynamics of how latent states evolve.
• Emission Model $p_\theta(x_t | z_t)$: The VAE's decoder network, $p_\theta(x_t | z_t)$, functions as the observation or emission model. It learns to generate the observed data point $x_t$ given the current latent state $z_t$.

The generative process of a sequential VAE often follows this SSM-like structure:

1. Sample initial latent state $z_1 \sim p_\theta(z_1)$.
2. For $t = 2, \ldots, T$: Sample $z_t \sim p_\theta(z_t | z_{t-1})$ using the prior network (transition).
3. For $t = 1, \ldots, T$: Sample $x_t \sim p_\theta(x_t | z_t)$ using the decoder network (emission).

The following diagram illustrates the structural similarities in their generative paths:

This diagram highlights the parallel components in the generative process of a State-Space Model and a sequential Variational Autoencoder. Both rely on a latent variable at time $t-1$ to inform the latent variable at time $t$, which then produces an observation.

Key Distinctions and Shared Aspects

While the structural analogy is strong, there are important differences, primarily stemming from the use of neural networks and variational inference in VAEs:

• Non-Linearity and Expressiveness:

  • SSMs: Classical SSMs often assume linear dynamics ($f$) and linear observation models ($g$), with Gaussian noise. While non-linear SSMs exist (e.g., Extended Kalman Filter, Unscented Kalman Filter, Particle Filters for inference), defining appropriate non-linear functions can be challenging.
  • VAEs: The transition (prior network) and emission (decoder network) functions in VAEs are typically deep neural networks. This allows VAEs to model highly complex, non-linear dynamics and observation processes with much greater flexibility than traditional parametric SSMs. They can also learn arbitrary conditional distributions, not just Gaussians.

• Inference:

  • SSMs: For linear-Gaussian SSMs, Kalman filtering and smoothing provide exact and efficient inference of $p(z_{1:T} | x_{1:T})$. For non-linear/non-Gaussian SSMs, inference is approximate and can be computationally intensive (e.g., particle filters).
  • VAEs: VAEs employ amortized variational inference. An encoder network $q_\phi(z_{1:T} | x_{1:T})$ (often factorized as $q_\phi(z_{1:T} | x_{1:T}) = \prod_t q_\phi(z_t | z_{<t}, x_{\le t})$ or similar) is trained to approximate the true posterior over latent states. This encoder learns to map observed sequences $x_{1:T}$ to distributions over latent sequences $z_{1:T}$. This process is akin to learning an approximate filtering or smoothing distribution.

• Learning:

  • SSMs: Parameters (e.g., matrices in linear SSMs, noise covariances) are often estimated using methods like the Expectation-Maximization (EM) algorithm or by maximizing the likelihood directly if tractable.
  • VAEs: All parameters of the encoder, decoder, and prior networks are learned jointly by maximizing the Evidence Lower Bound (ELBO) via stochastic gradient descent. This end-to-end learning is a hallmark of deep generative models.

Examples and Model Variants

Several VAE architectures for sequential data explicitly or implicitly embody SSM principles:

• Variational Recurrent Neural Network (VRNN): Integrates VAE principles within an RNN. At each time step, the RNN state $h_t$ influences the prior for $z_t$. The latent $z_t$ and observation $x_t$ then update $h_t$. This creates a dynamic system where latent variables guide sequence generation.

  • $z_t \sim \text{prior}(z_t | h_{t-1})$
  • $x_t \sim \text{decoder}(x_t | z_t, h_{t-1})$
  • h_t = \text{RNN_cell}(x_t, z_t, h_{t-1}) The encoder also uses $h_{t-1}$ and $x_t$ to approximate $p(z_t | x_{\le t}, h_{<t})$.

• Deep Kalman Filters (DKFs) and Kalman VAEs (KVAEs): These models attempt to combine the structured probabilistic inference of Kalman filters with the expressive power of neural networks. For instance, a DKF might assume linear Gaussian transitions but use a neural network for the emission model:

  $$p(z_t | z_{t-1}) = \mathcal{N}(z_t | A z_{t-1}, Q)$$ $$p_\theta(x_t | z_t) = \mathcal{N}(x_t | \text{NN}_\theta(z_t), R)$$

  Inference in such models can still be challenging, and VAE-based approaches (like KVAEs) use variational methods to approximate the posterior over the structured latent states.

• Deep Markov Models (DMMs): This term is often used for sequential VAEs where the latent states $z_t$ are assumed to follow a first-order Markov process, $p(z_t | z_{t-1})$, typically parameterized by a neural network.

Implications of the SSM Connection

Recognizing sequential VAEs as sophisticated SSMs offers several benefits:

1. Structured Understanding: Viewing VAEs through the SSM lens provides a more structured way to reason about the learned latent dynamics. It helps in thinking about what aspects of the temporal process the VAE is trying to capture.
2. Principled Architecture Design: Knowledge from SSM theory can inspire new VAE architectures. For example, incorporating specific types of recurrent structures in the prior or encoder that mimic known dynamical systems or imposing certain structures on the latent space.
3. Tool for Analysis: Techniques from SSM analysis, such as examining state transitions or long-term dependencies, might be adapted to understand the behavior of trained sequential VAEs.
4. Bridging Paradigms: This connection allows for cross-pollination of ideas. Concepts like controllability and observability from SSMs might find new interpretations or analogous principles in the VAE context.

However, the high dimensionality and non-linear nature of latent spaces in VAEs mean that direct application of classical SSM analysis tools can be difficult. The interpretability of learned dynamics in complex VAEs remains an active area of research.

In summary, the relationship between VAEs for temporal data and State-Space Models is profound. VAEs extend the SSM framework by incorporating powerful non-linear function approximators (neural networks) and scalable inference techniques (amortized variational inference). This allows them to model significantly more complex sequential data than traditional SSMs, while the SSM perspective provides a valuable framework for understanding and designing these advanced generative models. As VAEs continue to evolve, this connection will likely inspire further innovations in modeling dynamic systems.