Variational Inference

Mathematics of Evidence Lower Bound (ELBO) objective formulation is inevitably encountered when dealing with latent dynamics models and VAEs. In this blog, we break it down along with necessary concepts like the reparameterization trick, importance sampling and Jensen's inequality.

Prior Distribution

Latent variable models aim to model observed data $o_{1:T}$ while leveraging hidden (latent) variables $s_{1:T}$​.

Joint Distribution

In these models, the joint distribution can be defined as: $$\begin{array}{}\text{(1)}& p(o_{1:T},s_{1:T}\mid a_{1:T})=\prod _{t=1}^{T}p(o_t\mid s_t)p(s_t\mid s_{t-1},a_{t-1})\end{array}$$ Where $p(o_t\mid s_t)$ is the likelihood of observing $o_t$ given the latent state $s_t$​ and $p(s_t\mid s_{t-1},a_{t-1})$ is transition dynamics of latent states.

Marginal Likelihood

To compute the marginal likelihood, we marginalize out $s_{1:T}$ as the following equation: $$\begin{array}{}\text{(2)}& p(o_{1:T}\mid a_{1:T})=\int p(o_{1:T},s_{1:T}\mid a_{1:T})d{s}_{1:T}\end{array}$$ This integral is typically intractable due to high dimensionality and complex dependencies between variables. So we instead approximate this intractable data likelihood distribution with a tractable posterior approximation.

Importance sampling

This statistical technique used to approximate expectations or probabilities by sampling from an easier-to-handle distribution (called the importance distribution) q(x) rather than the true, often intractable, distribution p(x). While sampling q(x) reweight the samples to account for the difference between q(x) and p(x), so it also known as importance weighting. $$\begin{array}{}\text{(1)}& {\mathbb{E}}_{p(x)}[f(x)]=\int f(x)p(x)dx\end{array}$$ Rewriting using q(x), we multiply and divide by q(x): $$\begin{array}{}\text{(2)}& {\mathbb{E}}_{p(x)}[f(x)]=\int f(x)\frac{p(x)}{q(x)}q(x)dx={\mathbb{E}}_{q(x)}[f(x)\frac{p(x)}{q(x)}]\end{array}$$ Where w(x)= $\frac{p(x)}{q(x)}$​ are the importance weights.

Variational Posterior q

Variational posterior approximating the true posterior distribution $p(s_{1:T}\mid o_{1:T},a_{1:T})$ is expressed as: $$\begin{array}{}\text{(3)}& q(s_{1:T}\mid o_{1:T},a_{1:T})=\prod _{t=1}^{T}q(s_t\mid o_t,a_{<t})\end{array}$$

Tractability

Why q is easier to tract?

This factorization assumes that each latent variable $s_t$ depends only on local observations $o_t$​ and preceding actions $a_{<t}$​, rather than the entire sequence of states. In the marginal likelihood or prior distribution, as in equation ($\text{2}$), one has to do the integral over all possible state values (probabilities of all states leading to an observation from $o_{1:T}$), which is infinite as the state value is continuous.

Why p is difficult to tract?

P when $s_t$ is continuous

Equation ($\text{2}$) involves integrating over $s_{1:T}$​, which consists of the T latent variables $s_1,s_2,\dots ,s_T$​. $$\begin{array}{}\text{(4)}& \int d{s}_{1:T}={\int }_{{\mathbb{R}}^n}d{s}_1{\int }_{{\mathbb{R}}^n}d{s}_2\cdots {\int }_{{\mathbb{R}}^n}d{s}_T\end{array}$$ The sequence $s_{1:T}$​ has a finite length T, but The integral aggregates over all possible continuous values of $s_1,s_2,\dots$ leading to a continuous T-fold integral.

P when $s_t$ is discrete

If $s_t$​ takes on discrete values $s_t^{(1)},s_t^{(2)},\dots ,s_t^{(K)}$, where K is the number of discrete states for $s_t$​, the marginal likelihood becomes: $$\begin{array}{}\text{(5)}& .\sum _{s_{1:T}}=\sum _{s_1}\sum _{s_2}\cdots \sum _{s_T}\end{array}$$ Each summation is over all possible values of $s_t$​. For categorical $s_t$​ might be with K states and T variables, we need to do a sum over $K^T$ combinations of $s_{1:T}$. Note: In discrete models (e.g., Hidden Markov Models), the marginal likelihood summation can often be computed efficiently using algorithms like the forward algorithm, avoiding an explicit sum over all $K^T$ possibilities.

This is why

Marginal P is computationally intractable as exact integration over continuous spaces or $K^T$ summations over high dimensional discrete state with K values is rarely feasible. Computation of marginal likelihood using analytical methods is not possible, numerical methods are needed. Note: This is why techniques like variational inference or Monte Carlo sampling are used to approximate the marginal likelihood. Before proceeding with understanding expressions of ELBO, apart from about importance sampling, we will also need to know Jensen's equality.

Jensen’s Inequality

Jensen's inequality states that for a convex (or concave) function f(x), and a random variable X: $$\begin{array}{}\text{(6)}& f(\mathbb{E}[X])\le \mathbb{E}\end{array}$$ If f(x) is concave, the inequality is flipped: $$\begin{array}{}\text{(7)}& \ln \mathbb{E}[x]\ge \mathbb{E}[\ln x]\end{array}$$

Deriving the Variational Bound ELBO

In equation ($\text{2}$), we introduce a variational approximation $q(s_{1:T}\mid o_{1:T},a_{1:T})$ to the true posterior $p(s_{1:T}\mid o_{1:T},a_{1:T})$ using importance sampling: $$\begin{array}{}\text{(8)}& p(o_{1:T}\mid a_{1:T})=\int q(s_{1:T}\mid o_{1:T},a_{1:T})\frac{p(o_{1:T},s_{1:T}\mid a_{1:T})}{q(s_{1:T}\mid o_{1:T},a_{1:T})}d{s}_{1:T}\end{array}$$ Taking the logarithm: $$\begin{array}{}\text{(9)}& \ln p(o_{1:T}\mid a_{1:T})=\ln \int q(s_{1:T}\mid o_{1:T},a_{1:T})\frac{p(o_{1:T},s_{1:T}\mid a_{1:T})}{q(s_{1:T}\mid o_{1:T},a_{1:T})}d{s}_{1:T}\end{array}$$ Since the logarithm is concave, we can apply Jensen's inequality as: $$\begin{array}{}\text{(10)}& \ln (\mathbb{E}[X])\ge \mathbb{E}[\ln (X)]\end{array}$$ so that $$\begin{array}{}\text{(11)}& \ln p(o_{1:T}\mid a_{1:T})\ge \int q(s_{1:T}\mid o_{1:T},a_{1:T})\ln \frac{p(o_{1:T},s_{1:T}\mid a_{1:T})}{q(s_{1:T}\mid o_{1:T},a_{1:T})}d{s}_{1:T}\end{array}$$ This forms the Evidence Lower Bound (ELBO): $$\begin{array}{}\text{(12)}& \ln p(o_{1:T}\mid a_{1:T})\ge {\mathbb{E}}_{q(s_{1:T}\mid o_{1:T},a_{1:T})}[\ln p(o_{1:T},s_{1:T}\mid a_{1:T})-\ln q(s_{1:T}\mid o_{1:T},a_{1:T})]\end{array}$$ Equation of joint likelihood: $$\begin{array}{}\text{(13)}& p(o_{1:T},s_{1:T}\mid a_{1:T})=\prod _{t=1}^{T}p(o_t\mid s_t)p(s_t\mid s_{t-1},a_{t-1})\end{array}$$ Using equation ($\text{1}$) of joint likelihood into this ELBO: $$\begin{array}{}\text{(14)}& \ln p(o_{1:T}\mid a_{1:T})\ge {\mathbb{E}}_{q(s_{1:T}\mid o_{1:T},a_{1:T})}[\sum _{t=1}^T\ln p(o_t\mid s_t)+\ln p(s_t\mid s_{t-1},a_{t-1})-\ln q(s_t\mid o_t,a_{<t})].\end{array}$$

ELBO Components

We can now split the ELBO into two main components: Reconstruction Term: measures how well the model reconstructs the observations given the latent states and KL Divergence measures how much the variational posterior $q(s_t\mid o_t,a_{tt})$ deviates from the prior dynamics $p(s_t\mid s_{t-1},a_{t-1})$. Final ELBO can be written as: $$\begin{array}{}\text{(15)}& \mathcal{L}=\underset{\text{Reconstruction Loss}}{\underbrace{\sum _{t=1}^T{\mathbb{E}}_{q(s_t\mid o_t)}[\ln p(o_t\mid s_t)]}}-\underset{\text{Complexity}}{\underbrace{\sum _{t=1}^T{D}_{\mathrm{KL}}(q(s_t\mid o_t,a_{<t})\parallel p(s_t\mid s_{t-1},a_{t-1}))}}\end{array}$$ The loss function becomes: $$\begin{array}{}\text{(16)}& \mathcal{L}=-\text{ELBO}=\text{Reconstruction Loss}+\text{Complexity Penalty}\end{array}$$

Non-differentiability Problem: Stochastic Sampling from variational posterior

While maximizing the ELBO if q is a complex function of parameters (e.g., neural network outputs), gradients for optimization can become problematic due to the non-differentiable nature of sampling operations in q. How? To optimize the parameters $\varphi$ of $q(s_{1:T}\mid o_{1:T},a_{1:T})$, we compute gradients of \mathcal{L} with respect to $\varphi$. The expectation is over $q(s_{1:T}\mid o_{1:T},a_{1:T}))$, meaning we must sample $s_{1:T}$​ to estimate this expectation. But sampling $s_{1:T}\sim q(s_{1:T}\mid o_{1:T},a_{1:T})$ is stochastic and hence non-differentiable - the stochasticity in sampling disrupts the computation of gradients via backpropagation.

Solution: Reparameterization Trick

It re-expresses the stochastic sampling as a deterministic transformation of a fixed random variable. Instead of directly sampling $s_t$​ from $q(s_t\mid o_t,a_{<t})$, It introduces a Noise $ϵ$ by sampling from a fixed distribution (e.g., $\mathcal{N}(0,1)$ - we move the randomness to $ϵ$ that is independent of learnable parameters $\varphi$ making the computation of ${\mathrm{\nabla }}_{\varphi }$ tractable. In the deterministic transformation as below gradients ${\mathrm{\nabla }}_{\varphi }$​ can flow through ${\mu }_{\varphi }$​ and ${\sigma }_{\varphi }$​ via backpropagation: $$\begin{array}{}\text{(17)}& s_t={\mu }_{\varphi }+{\sigma }_{\varphi }\cdot ϵ\sim \mathcal{N}(0,1)\end{array}$$