Chapter 13 · I.13

Problems

A chapter that uses mathematics has to teach that mathematics by making the reader compute. A chapter that only displays equations has failed, however correct the equations are.

M3Load-bearing

5/5 problems5/4 variants10/10 exercisesquota met, and enforced

The contract

5 worked problems, minimum.
4 distinct variants, and no variant more than half of them.
10 exercises, every one with a published solution.
At least one numeric problem — present.
At least one symbolic problem — present.
At least one limit or counterexample problem — present.
At least one complexity or shape problem — present.
At least one ▲▲▲ problem — present.

Numerical instantiationSymbolic derivationLimiting caseProbabilisticCost accounting

Problem I.13.B01

One point, four bounds

numeric▲△△

Exact values where possible; six decimal places otherwise.

STATEMENT

Compute the log-evidence of one point under probabilistic PCA, then the lower bound that four different encoder distributions give it, and confirm that each gap is a KL divergence.

GIVEN

The four points of I.12 have covariance eigenvalues $4$ and $1$ along $\vec u_1=(0.8,0.6)$ and $\vec u_2=(-0.6,0.8)$ . The model has one latent dimension, $z\sim\mathcal N(0,1)$ and $\vec x\mid z\sim\mathcal N(\vec wz,\sigma^2\mat I)$ , at its maximum likelihood $\sigma^2=1$ and $\vec w=\sqrt3\,\vec u_1$ . Take $\vec x=(1,2)$ . The four candidate encoder distributions are $\mathcal N(0,1)$ , $\mathcal N(0,\tfrac14)$ , $\mathcal N(\tfrac{\sqrt3}2,1)$ and $\mathcal N(\tfrac{\sqrt3}2,\tfrac14)$ .

FIND

$\log p(\vec x)$ ; the posterior of $z$ ; the ELBO for each candidate; the gaps; and the two terms of the ELBO at the posterior.

STRATEGY

Use (I.13.2) for the evidence and the posterior, written in the eigenbasis. For the ELBO, both terms of (I.13.4) have closed forms for a Gaussian $q$ with mean $m$ and variance $v$ .

SOLUTION

The evidence. The point has coordinates $\vec u_1^{\mathsf T}\vec x=2$ and $\vec u_2^{\mathsf T}\vec x=1$ . The model’s covariance $\vec w\vec w^{\mathsf T}+\mat I$ has eigenvalues $4$ along $\vec u_1$ and $1$ along $\vec u_2$ , so its determinant is $4$ and

\begin{aligned} &\log p(\vec x)\\ &\;=-\log2\pi-\tfrac12\log4\\ &\;\quad-\tfrac12\Big(\frac{2^2}{4}+\frac{1^2}{1}\Big)\\ &\;=-\log2\pi-\log2-1\\ &\;=-3.531024. \end{aligned}

The posterior. With $\|\vec w\|^2=3$ and $\vec w^{\mathsf T}\vec x=2\sqrt3$ , (I.13.2) gives variance $\sigma^2/(\|\vec w\|^2+\sigma^2)=1/4$ and mean $2\sqrt3/4=0.866025$ .

The bounds. For $q=\mathcal N(m,v)$ ,

\begin{aligned} &\E_q\|\vec x-\vec wz\|^2\\ &\;=\|\vec x\|^2-2m\,\vec w^{\mathsf T}\vec x\\ &\;\quad+(m^2+v)\|\vec w\|^2\\ &\;=5-4\sqrt3\,m+3(m^2+v), \end{aligned}

and the reconstruction term is $-\log2\pi$ minus half of that. The price of the code is $\tfrac12(m^2+v-1-\log v)$ by (I.13.5). For the prior, $m=0$ and $v=1$ : the squared error averages to $8$ , the reconstruction term is $-\log2\pi-4$ , the price is $0$ , and

\operatorname{ELBO}=-5.837877.

The other three, computed the same way, are $-5.031024$ , $-4.337877$ and $-3.531024$ .

The gaps. Subtracting from $\log p(\vec x)$ gives $2.306853$ , $1.5$ , $0.806853$ and $0$ . By (I.13.4) each must equal the KL divergence from $\mathcal N(m,v)$ to the posterior $\mathcal N(0.866025,0.25)$ . For the prior,

\begin{aligned} &\tfrac12\Big(\frac{1}{0.25}+\frac{0.75}{0.25}\\ &\qquad-1+\log0.25\Big)\\ &\;\;=\tfrac12(4+3-1-1.386294)\\ &\;\;=2.306853. \end{aligned}

At the posterior. The squared error averages to $5-6+3=2$ , so the reconstruction term is $-\log2\pi-1=-2.837877$ . The price is $\tfrac12(0.75+0.25-1-\log0.25)=\log2=0.693147$ . Their difference is the evidence.

Answer

$\log p(\vec x)=-3.531024$ ; posterior $\mathcal N(0.866025,0.25)$ ; ELBOs $-5.837877$ , $-5.031024$ , $-4.337877$ , $-3.531024$ ; gaps $2.306853$ , $1.5$ , $0.806853$ , $0$ ; at the posterior, reconstruction $-2.837877$ and price $0.693147$ .

Check — sanity

Every bound is below the evidence and the gaps are non-negative. The correct width with mean $0$ costs $1.5$ , which is $\tfrac12(0.866025^2)/0.25$ : with equal variances, the KL is the squared distance between the means over twice the variance.

Check — numeric · i-13-vae-checks.py

The complete reproduction script uses only the Python standard library. From a checkout of the repository, run:

python3 scripts/snippets/i-13-vae-checks.py

It checks all fifteen problems of the chapter and the numbers quoted in its four propositions. It confirms the evidence by integrating $p(z)\,p(\vec x\mid z)$ numerically, and each gap against a separately computed KL divergence.

Executed in CI. The digits above are the digits it printed.

Where this breaks

With a nonlinear decoder neither the evidence nor the posterior has a closed form, and only the ELBO can be computed. The gaps above can be seen only because this model is linear.

Variation

Every corner of the rectangle has $|\vec u_1^{\mathsf T}\vec x|=2$ and $|\vec u_2^{\mathsf T}\vec x|=1$ , so all four have the same evidence, and the posterior means are $\pm0.866025$ . The average log-likelihood over the dataset is therefore also $-3.531024$ .

Problem I.13.B02

The bound, its gap and the price of a Gaussian code

symbolic▲▲▲

Exact values where possible; six decimal places otherwise.

STATEMENT

Derive the ELBO identity, derive the KL divergence of a Gaussian code from the standard normal and its derivatives, and show that the encoder that maximises the ELBO is the posterior.

GIVEN

A model $p(\vec z)\,p(\vec x\mid\vec z)$ with evidence $p(\vec x)>0$ , and a distribution $q(\vec z\mid\vec x)$ that is positive exactly where the posterior is. For the last part, the linear model of I.13.B01.

FIND

(a) $\log p(\vec x)-\operatorname{ELBO}(q)$ as a single divergence. (b) $\KL(\mathcal N(\mu,s^2)\,\|\,\mathcal N(0,1))$ and its derivatives with respect to $\mu$ and $\log s^2$ . (c) The $q=\mathcal N(m,v)$ that maximises the ELBO of the linear model, and the ELBO there.

STRATEGY

For (a), write $\log p(\vec x)$ as an average under $q$ of something that does not depend on $\vec z$ . For (b), average the log-ratio of densities. For (c), differentiate the closed-form ELBO.

SOLUTION

(a) Bayes’ rule gives $p(\vec x)=p(\vec x,\vec z)/p(\vec z\mid\vec x)$ for every $\vec z$ with positive posterior. The left side does not depend on $\vec z$ , so it equals its average under $q$ :

\begin{aligned} &\log p(\vec x)\\ &\;=\E_q\Big[\log\frac{p(\vec x,\vec z)}{p(\vec z\mid\vec x)}\Big]\\ &\;=\E_q\Big[\log\frac{p(\vec x,\vec z)}{q(\vec z\mid\vec x)}\Big]\\ &\;\quad+\E_q\Big[\log\frac{q(\vec z\mid\vec x)}{p(\vec z\mid\vec x)}\Big]. \end{aligned}

The first average is the ELBO and the second is $\KL(q(\vec z\mid\vec x)\,\|\,p(\vec z\mid\vec x))$ . Splitting $p(\vec x,\vec z)=p(\vec x\mid\vec z)\,p(\vec z)$ in the first gives the form in (I.13.4), with $-\KL(q\,\|\,p(\vec z))$ as its second term. Since a KL divergence is non-negative and zero only for equal distributions, the ELBO is a lower bound, tight exactly at the posterior.

(b) The log-densities are

\begin{aligned} \log q(z)&=-\tfrac12\log2\pi s^2\\ &\quad-\frac{(z-\mu)^2}{2s^2},\\ \log p(z)&=-\tfrac12\log2\pi-\frac{z^2}{2}. \end{aligned}

Under $q$ , $(z-\mu)^2$ averages to $s^2$ and $z^2$ to $\mu^2+s^2$ . Therefore

\begin{aligned} \KL&=-\tfrac12\log s^2-\tfrac12\\ &\quad+\tfrac12(\mu^2+s^2)\\ &=\tfrac12\big(\mu^2+s^2\\ &\qquad-1-\log s^2\big). \end{aligned}

This is (I.13.5) for one coordinate. With $\ell=\log s^2$ , so that $s^2=e^\ell$ ,

\begin{aligned} \frac{\partial\KL}{\partial\mu}&=\mu,\\[2pt] \frac{\partial\KL}{\partial\ell}&=\tfrac12(s^2-1). \end{aligned}

Both vanish only at $\mu=0$ , $s^2=1$ . For example $\KL(\mathcal N(1,0.25)\,\|\,\mathcal N(0,1))=\tfrac12(0.25+1.386294)=0.818147$ .

(c) From B01, with $a=\vec w^{\mathsf T}\vec x$ and $r=\|\vec w\|^2$ , the ELBO of $\mathcal N(m,v)$ is, up to a constant,

\begin{aligned} &-\frac{-2am+(m^2+v)r}{2\sigma^2}\\ &\quad-\tfrac12\big(m^2+v-1-\log v\big). \end{aligned}

Setting the derivative in $m$ to zero gives $a/\sigma^2=m(r/\sigma^2+1)$ , so $m=a/(r+\sigma^2)$ . In $v$ it gives $-r/2\sigma^2-\tfrac12+1/(2v)=0$ , so $v=\sigma^2/(r+\sigma^2)$ . These are the posterior’s mean and variance in (I.13.2). For the point $(1,2)$ they are $0.866025$ and $0.25$ , and by (a) the ELBO there equals $\log p(\vec x)=-3.531024$ .

Answer

(a) $\log p(\vec x)-\operatorname{ELBO}(q)=\KL(q(\vec z\mid\vec x)\,\|\,p(\vec z\mid\vec x))$ . (b) $\tfrac12(\mu^2+s^2-1-\log s^2)$ , with derivatives $\mu$ and $\tfrac12(s^2-1)$ . (c) $m=a/(r+\sigma^2)$ and $v=\sigma^2/(r+\sigma^2)$ , the exact posterior, where the ELBO is $-3.531024$ .

Check — sanity

The KL in (b) is zero at the prior and positive elsewhere: for fixed $s$ it grows with $\mu^2$ , and $s^2-1-\log s^2\geq0$ with equality only at $s^2=1$ , by $\log t\leq t-1$ .

Where this breaks

If $q$ is positive somewhere the posterior is zero, the KL in (a) is infinite and the ELBO is minus infinity. If $q$ is zero somewhere the posterior is not, (a) still holds but $q$ cannot be the posterior, so the gap never closes. A Gaussian prior and decoder make the posterior positive everywhere, and neither case arises. If the decoder variance also depends on $z$ , (c) has no closed form.

Variation

For a code with $k$ independent coordinates the KL is a sum of $k$ copies of (b). At the posterior of the linear model the price is $\tfrac12\log\big((r+\sigma^2)/\sigma^2\big)+\tfrac12(m^2+v-1)$ . At the maximum-likelihood $\vec w$ the second part averages to zero over the data, and at every corner of the rectangle it is exactly zero, leaving $\tfrac12\log4=\log2$ .

Problem I.13.B03

When the code switches off

limit▲▲△

Exact values where possible; six decimal places otherwise.

STATEMENT

Solve the linear β-VAE exactly, find its rate and distortion, and follow them as β goes to zero and as β grows past the point where the code is abandoned.

GIVEN

Centred data whose covariance has eigenvalues $\lambda_1>\lambda_2\geq\cdots$ , the top one along $\vec u_1$ . One latent dimension with prior $\mathcal N(0,1)$ , a decoder $\vec x\mid z\sim\mathcal N(\vec wz,\sigma^2\mat I)$ with $\sigma^2$ fixed, and an encoder $\mathcal N(m,v)$ chosen freely for each input. The objective is the average of $\E_q\log p(\vec x\mid z)-\beta\,\KL(q\,\|\,\mathcal N(0,1))$ . For numbers use the four points of I.12: $\lambda_1=4$ , $\lambda_2=1$ , $\sigma^2=1$ .

FIND

(a) The best encoder for a given $\vec w$ . (b) The best $\|\vec w\|^2$ . (c) The rate and the mean squared distortion there. (d) The limits $\beta\to0$ and $\beta\sigma^2\geq\lambda_1$ . (e) A table for the four points at $\beta=0.25$ , $1$ , $2$ and $4$ .

STRATEGY

Write $a=\vec w^{\mathsf T}\vec x$ , $r=\|\vec w\|^2$ and $c=r+\beta\sigma^2$ . Optimise the encoder first; what is left depends on $\vec w$ only through $r$ and its direction.

SOLUTION

(a) For one input, up to terms that do not involve $m$ or $v$ , the objective is

\begin{aligned} &\frac{2am-(m^2+v)r}{2\sigma^2}\\ &\quad-\frac\beta2\big(m^2+v-1-\log v\big). \end{aligned}

Setting the derivatives in $m$ and $v$ to zero,

m=\frac{a}{c},\qquad v=\frac{\beta\sigma^2}{c}.

(b) Substituting, the terms in $a^2$ combine to $a^2/(2\sigma^2c)$ and the terms in $v$ to $\tfrac\beta2\log(\beta\sigma^2/c)$ , up to constants. Averaged over the data, $\E[a^2]=\vec w^{\mathsf T}\mat S\vec w$ , which for a given $r$ is largest, $r\lambda_1$ , when $\vec w$ lies along $\vec u_1$ . What remains is

\begin{aligned} J(r)&=\frac{r\lambda_1}{2\sigma^2c}-\frac\beta2\log c\\ &\quad+\text{const}, \end{aligned}

with derivative

J'(r)=\frac{\beta}{2c}\Big(\frac{\lambda_1}{c}-1\Big).

It is positive while $c<\lambda_1$ and negative after, so

r=\max(\lambda_1-\beta\sigma^2,\;0).

(c) While the code is in use, $c=\lambda_1$ . Then $\E[m^2]=r/\lambda_1$ and $v=\beta\sigma^2/\lambda_1$ , and the average KL is

\begin{aligned} R&=\tfrac12\Big(\frac{r+\beta\sigma^2}{\lambda_1}-1\\ &\qquad-\log\frac{\beta\sigma^2}{\lambda_1}\Big)\\ &=\tfrac12\log\frac{\lambda_1}{\beta\sigma^2}. \end{aligned}

The squared error averages to $\operatorname{tr}\mat S-2r+(\E[m^2]+v)\,r=\operatorname{tr}\mat S-r$ , so

D=\sum_{i>1}\lambda_i+\beta\sigma^2.

(d) As $\beta\to0$ , $r\to\lambda_1$ , the encoder’s variance $\beta\sigma^2/\lambda_1$ goes to zero, the rate grows like $-\tfrac12\log\beta$ without bound, and the distortion falls to $\sum_{i>1}\lambda_i$ , the loss of the principal-axis autoencoder of I.12.T1. The code tends to $\vec u_1^{\mathsf T}\vec x/\sqrt{\lambda_1}$ : the principal score, scaled to the prior’s unit variance.

As $\beta\sigma^2$ rises to $\lambda_1$ , $r$ and $R$ fall to zero and $D$ rises to $\operatorname{tr}\mat S$ . For any larger $\beta$ , $J'(r)<0$ for all $r\geq0$ : the optimum is $\vec w=\vec 0$ , the encoder returns $\mathcal N(0,1)$ for every input, and the code is ignored.

(e) With $\lambda_1=4$ , $\lambda_2=1$ and $\sigma^2=1$ , $R=\tfrac12\log(4/\beta)$ and $D=1+\beta$ :

$\beta$	$r$	$v$	$R$	$D$
$0.25$	$3.75$	$0.0625$	$1.386294$	$1.25$
$1$	$3$	$0.25$	$0.693147$	$2$
$2$	$2$	$0.5$	$0.346574$	$3$
$4$	$0$	$1$	$0$	$5$

Answer

$m=a/(r+\beta\sigma^2)$ and $v=\beta\sigma^2/(r+\beta\sigma^2)$ ; $\|\vec w\|^2=\max(\lambda_1-\beta\sigma^2,0)$ along $\vec u_1$ ; while in use, $R=\tfrac12\log(\lambda_1/\beta\sigma^2)$ and $D=\sum_{i>1}\lambda_i+\beta\sigma^2$ . As $\beta\to0$ , $D\to\sum_{i>1}\lambda_i$ and $R\to\infty$ ; from $\beta\sigma^2=\lambda_1$ on, the code is unused, with $R=0$ and $D=\operatorname{tr}\mat S$ .

Check — sanity

At $\beta=1$ the encoder is the exact posterior of (I.13.2), $r=3$ is probabilistic PCA’s $\|\vec w\|^2$ , and the objective equals the log-evidence $-3.531024$ of every corner. The curve is continuous at the threshold: $1+4=5=\operatorname{tr}\mat S$ .

Where this breaks

The solution needs $\sigma^2$ fixed. If it is learned, at $\beta=1$ it settles at the average discarded eigenvalue, and a direction collapses when its eigenvalue does not exceed that. With several latent dimensions the same threshold applies to each eigen-direction separately.

Variation

At $\beta=8$ the code is unused, as at $\beta=4$ . At $\beta=0.5$ the rate is $\tfrac12\log8=1.039721$ and the distortion $1.5$ . Halving $\beta$ always adds $\tfrac12\log2=0.346574$ nats to the rate.

Problem I.13.B04

Two estimators of one gradient

probability▲▲△

Exact values where possible; six decimal places otherwise.

STATEMENT

Show that the pathwise and score-function estimators of one gradient are both unbiased, compute their variances, and follow them as the encoder’s distribution narrows.

GIVEN

$q=\mathcal N(\mu,s^2)$ , $z=\mu+s\varepsilon$ with $\varepsilon\sim\mathcal N(0,1)$ , and $f(z)=z^2$ . The target is $\partial\,\E_q[f(z)]/\partial\mu$ . The moments $\E[\varepsilon^2]=1$ , $\E[\varepsilon^4]=3$ and $\E[\varepsilon^6]=15$ , and odd moments vanish.

FIND

Both one-draw estimates, their means, their variances in general and at $\mu=1$ , $s=0.5$ , and their behaviour as $s\to0$ .

STRATEGY

Write each estimate as a polynomial in $\varepsilon$ and take moments term by term.

SOLUTION

The target. $\E_q[z^2]=\mu^2+s^2$ , so the derivative is $2\mu$ .

Pathwise. By (I.13.6) the estimate is $g_{\mathrm p}=f'(z)=2(\mu+s\varepsilon)$ . Its mean is $2\mu$ and its variance is

\operatorname{Var}g_{\mathrm p}=4s^2.

Score function. By (I.13.7), with $\partial\log q/\partial\mu=(z-\mu)/s^2=\varepsilon/s$ ,

g_{\mathrm s}=\frac{(\mu+s\varepsilon)^2\,\varepsilon}{s}.

Expanding, $(\mu+s\varepsilon)^2\varepsilon =\mu^2\varepsilon+2\mu s\varepsilon^2+s^2\varepsilon^3$ , whose mean is $2\mu s$ . So $\E g_{\mathrm s}=2\mu$ : unbiased. For the second moment,

\begin{aligned} &(\mu+s\varepsilon)^4\varepsilon^2\\ &\;=\mu^4\varepsilon^2+4\mu^3s\varepsilon^3\\ &\;\quad+6\mu^2s^2\varepsilon^4+4\mu s^3\varepsilon^5\\ &\;\quad+s^4\varepsilon^6, \end{aligned}

which averages to $\mu^4+18\mu^2s^2+15s^4$ . Dividing by $s^2$ and subtracting the squared mean,

\begin{aligned} &\operatorname{Var}g_{\mathrm s}\\ &\;=\frac{\mu^4+18\mu^2s^2+15s^4}{s^2}\\ &\;\quad-4\mu^2. \end{aligned}

The numbers. At $\mu=1$ , $s=0.5$ the pathwise variance is $1$ and the score-function variance is $(1+4.5+0.9375)/0.25-4=21.75$ .

As $s\to0$ . The pathwise variance $4s^2$ goes to zero. The score-function variance behaves like $\mu^4/s^2$ and grows without bound:

$s$	pathwise	score function
$0.5$	$1$	$21.75$
$0.1$	$0.04$	$114.15$
$0.01$	$0.0004$	$10{,}014$

Answer

Both have mean $2\mu$ . Variances $4s^2$ and $(\mu^4+18\mu^2s^2+15s^4)/s^2-4\mu^2$ : at $\mu=1$ , $s=0.5$ , $1$ against $21.75$ . As $s\to0$ the first vanishes and the second grows like $\mu^4/s^2$ .

Check — sanity

At $\mu=0$ the score-function variance is $15s^2$ , finite and small, while the pathwise variance is $4s^2$ . The score function is worst when $f$ is large where $q$ sits, because it multiplies $f$ itself by $\varepsilon/s$ .

Where this breaks

The comparison depends on $f$ . For an integrand that oscillates quickly, such as $\sin(10z)$ in I.13.X10, the order reverses. A baseline subtracted from $f$ can reduce the score function’s variance a great deal without biasing it.

Variation

Subtract the baseline $b=\E_q f=\mu^2+s^2$ from $f$ before multiplying by $\varepsilon/s$ . The estimate stays unbiased because $\E[\varepsilon/s]=0$ . The estimate becomes $2\mu\varepsilon^2+s(\varepsilon^3-\varepsilon)$ , with variance $8\mu^2+10s^2$ . At $\mu=1$ , $s=0.5$ that is $10.5$ : half of $21.75$ , and still ten times the pathwise $1$ .

Problem I.13.B05

What sampling the prior costs

complexity▲▲△

Three significant figures for sample counts; six decimal places otherwise.

STATEMENT

Count how many prior samples an unbiased estimate of the evidence needs, in one dimension and in ten, and set that against the cost of one ELBO step.

GIVEN

The estimate is the average of $p(\vec x\mid z_s)$ over $n$ draws $z_s$ from the prior. By (I.13.3) one draw has relative variance $r=\int p(z\mid\vec x)^2/p(z)\,dz-1$ , and for a posterior $\mathcal N(\mu,s^2)$ with $s^2<2$ the integral is $e^{\mu^2/(2-s^2)}/(s\sqrt{2-s^2})$ . Use the point $(1,2)$ under the model of I.13.B01, and again with noise variance $0.01$ , where $\|\vec w\|^2=3.99$ . The target is a relative standard error of ten per cent, so $n=r/0.1^2$ . For the encoder, take a hidden layer of $400$ units and a code of $k=20$ dimensions.

FIND

$r$ and $n$ for both noise levels in one dimension; $n$ in ten independent latent dimensions; the cost of one ELBO step; the size of the encoder’s output layer.

STRATEGY

Get the posterior from (I.13.2), apply the formula, and use the fact that for independent dimensions the integral in (I.13.3) factorises.

SOLUTION

Noise variance 1. The posterior is $\mathcal N(0.866025,0.25)$ , so $\mu^2=0.75$ and $2-s^2=1.75$ :

\begin{aligned} 1+r&=\frac{e^{0.75/1.75}}{0.5\sqrt{1.75}}\\ &=\frac{1.535063}{0.661438}\\ &=2.320797. \end{aligned}

So $r=1.320797$ and $n=132$ draws.

Noise variance 0.01. The posterior variance is $0.01/(3.99+0.01)=0.0025$ , a standard deviation of $0.05$ , and its mean is $2\sqrt{3.99}/4=0.998749$ . The same formula gives $r=22.316431$ and $n=2{,}230$ . Narrowing the posterior by a factor of ten multiplied the cost by about seventeen.

Ten dimensions. For a factorised prior and posterior the integral is a product, one factor $1+r$ per dimension:

1+r_{10}=(1+r)^{10}.

At noise variance $1$ that is $2.320797^{10}=4{,}530$ , so $n=4.53\times10^5$ . At $0.01$ it is $23.316431^{10}=4.75\times10^{13}$ , so $n=4.75\times10^{15}$ . The cost is exponential in the number of latent dimensions.

One ELBO step. The bound needs one draw from the encoder’s distribution per example: one encoder pass, one draw of $\vec\varepsilon$ , one decoder pass, and the closed form (I.13.5) for the KL. Its cost does not grow with $k$ except through the width of two layers.

The encoder’s output layer. It returns $\vec\mu$ and $\log\vec\sigma^2$ , $2k=40$ numbers, from $400$ hidden units: $400\cdot40+40=16{,}040$ parameters, twice what a deterministic code of the same width would need.

Answer

$r=1.320797$ and $n\approx132$ at noise variance $1$ ; $r=22.316431$ and $n\approx2{,}230$ at $0.01$ . In ten dimensions, $n\approx4.53\times10^5$ and $4.75\times10^{15}$ . One ELBO step costs one encoder and one decoder pass per example. The output layer has $16{,}040$ parameters.

Check — sanity

If the posterior were the prior, $\mu=0$ and $s=1$ , the formula gives $1+r=1/(1\cdot1)=1$ and $r=0$ : every draw would return $p(\vec x)$ exactly. The cost comes entirely from the posterior differing from the prior.

Where this breaks

The closed form needs $s^2<2$ ; a posterior wider than twice the prior’s variance makes the integral infinite. Real posteriors are not factorised, and the exponent of the dimension count then depends on their shape, but the growth remains exponential in general.

Variation

Draw from the posterior instead and weight each draw by $p(z)/p(z\mid\vec x)$ . Every draw then returns $p(\vec x)$ exactly and the variance is zero. An encoder close to the posterior is a cheap approximation to that ideal proposal.

Exercises

Every one has a published solution. A hidden solution is a solution; a missing one is an abandonment.

I.13.X01A Gaussian KL by hand, both ways roundnumeric▲△△

Compute $\KL(\mathcal N(1,0.25)\,\|\,\mathcal N(0,1))$ with (I.13.5). Then compute $\KL(\mathcal N(0,1)\,\|\,\mathcal N(1,0.25))$ with the general formula $\tfrac12\big(v_1/v_2+(m_1-m_2)^2/v_2-1+\log(v_2/v_1)\big)$ . Which direction does a VAE’s loss use, and what does the asymmetry mean?

Hint

The first is a narrow distribution measured against a wide one, the second a wide one against a narrow one.

Solution

The first is

\begin{aligned} &\tfrac12\big(1+0.25-1-\log0.25\big)\\ &\;\;=\tfrac12(0.25+1.386294)\\ &\;\;=0.818147. \end{aligned}

The second is

\begin{aligned} &\tfrac12\Big(\frac{1}{0.25}+\frac{1}{0.25}\\ &\qquad-1+\log0.25\Big)\\ &\;\;=\tfrac12(4+4-1-1.386294)\\ &\;\;=2.806853. \end{aligned}

The VAE’s loss uses the first direction: the encoder’s distribution measured against the prior, $\KL(q\,\|\,p)$ . It is cheap for $q$ to be narrower than the prior and costly to put mass where the prior has little. The reverse direction punishes the narrow distribution for every place the wide one reaches and it does not, so it is more than three times larger here.

The same direction appears in the gap of (I.13.4), $\KL(q\,\|\,p(\vec z\mid\vec x))$ . That is why an encoder fitted by the ELBO tends to sit inside the posterior rather than cover it, the effect I.13.X03 measures.

I.13.X02One reparameterised gradientgradient▲△△

Use the model of I.13.B01: $\vec w=\sqrt3\,\vec u_1$ , $\sigma^2=1$ , $\vec x=(1,2)$ , so $\vec w^{\mathsf T}\vec x=2\sqrt3$ and $\|\vec w\|^2=3$ . The encoder returns $m=0.5$ and $s=0.5$ , and the draw is $\varepsilon=0.6$ . Compute the one-draw gradient of the ELBO with respect to $m$ and $s$ , then the exact gradient with respect to $m$ .

Hint

The reconstruction term is $-\|\vec x-\vec wz\|^2/2$ plus a constant, with $z=m+s\varepsilon$ . The KL term is $\tfrac12(m^2+s^2-1-\log s^2)$ .

Solution

The code is $z=0.5+0.5\cdot0.6=0.8$ . The derivative of the reconstruction term with respect to $z$ is

\begin{aligned} \vec w^{\mathsf T}(\vec x-\vec wz) &=2\sqrt3-3\cdot0.8\\ &=1.064102. \end{aligned}

Since $\partial z/\partial m=1$ and $\partial z/\partial s=\varepsilon$ , and the KL has derivatives $m$ and $s-1/s$ ,

\begin{aligned} \frac{\partial}{\partial m}&=1.064102-0.5=0.564102,\\ \frac{\partial}{\partial s}&=1.064102\cdot0.6-(0.5-2)\\ &=2.138461. \end{aligned}

The exact gradient with respect to $m$ averages over $\varepsilon$ . The reconstruction derivative is linear in $z$ , so its average is its value at $z=m$ : $2\sqrt3-1.5=1.964102$ . Subtracting the KL’s $0.5$ gives $1.464102$ . The single draw underestimated it by $0.9$ , because $\varepsilon=0.6$ moved the code above its mean; an $\varepsilon$ of $-0.6$ would have overestimated it by the same amount.

Both one-draw derivatives point uphill. An unbiased estimate need not be close; it needs only to be right on average, which is all stochastic gradient ascent requires.

I.13.X03A factorised encoder misreports its uncertaintycounterexample▲▲△

A two-dimensional posterior is $\mathcal N(\vec 0,\mat\Sigma)$ with unit variances and correlation $\rho=0.9$ . The encoder’s family is factorised: $q=\mathcal N(\vec m,\operatorname{diag}(v_1,v_2))$ . Find the member of the family that minimises $\KL(q\,\|\,p)$ , the variance it reports for each coordinate, and the gap that remains.

Hint

For Gaussians in $d$ dimensions,

\begin{aligned} &\KL(q\,\|\,p)=\tfrac12\big(\operatorname{tr}(\mat\Sigma^{-1}\mat\Sigma_q)\\ &\quad+\vec m^{\mathsf T}\mat\Sigma^{-1}\vec m-d\\ &\quad+\log\det\mat\Sigma-\log\det\mat\Sigma_q\big), \end{aligned}

and the diagonal of $\mat\Sigma^{-1}$ is $1/(1-\rho^2)$ .

Solution

The mean term is minimised at $\vec m=\vec 0$ . With $\mat\Sigma_q$ diagonal, only the diagonal of $\mat\Sigma^{-1}$ enters the trace:

\begin{aligned} &\KL=\tfrac12\Big(\frac{v_1+v_2}{1-\rho^2}-2\\ &\quad+\log(1-\rho^2)\\ &\quad-\log v_1-\log v_2\Big). \end{aligned}

Each $v_i$ appears in its own term. Setting $1/(1-\rho^2)-1/v_i=0$ gives $v_i=1-\rho^2=0.19$ . Substituting, the trace term becomes $2$ and cancels the $-2$ :

\begin{aligned} \KL&=\tfrac12\big(\log0.19-2\log0.19\big)\\ &=-\tfrac12\log0.19\\ &=0.830366. \end{aligned}

So the best factorised encoder reports a variance of $0.19$ for each coordinate, whose true posterior variance is $1$ . It matches the conditional variance of one coordinate given the other, not the marginal one, and so understates the uncertainty more than fivefold.

The claim this breaks is that a trained encoder’s variance measures how uncertain the code is. The gap of $0.830366$ nats is the approximation gap: it does not shrink with training, because no member of the family is closer. Only a richer family, one with correlations, can remove it.

I.13.X04A point code has unbounded ratelimit▲▲△

Keep the encoder’s mean at $0.866025$ , the posterior mean of the point $(1,2)$ , and shrink its standard deviation $s$ . Compute the price of the code, $\KL(\mathcal N(0.866025,s^2)\,\|\,\mathcal N(0,1))$ , at $s=0.1$ , $0.01$ and $0.001$ . How fast does it grow, and what does the limit say about the autoencoder of Chapter I.12?

Hint

Only the $-\log s^2$ term of (I.13.5) grows as $s\to0$ .

Solution

With $\mu^2=0.75$ , (I.13.5) gives $\tfrac12(s^2-0.25-\log s^2)$ :

$s$	price of the code
$0.1$	$2.182585$
$0.01$	$4.480220$
$0.001$	$6.782756$

Each tenfold narrowing adds $\log10=2.302585$ nats, apart from the vanishing $s^2$ term: the price grows like $-\log s$ without limit.

A deterministic encoder is the limit $s\to0$ , so its rate is infinite. The autoencoder of I.12 is a VAE whose code costs nothing because its price is never charged. Proposition 4 reaches the same corner from the other side: in the linear model the encoder’s variance is $\beta\sigma^2/\lambda_1$ , which goes to zero only as $\beta\to0$ , and the rate $\tfrac12\log(\lambda_1/\beta\sigma^2)$ diverges with it.

The noise in the code is what makes the space between codes trained. A point code puts the decoder back where it was in I.12.

I.13.X05Where Jensen's inequality is tightsymbolic▲▲△

Derive the ELBO from Jensen’s inequality, and state exactly when the inequality is an equality. Then check the condition for the model of I.13.B01: show that $\log p(\vec x,z)-\log p(z\mid\vec x)$ takes the same value at every $z$ , and find it.

Hint

The logarithm is strictly concave, so $\log\E[Y]=\E[\log Y]$ only when $Y$ is constant.

Solution

Write $Y=p(\vec x,\vec z)/q(\vec z\mid\vec x)$ with $\vec z\sim q$ . Its mean is $\int p(\vec x,\vec z)\,d\vec z=p(\vec x)$ , so by Jensen’s inequality (Jensen's inequality 0.IT.05)

\begin{aligned} \log p(\vec x)&=\log\E_q[Y]\\ &\geq\E_q[\log Y]=\operatorname{ELBO}(q). \end{aligned}

Because $\log$ is strictly concave, equality holds exactly when $Y$ is the same number for every $\vec z$ that $q$ can draw. Then $q(\vec z\mid\vec x)\propto p(\vec x,\vec z)$ , and normalising gives $q=p(\vec x,\vec z)/p(\vec x)$ , the posterior. The constant is $p(\vec x)$ .

For the linear model, with $q$ the posterior, $Y$ is $p(\vec x,z)/p(z\mid\vec x)$ , which by Bayes’ rule is $p(\vec x)$ at every $z$ . The script evaluates $\log p(z)+\log p(\vec x\mid z)-\log\mathcal N(z;0.866025,0.25)$ at $z=-1$ , $0$ , $0.5$ and $2$ and gets $-3.531024$ each time, the log-evidence.

For any other $q$ the ratio varies with $z$ , and the spread of $\log Y$ is what Jensen’s inequality charges. The first derivation in proposition 2 shows that the charge is exactly $\KL(q\,\|\,p(\vec z\mid\vec x))$ .

I.13.X06The rate splits into information and a mismatchprobability▲▲△

At $\beta=1$ the four points have encoder distributions $\mathcal N(\pm0.866025,\,0.25)$ , two corners at each sign. Find the average KL to the prior, the mutual information between the example and its code, and the KL from the aggregate posterior to the prior, and check that (I.13.9) holds. Compare the aggregate posterior with the prior at $z=0$ .

Hint

Two corners share each distribution, so the information between the example and the code equals the information between the sign and the code. That is $\mathcal H(q(z))-\mathcal H(q(z\mid\text{sign}))$ (Mutual information 0.IT.04). The first entropy has no closed form; integrate it numerically.

Solution

Average KL. Every corner has $m^2=0.75$ and $v=0.25$ , so each KL is $\tfrac12(0.75+0.25-1-\log0.25)=\log2=0.693147$ , and so is the average.

Information. The aggregate posterior is the mixture $q(z)=\tfrac12\mathcal N(0.866025,0.25)+\tfrac12\mathcal N(-0.866025,0.25)$ . Each component has entropy $\tfrac12\log(2\pi e\cdot0.25)=0.725791$ . Numerical integration of $-q\log q$ gives the mixture’s entropy, and the difference is

I_q(n;z)=0.585939.

It is below $\log2=0.693147$ , the most a binary sign could carry, because the two components overlap.

Mismatch. $\KL(q(z)\,\|\,\mathcal N(0,1)) =-\mathcal H(q)+\tfrac12\log2\pi+\tfrac12\E_q[z^2]$ , and $\E_q[z^2]=0.75+0.25=1$ . Numerically it is $0.107208$ .

The identity. $0.585939+0.107208=0.693147$ , the average KL. Most of the rate pays for information about the example; about a sixth pays for the aggregate posterior not being the prior.

At $z=0$ . The mixture’s density is $2\cdot\tfrac12\mathcal N(0;0.866025,0.25)=0.178032$ , and the prior’s is $1/\sqrt{2\pi}=0.398942$ . The mixture has the prior’s variance but not its shape. Codes drawn from the prior near zero, which are the commonest, are codes the encoder produced less than half as often. That is a hole, and the second term of (I.13.9) is its price.

I.13.X07A mean reduction weights the KL by the input widthshape▲▲△

A VAE for $28\times28$ images in $[0,1]$ has a batch of $B$ inputs of shape $(B,1,28,28)$ , a code of width $k=32$ , and a Gaussian decoder with $\sigma^2=\tfrac12$ , so that $-\log p(\vec x\mid\vec z)$ is the summed squared error plus a constant. Give the shape of every tensor from the encoder’s output to the loss. Then compare the correct loss with the one written as the mean squared error over all elements plus the KL summed over the code and averaged over the batch. What $\beta$ does the second one train with, and by I.13.T3 what happens to the code?

Hint

A mean over all elements divides by $B\cdot784$ . The KL was divided only by $B$ .

Solution

Shapes. The encoder’s last layer returns $(B,64)$ , split into $\vec\mu$ and $\log\vec\sigma^2$ , each $(B,32)$ . The noise $\vec\varepsilon$ and the code $\vec z$ are $(B,32)$ . The decoder returns $(B,784)$ , compared with the input flattened to $(B,784)$ . The KL by (I.13.5) is summed over the $32$ coordinates to give $(B,)$ , as is the squared error over the $784$ pixels, and the loss averages both over the batch.

The two losses. The correct negative ELBO, up to a constant, is

\begin{aligned} \frac1B\sum_b\Big(&\sum_{i=1}^{784}(x_{bi}-\hat x_{bi})^2\\ &+\sum_{j=1}^{32}\KL_{bj}\Big). \end{aligned}

The second divides the squared error by $B\cdot784$ and the KL by $B$ . Multiplying it by $784$ gives the correct reconstruction term plus $784$ times the KL. The two have the same minimisers only if $\beta=784$ .

What happens. In the linear model of I.13.T3 a direction survives only if its variance exceeds $\beta\sigma^2=392$ . A pixel in $[0,1]$ has variance at most $\tfrac14$ , so the total variance of all $784$ pixels is at most $196$ , and no direction exceeds that. Every direction collapses: the encoder returns the prior and the decoder learns the mean image.

A nonlinear VAE is not bound by this exactly, but the pressure is the same, and blurred reconstructions with a KL near zero are the usual symptom. The shapes were all correct. Only the reductions disagreed.

I.13.X08A discrete code has no pathwise derivativecounterexample▲▲△

A one-bit code is $z\sim\operatorname{Bernoulli}(\pi)$ , and $f(z)=(z-0.3)^2$ . Find the true derivative of $\E[f(z)]$ with respect to $\pi$ . Write the draw as $z=\mathbf 1[u<\pi]$ with $u$ uniform on $(0,1)$ , and find what differentiating through it returns. Then find the score-function estimate at $\pi=0.5$ , its mean and its variance.

Hint

$\E[f(z)]=\pi f(1)+(1-\pi)f(0)$ . The function $\mathbf 1[u<\pi]$ is a step in $\pi$ .

Solution

The truth. $f(1)=0.49$ and $f(0)=0.09$ , so the derivative is $f(1)-f(0)=0.4$ .

Through the draw. For fixed $u$ , $\mathbf 1[u<\pi]$ is constant in $\pi$ except at $\pi=u$ , where it jumps. Its derivative is zero with probability one, so backpropagating through the draw returns $0$ . The claim broken is that writing a draw as a function of noise always gives a usable pathwise gradient: the function must be differentiable, and a step is not.

Score function. $\log q(z)=z\log\pi+(1-z)\log(1-\pi)$ , so its derivative is $z/\pi-(1-z)/(1-\pi)$ . At $\pi=0.5$ the estimate is $f(1)/0.5=0.98$ when $z=1$ and $-f(0)/0.5=-0.18$ when $z=0$ . Its mean is $\tfrac12(0.98-0.18)=0.4$ , correct, and its variance is

\begin{aligned} &\tfrac12(0.98-0.4)^2\\ &\;\;+\tfrac12(-0.18-0.4)^2=0.3364. \end{aligned}

The standard deviation $0.58$ exceeds the gradient itself, which is typical of the score function. Relaxations such as the Gumbel-softmax replace the step by a smooth function of the noise, accepting bias for a pathwise gradient.

I.13.X10The pathwise estimator is not always the quieter onecounterexample▲▲△

Let $q=\mathcal N(\mu,1)$ at $\mu=0$ and $f(z)=\sin(10z)$ . Find the true derivative of $\E_q[f(z)]$ with respect to $\mu$ , and the variances of the pathwise estimate $f'(z)$ and the score-function estimate $f(z)(z-\mu)$ .

Hint

For $z\sim\mathcal N(\mu,1)$ , $\E[\cos(\omega z)]=\cos(\omega\mu)e^{-\omega^2/2}$ . Use $\cos^2t=\tfrac12(1+\cos2t)$ .

Solution

The truth. $\E_q[\sin(10z)]=\sin(10\mu)e^{-50}$ , so the derivative at $\mu=0$ is $10e^{-50}=1.9\times10^{-21}$ : zero to any printed precision.

Pathwise. The estimate is $10\cos(10z)$ . Its mean is the derivative, and its second moment is $100\,\E[\cos^2(10z)]=50\big(1+e^{-200}\big)$ . Its variance is therefore $50$ to every printed digit.

Score function. The estimate is $\sin(10z)\,z$ . Its mean is also the derivative. Its second moment is $\E[\sin^2(10z)\,z^2]=\tfrac12\E[z^2]-\tfrac12\E[z^2\cos(20z)]$ , and the second average is $-399\,e^{-200}$ , so the variance is $0.5$ to every printed digit.

So the pathwise estimate is a hundred times noisier here. It reads the slope of $f$ at the draw, and a function that oscillates ten times per unit has slopes of size ten everywhere, which cancel only on average. The score function reads the value of $f$ , which never exceeds one. The claim broken is that reparameterisation always lowers the variance. It does for the smooth, slowly varying log-likelihoods of a VAE decoder; it does not in general.

I.13.X09The rate is information plus mismatch, for any encoderproof▲▲▲

Let the training set have $N$ examples, chosen with equal probability $q(n)=1/N$ , and let the encoder give $q(\vec z\mid n)$ . Prove (I.13.9): the average of $\KL(q(\vec z\mid n)\,\|\,p(\vec z))$ equals $I_q(n;\vec z)+\KL(q(\vec z)\,\|\,p(\vec z))$ , with $q(\vec z)=\frac1N\sum_nq(\vec z\mid n)$ . Deduce that a zero average KL means a code that carries no information, and that the rate bounds the information from above.

Hint

Insert $q(\vec z)$ into the log-ratio: $\log\frac{q(\vec z\mid n)}{p(\vec z)} =\log\frac{q(\vec z\mid n)}{q(\vec z)}+\log\frac{q(\vec z)}{p(\vec z)}$ .

Solution

The identity. Average the split in the hint over $q(n,\vec z)=q(n)\,q(\vec z\mid n)$ :

\begin{aligned} &\frac1N\sum_n\KL\big(q(\vec z\mid n)\,\|\,p(\vec z)\big)\\ &\;=\E_{q(n,\vec z)}\Big[\log\frac{q(\vec z\mid n)}{q(\vec z)}\Big]\\ &\;\quad+\E_{q(n,\vec z)}\Big[\log\frac{q(\vec z)}{p(\vec z)}\Big]. \end{aligned}

In the first term, $q(\vec z\mid n)/q(\vec z)=q(n,\vec z)/\big(q(n)\,q(\vec z)\big)$ , so it is $\KL(q(n,\vec z)\,\|\,q(n)q(\vec z))=I_q(n;\vec z)$ (Mutual information 0.IT.04). The second term depends on $\vec z$ alone, and the marginal of $q(n,\vec z)$ over $\vec z$ is $q(\vec z)$ , so it is $\KL(q(\vec z)\,\|\,p(\vec z))$ . $\blacksquare$

Zero rate means no information. Both terms on the right are non-negative. If the left side is zero, both are zero. Then $I_q(n;\vec z)=0$ : the code is independent of the example, so it tells the decoder nothing about which input it came from. Equivalently, each $\KL(q(\vec z\mid n)\,\|\,p)$ is non-negative and their average is zero, so each is zero and every $q(\vec z\mid n)$ is the prior. That is posterior collapse.

The rate bounds the information. Dropping the non-negative second term gives $I_q(n;\vec z)\leq\frac1N\sum_n\KL(q(\vec z\mid n)\,\|\,p)$ . The rate is an upper bound on what the code says about the example. Since $n$ takes $N$ values, the information is also at most $\log N$ ; any rate above that is spent on the mismatch.

The identity holds for every encoder, not only at an optimum. The script checks it for three encoders on the four points, including the collapsed one, where both sides are zero.