A Comparison of Probabilistic and Statistic View of Deep Learning Models
Published:
Background
We use lower case $x$ to denote the observed sample, and upper case $X$ to denote the associated unknown variable. We denote by $X, Y, Z$ the observation, label, and latent, respectively. In most cases, only the variable $X$ can be fully observed, the label $Y$ might be partially observed, and the latent $Z$ is unobservable.
Evidence Lower Bound (ELBO)
We have an observed variable $X$ and one latent variable $Z$, the joint distribution $p(X, Z)$ can be parameterized by
\[p(X, Z) = p (Z) p (X|Z) \approx \theta (Z) \theta (X|Z), \text{ where } \theta (Z) = \mathcal{N} (Z; 0, I), \; \theta (X|Z) = \mathcal{N} (X; \mu_{\theta}, \Sigma_{\theta})\]or
\[p(X, Z) = p (X) p (Z \vert X) \approx \phi (X) \phi (Z \vert X), \text{ where } \phi (Z \vert X) = \mathcal{N} (Z; \mu_{\phi}, \Sigma_{\phi})\]Now we have three probabilistic worlds:
- Real world: $p(X)$;
- Observed world: $\hat{p} (X)$;
- Approximated world: $\phi (X), \theta(Z), \theta(X\vert Z), \phi(Z\vert X)$.
The true world is never observable, and the approximated world is derived from our prior and posterior knowledge. Therefore, the goal is to make our knowledge as close to the ground truth as possible, i.e., minimize the probability discrepancy between $\hat{p} (x)$ and $p(x)$. In the cases where computing the evidence $\log p (x)$ is intractable (due to mixture of distributions), we would rather use the marginal
\[\log p(x) = \log \int p (x, z) \text{d} z = \log \int \theta (x \vert z) \theta (z) \text{d} z\]If the integral over $z$ is tractable, then EM algorithm can be applied. Otherwise, we need to optimize its lower bound, called evidence lower bound. For any observed sample $x$, the log likelihood is lower bounded by
\[\log p (x) \ge \mathbb{E}_{Z \sim \phi (\cdot \vert x)} \bigg[\log \frac{p (x, Z)}{\phi (Z \vert x)}\bigg] = ELBO (x)\]Actually, $\log p (x) = ELBO (x) + \mathcal{D}_{KL} (\phi (\cdot\vert x) \vert \vert p(\cdot\vert x))$, which means we must pay a price if we use an approximation of the posterior $p(\cdot\vert x)$. Furthermore, ELBO can be written as the difference between a reconstruction error and a prior matching error:
\[ELBO (x) = \underbrace{\mathbb{E}_{Z \sim \phi (\cdot \vert x)} [\log \theta (x \vert Z)]}_{\text{reconstruction error}} - \underbrace{\mathcal{D}_{KL} (\phi (\cdot \vert x) \vert\vert \theta (\cdot))}_{\text{prior matching error}}\]On the other hand,
\[ELBO (x) = \underbrace{\log p(x)}_{\text{likelihood}} - \underbrace{\mathcal{D}_{KL} (\phi (\cdot \vert x) \vert p (\cdot \vert x))}_{\text{posterior matching error}}\]We see that, by maximizing ELBO, we simultaneously minimize the KL-divergence between pair $(\phi(\cdot\vert x), \theta(\cdot))$ and $(\phi (\cdot\vert x), p(\cdot\vert x))$. The former one ensures that $\phi$ is indeed learning some distribution, and the latter one ensures that $\phi$ is getting close to the real distribution. In practice, since both $\phi (\cdot \vert x)$ and $\theta (\cdot)$ are assumed to be Gaussian, the KL divergence between them can be computed in closed form. The reconstruction error is computed via Monte-Carlo sampling:
\[\mathbb{E}_{Z \sim \phi (\cdot \vert x)} [\log \theta (x \vert Z)] \approx \frac{1}{N} \sum_{i = 1}^N \log \theta (x \vert z_i).\]Interactive illustration: the two faces of the ELBO
Consider the toy model $Z \sim \mathcal{N} (0, 1)$, $X \vert Z \sim \mathcal{N} (Z, \sigma^2)$, for which everything is closed form: the evidence is $p (x) = \mathcal{N} (x; 0, 1 + \sigma^2)$ and the true posterior is $p (z \vert x) = \mathcal{N} \big(\frac{x}{1 + \sigma^2}, \frac{\sigma^2}{1 + \sigma^2}\big)$. Choose the approximate posterior $\phi (z \vert x) = \mathcal{N} (m, s^2)$ by hand. The bars verify both decompositions at once:
\[ELBO = \text{reconstruction} - \mathcal{D}_{KL} (\phi \vert\vert \text{prior}) = \log p (x) - \mathcal{D}_{KL} (\phi \vert\vert p (\cdot \vert x)).\]Pushing the ELBO up is the same as shrinking the posterior gap; it reaches $\log p (x)$ exactly when $\phi$ coincides with the true posterior.
Density Estimation
With observation $X$, we minimize the KL-divergence between $\hat{p} (X)$ and $\phi(X)$:
\[\min KL (\hat{p} \vert \phi) = \max \mathbb{E}_{X \sim \hat{p}} [\log \phi(X)]\]Classification/Regression
With observation $X, Y$, we minimize the KL-divergence between $\hat{p} (Y\vert X)$ and $\phi(Y\vert X)$:
\[\min KL (\hat{p} \vert \phi) = \max \mathbb{E}_{(X, Y) \sim \hat{p}} [\log \phi(Y \vert X)]\]Variational Autoencoder (VAE)
With observation $X$ associated with latent variable $Z$, we minimize the KL-divergence between $\hat{p} (X)$ and $\phi(X)$:
\[\begin{align*} & \min KL (\hat{p} \vert \phi) \\ = & \max \mathbb{E}_{X \sim \hat{p}} [\log \phi(X)] \\ \ge & \max \mathbb{E}_{X \sim \hat{p}} [ELBO (X)] \\ = & \max \mathbb{E}_{X \sim \hat{p}} \big[\mathbb{E}_{Z \sim \phi (\cdot \vert X)} [\log \theta (X \vert Z)] - \mathcal{D}_{KL} (\phi (\cdot \vert X) \vert\vert \theta (\cdot))\big] \end{align*}\]Markovian Hierarchical VAE (MHVAE)
With observation $X$ associated with latent variables $Z_{1:T}$, we minimize the KL-divergence between $\hat{p} (X)$ and $\phi(X)$:
\[\begin{align*} & \min KL (\hat{p} \vert \phi) \\ = & \max \mathbb{E}_{X \sim \hat{p}} [\log \phi(X)] \\ \ge & \max \mathbb{E}_{X \sim \hat{p}} [ELBO (X)] \\ = & \max \mathbb{E}_{X \sim \hat{p}} \bigg[\mathbb{E}_{Z_{1:T} \sim \phi (\cdot \vert X)} \bigg[\log \frac{p (X, Z_{1:T})}{\phi (Z_{1:T} \vert X)}\bigg]\bigg] \\ = & \max \mathbb{E}_{X \sim \hat{p}} \bigg[\mathbb{E}_{Z_{1:T} \sim \phi (\cdot \vert X)} \bigg[\log \frac{\theta (Z_T) \theta (X \vert Z_1) \prod_{t = 2}^T \theta (Z_{t-1} \vert Z_t)}{\phi (Z_1 \vert X) \prod_{t = 2}^T \phi (Z_t \vert Z_{t-1})}\bigg]\bigg] \end{align*}\]Variational Diffusion Model (VDM)
With observation $X_0$ associated with latent variables $X_{1:T}$, we additionally know (pre-define) the forward encoder conditional distribution $p (X_t \vert X_{t-1})$. We minimize the KL-divergence between $\hat{p} (X_0)$ and $\phi(X_0)$:
\[\begin{align*} & \min KL (\hat{p} \vert \phi) \\ = & \max \mathbb{E}_{X_0 \sim \hat{p}} [\log \phi(X_0)] \\ \ge & \max \mathbb{E}_{X_0 \sim \hat{p}} [ELBO (X_0)] \\ = & \max \mathbb{E}_{X_0 \sim \hat{p}} \bigg[\mathbb{E}_{X_{1:T} \sim \phi (\cdot \vert X_0)} \bigg[\log \frac{p (X_{0:T})}{\phi (X_{1:T} \vert X_0)}\bigg]\bigg] \\ = & \max \mathbb{E}_{X_0 \sim \hat{p}} \bigg[\mathbb{E}_{X_{1:T} \sim \phi (\cdot \vert X_0)} \bigg[\log \frac{\theta (X_T) \prod_{t = 1}^T \theta (X_{t-1} \vert X_t)}{\prod_{t = 1}^T p (X_t \vert X_{t-1})}\bigg]\bigg] \end{align*}\]The ELBO of VDM has two different decompositions. The first one is
\[\begin{align*} ELBO (x_0) = & \underbrace{\mathbb{E}_{X_1 \sim p (\cdot \vert x_0)} [\log \theta (x_0 \vert X_1)]}_{\text{reconstruction error}} - \underbrace{\mathbb{E}_{X_{T-1} \sim p (\cdot \vert x_0)} [\mathcal{D}_{KL} (p (\cdot \vert X_{T-1}) \vert\vert \theta)]}_{\text{prior matching error of } X_T} \\ & - \sum_{t = 1}^{T-1} \underbrace{\mathbb{E}_{(X_{t-1}, X_{t+1}) \sim p (\cdot \vert x_0)} [\mathcal{D}_{KL} (p (\cdot \vert X_{t-1}) \vert\vert \theta (\cdot \vert X_{t+1}))]}_{\text{consistency error of } X_t} \end{align*}\]The second one is
\[\begin{align*} ELBO (x_0) = & \underbrace{\mathbb{E}_{X_1 \sim p (\cdot \vert x_0)} [\log \theta (x_0 \vert X_1)]}_{\text{reconstruction error}} - \underbrace{\mathcal{D}_{KL} (p (\cdot \vert x_0) \vert\vert \theta)}_{\text{prior matching error of } X_T} \\ & - \sum_{t = 2}^{T} \underbrace{\mathbb{E}_{X_t \sim p (\cdot \vert x_0)} [\mathcal{D}_{KL} (p (\cdot \vert X_{t}, x_0) \vert\vert \theta (\cdot \vert X_{t}))]}_{\text{denoising matching error of } X_{t-1}} \end{align*}\]Understanding Denoising Matching Error
The denoising matching error of $X_{t-1}$ can be interpreted using three different formulations:
\[\begin{align*} \mathcal{D}_{KL} (p (\cdot | X_{t}, x_0) \vert \theta (\cdot | X_{t})) = \begin{cases} C_1 \cdot \vert\vert\hat{x}_{\theta} (X_t, t) - x_0\vert\vert_2^2, & \text{matching the input}; \\ C_2 \cdot \vert\vert\hat{\varepsilon}_{\theta} (X_t, t) - \varepsilon_0\vert\vert_2^2, & \text{matching the noise}; \\ C_3 \cdot \vert\vert\hat{s}_{\theta} (X_t, t) - \nabla \log p (X_t)\vert\vert_2^2, & \text{matching the score}. \end{cases} \end{align*}\]Diffusion Model from the Physical Point of View
The forward process in diffusion models can be considered as the Ornstein-Uhlenbeck process, which is described by the Stochastic Differential Equation (SDE):
\[\text{d} X_t = - \frac{1}{2} g(t) X_t \text{d} t + \sqrt{g (t)} \text{d} W_t, \text{ for } g(t) > 0,\]where $(W_t)_{t \ge 0}$ is a standard Wiener process, and $g(t)$ is a non-decreasing weighting function. Given initial $X_0$, the conditional distribution of $X_t$ is
\[X_t \vert X_0 \sim \mathcal{N} (\alpha (t) X_0, h(t) I),\]where $\alpha (t) = \exp (- \frac{1}{2} \int_0^t g (s) \text{d} s)$ and $h(t) = 1 - \alpha^2 (t)$. Under mild conditions, the initial $X_0$ will be transformed to Gaussian, i.e., $X_{\infty} \vert X_0 \sim \mathcal{N} (0, I)$. The forward process will terminate at a sufficiently large $T > 0$, and the backward process is described by
\[\text{d} X_t = \bigg(\frac{1}{2} X_t + \nabla \log p_{T-t} (X_t)\bigg) \text{d} t + \text{d} W_t\](here $g \equiv 1$ for notational simplicity; for general $g$ the drift becomes $\frac{1}{2} g X_t + g \nabla \log p_{T-t} (X_t)$ and the noise $\sqrt{g} \, \text{d} W_t$). The only unknown ingredient is the score $\nabla \log p_t$, which is exactly what the denoising matching error above trains $\hat{s}{\theta}$ to approximate. There is also a deterministic alternative with the same marginals, the probability-flow ODE $\text{d} X_t = \big(\frac{1}{2} X_t + \frac{1}{2} \nabla \log p{T-t} (X_t)\big) \text{d} t$.
Interactive illustration: forward and reverse diffusion
Each point is a sample from a bimodal data distribution (two clusters). Here $g (t) = 10 t$, so $\alpha (t) = e^{-2.5 t^2}$. Drag the time slider to scrub the forward process $X_t = \alpha(t) X_0 + \sqrt{1 - \alpha^2(t)}\, \varepsilon$ (a fixed $\varepsilon$ per point, so you can go back and forth): at large $t$ the data is transported to an isotropic Gaussian.
“Generate” runs the reverse SDE starting from pure noise at $t = 1$. Because the data is a two-component Gaussian mixture, the score $\nabla \log p_t$ is available in closed form, so the demo uses the exact score and no neural network. Hollow points are the generated samples; they should recover the two clusters. Switch the sampler to the deterministic probability-flow ODE, or add score error to mimic an imperfect learned model and watch the clusters blur.
Diffusion Model for Optimization [Li et al. 2024]
Input: Datasets $\mathcal{T}, \mathcal{T}_{label}$, target reward value $a$, early-stopping time $t_0$, noise level $\nu$.
Reward learning: $\hat{f} \in \arg\min_{f \in \mathcal{F}} L (f, \mathcal{T}_{label})$.
Pseudo labeling:
- Conditional score matching:
- Conditional generation: sample from the backward SDE,
- Return: generated population $\mathbf{Pr} (\cdot \vert \hat{y} = a)$, learned subspace representation $V$ contained in $\hat{s}$.