Exercises — ELBO objective and KL term

(b) is the ELBO. The KL term is subtracted, so it is the penalty. Reason: we want the encoder blob $q_{\varphi }(z|x)$ to stay close to the prior $p(z)$; the larger the KL (the more it drifts), the more we shave off the objective. Choice (a) would reward drifting, which is backwards. Answer ::: (b), and $D_{KL}(q\Vert p(z))$ is the subtracted penalty.

(i) True — Gibbs' inequality; KL is a "wasted-bits" count, never negative. (ii) False — KL is not symmetric. VAEs use the reverse KL $D_{KL}(q\Vert p(z))$ because we take expectations under $q$, the distribution we can sample from. (iii) True — that is why it is called a lower bound. The parent note proves $\log p_{\theta }(x)=\mathcal{L}+D_{KL}(q_{\varphi }(z|x)\Vert p_{\theta }(z|x))$, and the extra KL is $\ge 0$.

Per dimension: ${\mu }_j^2+{\sigma }_j^2-\ln {\sigma }_j^2-1=0+1-\ln 1-1=0$. Sum over 2 dims and halve: $\frac{1}{2}(0+0)=0$. Value ::: $0$ This is the only configuration giving KL $=0$: the encoder blob equals the prior exactly.

Dim 1: $0.5^2+0.8^2-\ln (0.8^2)-1=0.25+0.64+0.4463-1=0.3363$. (Here $-\ln (0.8^2)=-2\ln 0.8=+0.4463$.) Dim 2: $(-0.3{)}^2+1.2^2-\ln (1.2^2)-1=0.09+1.44-0.3646-1=0.1654$. Sum $=0.5017$, halve $\Rightarrow 0.2509$. Value ::: $\approx 0.2509$

Why squared error? For a unit-variance Gaussian, $\log p_{\theta }(x|z)=\log \mathcal{N}(x;{\mu }_{\theta }(z),I)=-\frac{d}{2}\ln (2\pi )-\frac{1}{2}\Vert x-{\mu }_{\theta }(z){\Vert }^2$. The first piece is a constant (does not depend on $\theta ,\varphi$), so we drop it, leaving $-\frac{1}{2}\Vert x-{\mu }_{\theta }(z){\Vert }^2$ — plain negative squared reconstruction error. $\Vert x-{\mu }_{\theta }{\Vert }^2=(0.2-0.19{)}^2+(0.8-0.79{)}^2=0.0001+0.0001=0.0002$. Reconstruction $=-\frac{1}{2}(0.0002)=-0.0001$. ELBO estimate $=\text{recon}-D_{KL}=-0.0001-0.2509=-0.2510$. ELBO ::: $\approx -0.2510$

$f(\sigma )=\frac{1}{2}(0+{\sigma }^2-\ln {\sigma }^2-1)$. Differentiate w.r.t. $\sigma$: $f^{\prime }(\sigma )=\frac{1}{2}(2\sigma -\frac{2}{\sigma })=\sigma -\frac{1}{\sigma }$. Set $=0\Rightarrow {\sigma }^2=1\Rightarrow \sigma =1$. Second derivative $f^{\prime \prime }(\sigma )=1+1/{\sigma }^2$, and at $\sigma =1$ this is $2>0$: a minimum. At $\sigma =1$, KL $=0$. Shape: a valley — squeezing ($\sigma <1$) blows up via $-\ln {\sigma }^2$, spreading ($\sigma >1$) blows up via ${\sigma }^2$. The prior's own width $\sigma =1$ is the sweet spot. Figure s02 (below) plots exactly this valley; check that your minimum lands on the orange dashed line at $\sigma =1$. Edge case $\sigma \to 0$: the term $-\ln {\sigma }^2=-2\ln \sigma \to +\infty$, so the KL diverges — a spike-like encoder is infinitely penalised. Numerically this is dangerous: as $\sigma \to 0$, $\ln {\sigma }^2\to -\infty$ overflows and $1/\sigma$ (in the gradient) blows up, causing NaNs. This is exactly why implementations parametrise the network to output $\log {\sigma }^2$ (or $\log \sigma$) directly instead of $\sigma$ — the log stays finite and differentiable, and $\sigma =e^{\frac{1}{2}\log {\sigma }^2}>0$ can never hit zero. Optimal sigma ::: $\sigma =1$, KL there $=0$; as $\sigma \to 0$ the KL $\to +\infty$.

As $\beta \to \infty$ the KL term dominates. The cheapest way to drive KL to $0$ is $\mu =0,\sigma =1$ for every $x$ — the encoder ignores the input entirely. Then $z$ carries no information about $x$, the decoder must reconstruct from noise, and reconstruction rots. This is Posterior Collapse: $q_{\varphi }(z|x)\to p(z)$ independent of $x$. The KL term, meant as a gentle regulariser, becomes a wrecking ball when over-weighted. One-line answer ::: encoder collapses to the prior for all $x$ — posterior collapse.

Bayes: $p_{\theta }(z|x)=\frac{p_{\theta }(x|z)p(z)}{p_{\theta }(x)}$. Take logs and solve for the evidence: $\log p_{\theta }(x)=\log p_{\theta }(x|z)+\log p(z)-\log p_{\theta }(z|x).$ This holds for every $z$; take ${\mathbb{E}}_{q_{\varphi }(z|x)}$ of both sides (the left side has no $z$, so it is unchanged): \log p_\theta(x)=\mathbb{E}_q[\log p_\theta(x|z)]+\mathbb{E}_q[\log p(z)]-\mathbb{E}_q[\log p_\theta(z|x)].\tag{1} The strategic goal (WHY the next move). We are aiming to reshape line (1) into "ELBO + gap". Look at what each target piece needs: the ELBO contains $-D_{KL}(q\Vert p(z))={\mathbb{E}}_q[\log p(z)]-{\mathbb{E}}_q[\log q_{\varphi }]$, and the gap $D_{KL}(q\Vert p_{\theta }(z|x))={\mathbb{E}}_q[\log q_{\varphi }]-{\mathbb{E}}_q[\log p_{\theta }(z|x)]$. Both need an ${\mathbb{E}}_q[\log q_{\varphi }(z|x)]$ term — one with a minus sign, one with a plus. But line (1) contains no $\log q_{\varphi }$ at all. The only way to conjure a $-\log q_{\varphi }$ and a $+\log q_{\varphi }$ without changing the value is to add zero in the form ${\mathbb{E}}_q[\log q_{\varphi }]-{\mathbb{E}}_q[\log q_{\varphi }]$. That is not a trick pulled from a hat — it is the forced consequence of wanting two KL divergences to appear. $\log p_{\theta }(x)={\mathbb{E}}_q[\log p_{\theta }(x|z)]+{\mathbb{E}}_q[\log p(z)]\underset{=0}{\underbrace{-{\mathbb{E}}_q[\log q_{\varphi }(z|x)]+{\mathbb{E}}_q[\log q_{\varphi }(z|x)]}}-{\mathbb{E}}_q[\log p_{\theta }(z|x)].$ Now hand the first $-{\mathbb{E}}_q[\log q_{\varphi }]$ to the prior term, and the second $+{\mathbb{E}}_q[\log q_{\varphi }]$ to the posterior term: $=\underset{\text{group A}}{\underbrace{{\mathbb{E}}_q[\log p_{\theta }(x|z)]+({\mathbb{E}}_q[\log p(z)]-{\mathbb{E}}_q[\log q_{\varphi }(z|x)])}}+\underset{\text{group B}}{\underbrace{({\mathbb{E}}_q[\log q_{\varphi }(z|x)]-{\mathbb{E}}_q[\log p_{\theta }(z|x)])}}.$ Group A: ${\mathbb{E}}_q[\log p(z)]-{\mathbb{E}}_q[\log q_{\varphi }]=-{\mathbb{E}}_q\left[\log \frac{q_{\varphi }}{p(z)}\right]=-D_{KL}(q\Vert p(z))$, so group A is exactly the ELBO $\mathcal{L}$. Group B: ${\mathbb{E}}_q[\log q_{\varphi }]-{\mathbb{E}}_q[\log p_{\theta }(z|x)]={\mathbb{E}}_q[\log \frac{q_{\varphi }}{p_{\theta }(z|x)}]=D_{KL}(q\Vert p_{\theta }(z|x))$, the gap. Hence $\log p_{\theta }(x)=\underset{\mathcal{L}}{\underbrace{{\mathbb{E}}_q[\log p_{\theta }(x|z)]-D_{KL}(q\Vert p(z))}}+\underset{\text{gap}\ge 0}{\underbrace{D_{KL}(q\Vert p_{\theta }(z|x))}}.$ Since $\log p_{\theta }(x)$ is fixed w.r.t. $\varphi$, and the gap $\ge 0$, pushing $\mathcal{L}$ up must push the gap down — so optimising $\varphi$ makes $q_{\varphi }$ approach the true posterior. This is exactly Variational Inference.

• E-step = maximise $\mathcal{L}$ over $\varphi$ (the approximation) with $\theta$ fixed. With a flexible $q$ this sets $q_{\varphi }(z|x)=p_{\theta }(z|x)$, closing the gap so $\mathcal{L}=\log p_{\theta }(x)$.
• M-step = maximise $\mathcal{L}$ over $\theta$ (the model) with $q$ fixed, raising the tightened bound. VAEs do a soft version: instead of an exact E-step they train an amortised encoder $q_{\varphi }$ by gradient ascent, jointly with $\theta$.

Let $w_k=\frac{p_{\theta }(x,z_k)}{q_{\varphi }(z_k|x)}$, i.i.d. with $\mathbb{E}[w_k]=\int q_{\varphi }\cdot \frac{p_{\theta }(x,z)}{q_{\varphi }}dz=\int p_{\theta }(x,z)dz=p_{\theta }(x)$ (so each weight is an unbiased estimator of the evidence). Define the $K$-sample average ${\bar{w}}_K=\frac{1}{K}{\sum }_{k=1}^K{w}_k$, so $\mathbb{E}[{\bar{w}}_K]=p_{\theta }(x)$ for every $K$, and ${\mathcal{L}}_K=\mathbb{E}[\log {\bar{w}}_K]$. Step 1 — every ${\mathcal{L}}_K$ is a valid bound. By Jensen for the concave $\log$: ${\mathcal{L}}_K=\mathbb{E}[\log {\bar{w}}_K]\le \log \mathbb{E}[{\bar{w}}_K]=\log p_{\theta }(x)$. So every ${\mathcal{L}}_K$ lies below the same ceiling. Step 2 — monotone in $K$ (the crux). Take $K+1$ samples $w_1,\dots ,w_{K+1}$. A short combinatorial fact: the full average equals the average of the $K+1$ leave-one-out $K$-sample averages, ${\bar{w}}_{K+1}=\frac{1}{K+1}{\sum }_{i=1}^{K+1}\left(\frac{1}{K}{\sum }_{k\ne i}{w}_k\right).$ Because $\log$ is concave, "log of an average $\ge$ average of the logs" gives $\log {\bar{w}}_{K+1}\ge \frac{1}{K+1}{\sum }_{i=1}^{K+1}\log \left(\frac{1}{K}{\sum }_{k\ne i}{w}_k\right).$ Take $\mathbb{E}$ of both sides; each leave-one-out term is a genuine $K$-sample estimator, so its expectation is ${\mathcal{L}}_K$, and the right side averages $K+1$ copies of ${\mathcal{L}}_K$ down to ${\mathcal{L}}_K$ itself. Hence ${\mathcal{L}}_{K+1}\ge {\mathcal{L}}_K$. Step 3 — the variance-to-tightness link, finished. The only looseness in Step 1 is the slack in Jensen's inequality, and that slack is governed entirely by how spread ${\bar{w}}_K$ is around its fixed mean $p_{\theta }(x)$: for a fixed mean, the more concentrated the random variable, the smaller the gap between $\mathbb{E}[\log \bar{w}]$ and $\log \mathbb{E}[\bar{w}]$ — with equality exactly when $\bar{w}$ is a constant. Averaging $K$ i.i.d. copies shrinks the variance, $\mathrm{Var}({\bar{w}}_K)=\frac{1}{K}\mathrm{Var}(w)\to 0$, so ${\bar{w}}_K\to p_{\theta }(x)$ (a constant) and the Jensen slack $\to 0$. Therefore the bound rises all the way to the ceiling as $K\to \infty$: ${\mathcal{L}}_K\uparrow \log p_{\theta }(x)$. Step 2 already made this rise monotone at every step rather than merely asymptotic, so the two arguments together give the full chain ${\mathcal{L}}_1\le {\mathcal{L}}_2\le \cdots \le {\mathcal{L}}_K\le \cdots ⟶\log p_{\theta }(x).$ Name the model. This tighter multi-sample bound is the objective of the Importance Weighted Autoencoder (IWAE); see Importance Weighted Autoencoders. Ordinary VAE training is the special case $K=1$. Name ::: Importance Weighted Autoencoder (IWAE) — the $K$-sample ELBO bound; VAE is the $K=1$ case.

The ELBO always equals ${\mathbb{E}}_q[\log p_{\theta }(x,z)]-{\mathbb{E}}_q[\log q_{\varphi }(z|x)]$. The second term is the negative entropy of $q$. For a flow, the change-of-variables rule gives the exact log-density (constants kept, since here they depend on $\varphi$ through $f_{\varphi }$ and cannot be dropped): $\log q_{\varphi }(z|x)=\log \mathcal{N}(ϵ;0,I)-\log |\det \frac{\partial f_{\varphi }}{\partial ϵ}|,\log \mathcal{N}(ϵ;0,I)=-\frac{d}{2}\ln (2\pi )-\frac{1}{2}\Vert ϵ{\Vert }^2.$ Note the contrast with Exercise 2.3: there we dropped the decoder's constant $-\frac{d}{2}\ln (2\pi )$ because it is independent of the parameters we optimise; here the Gaussian constant sits inside the entropy of $q_{\varphi }$ and, while $-\frac{d}{2}\ln (2\pi )$ itself is still constant, the surrounding log-Jacobian is not, so we carry the full expression to avoid mistakes. Substituting, $\mathcal{L}={\mathbb{E}}_{ϵ}[\log p_{\theta }(x,f_{\varphi }(ϵ))+\frac{d}{2}\ln (2\pi )+\frac{1}{2}\Vert ϵ{\Vert }^2+\log |\det J_{f_{\varphi }}(ϵ)\left|\right],$ where $J_{f_{\varphi }}(ϵ)=\partial f_{\varphi }/\partial ϵ$. The closed-form Gaussian KL is replaced by this log-Jacobian entropy. Benefit: $q_{\varphi }$ can now bend into non-Gaussian shapes, shrinking the gap KL against the true posterior — richer approximation, same ELBO scaffold. (If one only needs gradients, the constant $\frac{d}{2}\ln (2\pi )$ may be dropped — but state that choice, do not assume it.)

What single configuration makes the Gaussian-vs-unit-prior KL exactly zero? ::: $\mu =0,\sigma =1$ in every dimension. What quantity is the gap between ELBO and $\log p_{\theta }(x)$? ::: $D_{KL}(q_{\varphi }(z|x)\Vert p_{\theta }(z|x))$ — KL to the true posterior. As $\beta \to \infty$ in a $\beta$-VAE, what fails? ::: Posterior collapse — the encoder ignores $x$. Does using $K$ importance samples change the target? ::: No — it only tightens the same lower bound on $\log p_{\theta }(x)$; the model is the IWAE. Why parametrise the network to output $\log {\sigma }^2$ instead of $\sigma$? ::: To keep $\sigma >0$ and avoid $-\ln {\sigma }^2\to +\infty$ blow-ups / NaNs as $\sigma \to 0$.