Variational RNN (VRNN)

The core innovation of a VRNN is the introduction of a stochastic latent variable z_t at each timestep. Unlike a standard RNN that maintains only a deterministic hidden state h_t, the VRNN models the conditional distribution p(z_t | x_≤t, z_<t). This allows the model to capture the inherent uncertainty and multi-modality in sequential data generation. For example, in video prediction, the next frame is not deterministic; the latent variable allows the model to represent plausible future states.

• Inference: The latent variable is inferred using a probabilistic encoder q(z_t | x_≤t, z_<t, h_{t-1}).
• Generation: The latent variable is used by a probabilistic decoder p(x_t | z_t, h_{t-1}) to generate the observation.
• Prior: The model also learns a dynamic prior p(z_t | x_<t, z_<t, h_{t-1}), which is the distribution of z_t before seeing x_t.

The VRNN is trained by maximizing a timestep-wise Evidence Lower Bound (ELBO) on the log-likelihood of the sequence. This objective decomposes the intractable marginal likelihood into a sum of tractable terms per timestep.

The ELBO for a sequence x_1:T is: ELBO = Σ_t [ E_{q(z_t)}[log p(x_t | z_t, h_{t-1})] - D_KL( q(z_t | x_≤t, h_{t-1}) || p(z_t | x_<t, h_{t-1}) ) ]

• Reconstruction Term: The first term encourages the decoder to accurately reconstruct (or predict) the observation x_t given the sampled latent z_t and previous state.
• KL Divergence Term: The second term acts as a regularizer, forcing the approximate posterior (the encoder's output) to stay close to the learned prior. This prior evolves based on the sequence history (h_{t-1}), making it a conditional prior.

A key distinction from a standard VAE is the VRNN's use of a conditional prior p(z_t | h_{t-1}). In a VAE, the prior is typically a static standard normal distribution N(0, I). In a VRNN, the prior at time t is parameterized by the RNN's previous hidden state h_{t-1}. This allows the model's belief about the latent state at time t to be informed by all previous observations and latents.

• Mechanism: A neural network takes h_{t-1} as input and outputs the parameters (e.g., mean and variance) of the Gaussian prior for z_t.
• Role in Training: The KL divergence in the ELBO is between the encoder's posterior and this dynamic prior. This encourages the latent variables to follow a temporally coherent trajectory dictated by the model's accumulated knowledge, rather than just being random.

The deterministic RNN core is the backbone that propagates information through time and conditions all probabilistic components. The update for the hidden state h_t is:

h_t = RNN( [φ_x(x_t), φ_z(z_t)], h_{t-1} )

Where φ_x and φ_z are feature extractor networks.

• Bidirectional Conditioning: The hidden state h_{t-1} conditions the prior and encoder for timestep t. Simultaneously, the new observation x_t and sampled latent z_t update the hidden state to h_t.
• Feature Extraction: The networks φ_x and φ_z transform raw observations and latent samples into feature vectors that the RNN can process effectively. This modularity allows the VRNN to handle diverse data types (e.g., audio spectrograms, image patches, text embeddings).

VRNNs excel at open-ended sequence generation and multi-step prediction. The generation process is ancestral sampling:

1. Initialize h_0.
2. For each timestep t:
   • Sample z_t from the conditional prior: z_t ~ p(z_t | h_{t-1}).
   • Sample a generated observation x_t from the decoder: x_t ~ p(x_t | z_t, h_{t-1}).
   • Update the RNN state: h_t = RNN([φ_x(x_t), φ_z(z_t)], h_{t-1}).

This process can continue indefinitely, generating coherent sequences. For prediction (e.g., forecasting the next 10 frames of video), the model uses the true observations up to time T and then switches to using its own generated samples x_t for t > T to condition future steps, propagating uncertainty through the latent variables.

The VRNN sits at the intersection of several model families:

• Vs. Standard RNNs: RNNs are purely deterministic. VRNNs add stochasticity via latent variables, enabling them to model distributions over sequences and generate diverse outputs.
• Vs. VAEs: A standard VAE models static data (e.g., images). A VRNN can be viewed as a temporal extension of the VAE, applying the VAE's variational framework per timestep with a recurrent state carrying memory.
• Vs. State-Space Models (SSMs): Like SSMs, VRNNs have latent states (z_t) and observations (x_t). However, VRNNs use flexible neural networks (RNNs, MLPs) for transition and emission functions, whereas classic SSMs use linear Gaussian models.
• Predecessor to Modern Models: The VRNN architecture is a conceptual precursor to more advanced models like Stochastic Recurrent Neural Networks (SRNNs) and certain types of Diffusion Models for sequences, which also blend deterministic recurrence with stochastic latent variables.

The Variational Autoencoder (VAE) is the foundational generative model upon which the VRNN is built. It consists of a probabilistic encoder that maps data to a latent distribution and a probabilistic decoder that reconstructs data from latent samples. The model is trained by maximizing the Evidence Lower Bound (ELBO), which balances reconstruction loss with a KL divergence regularization term that pushes the latent distribution toward a simple prior (e.g., a standard normal). The VRNN essentially instantiates a VAE at each timestep of a recurrent sequence.

A Recurrent Neural Network (RNN) is a class of neural networks designed for sequential data, where connections between nodes form a directed graph along a temporal sequence. This allows it to maintain a hidden state that acts as a memory of previous inputs. Core variants include:

• Long Short-Term Memory (LSTM): Uses gating mechanisms (input, forget, output) to control information flow and mitigate the vanishing gradient problem.
• Gated Recurrent Unit (GRU): A simplified LSTM with fewer gates, often faster to train. The VRNN uses an RNN (typically an LSTM or GRU) to model the temporal dependencies between the stochastic latent variables and observed data across the sequence.

The Evidence Lower Bound (ELBO) is the core objective function for training VAEs and, by extension, VRNNs. It is a tractable lower bound on the log-likelihood of the data. For a VRNN, the ELBO is a sum over timesteps, with each term containing:

• A reconstruction term: The expected log-likelihood of the observed data given the latent state.
• A KL divergence term: Measures the divergence between the approximate posterior (from the encoder) and the prior distribution over the latent variable at that step. Maximizing the ELBO simultaneously improves the quality of generated sequences and ensures a well-structured, regularized latent space over time.

A stochastic latent variable is a random, unobserved variable that represents a compressed, probabilistic summary of the data. In a VRNN, a new latent variable z_t is sampled at each timestep t from a distribution (e.g., Gaussian) whose parameters are output by the encoder network. This introduces inherent uncertainty into the model, which is crucial for generating diverse and realistic sequences, as opposed to deterministic RNNs that produce a single output for a given input. The reparameterization trick is used to allow gradient flow through these stochastic sampling steps during training.

A Deep Markov Model (DMM) is a closely related generative model for sequences that also combines latent variables with neural networks. Like the VRNN, it models sequences with a chain of stochastic latent states. The key architectural difference is that in a DMM, the transition between latent states z_{t-1} -> z_t is typically modeled by a simple network (e.g., a multilayer perceptron), while in a VRNN, this transition is mediated by the deterministic RNN hidden state h_t, which accumulates historical context. This gives the VRNN more powerful long-term memory capabilities.

Sequential Variational Autoencoder is a broad term encompassing models that apply variational inference to sequential data, including the VRNN and Deep Markov Model. These models extend the VAE framework by defining a joint distribution over a sequence of observed data x_{1:T} and a corresponding sequence of latent variables z_{1:T}. The training objective is the sequential ELBO. The VRNN is a prominent instance where the inference (encoder) and generative (decoder) networks are conditioned on the RNN's hidden state, enabling rich, history-dependent modeling of both the prior and posterior distributions.