When the code switches off

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

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

(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

with derivative

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

(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

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

(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$:

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$.