VAE variants: β-VAE and VQ-VAE

The VAE ended with two complaints. Its samples are blurry, and its fixed Gaussian prior is only roughly matched by the codes the encoder actually produces, so prior samples can land where the decoder was never trained. This note follows three answers. The first keeps the model and changes the objective: a weight $\beta$ on the KL term decides how much information the code may carry. The second keeps the autoencoder and changes the code: a grid of discrete symbols from a learned codebook, with the prior fitted afterwards by a model strong enough to match them. The third keeps the Gaussian code and changes the decoder’s likelihood: an exact likelihood computed by a normalizing flow instead of a squared error summed over pixels.

Notation. As in the VAE note: an image $x$, a code $z$, an encoder $q_\varphi(z \mid x)$, a decoder $p_\theta(x \mid z)$, a prior $p(z)$. $\mathrm{sg}[\cdot]$ is the stop-gradient operator: it returns its argument in the forward pass and has zero derivative in the backward pass. A codebook is a set of $K$ vectors $e_1, \dots, e_K \in \mathbb R^d$.

Three ways to change a VAE#

The ELBO of one image has two terms:

$$\mathcal L(x) = \underbrace{\mathbb E_{q_\varphi(z \mid x)}\big[\log p_\theta(x \mid z)\big]}_{\text{reconstruction}} - \underbrace{D_{\mathrm{KL}}\big(q_\varphi(z \mid x)\,\|\,p(z)\big)}_{\text{KL to the prior}} .$$

The first term wants the code to say as much about $x$ as possible; the second charges for every bit it says. Changing their balance is the β-VAE. Changing what a code is, a continuous Gaussian or one symbol out of $K$, is the VQ-VAE. The two are independent: one moves along a trade-off, the other changes where the difficulty of generation lives. A third change, an optional extension at the end of the note, keeps both terms and replaces the form of the reconstruction term itself, $p_\theta(x \mid z)$.

β-VAE: a weight on the KL term#

The β-VAE (Higgins et al., 2017) maximizes

$$\mathcal L_\beta(x) = \mathbb E_{q_\varphi(z \mid x)}\big[\log p_\theta(x \mid z)\big] - \beta\, D_{\mathrm{KL}}\big(q_\varphi(z \mid x)\,\|\,p(z)\big) .$$

$\beta = 1$ is the plain VAE. For $\beta > 1$ the objective is still a lower bound on $\log p_\theta(x)$, because it subtracts more of a non-negative quantity than the ELBO does; for $\beta < 1$ it can exceed $\log p_\theta(x)$ and is no longer a bound. In both cases $\beta$ changes which model training prefers, not only the scale of the loss.

Rate and distortion#

Averaged over the data, the two terms have names from information theory (Alemi et al., 2018). The rate is the average KL,

$$R = \mathbb E_{x}\, D_{\mathrm{KL}}\big(q_\varphi(z \mid x)\,\|\,p(z)\big),$$

measured in nats. It bounds from above the mutual information between images and their codes under the encoder: the identity behind this, rate equals mutual information plus the KL from the aggregate posterior to the prior, is derived in the VAE note. The distortion $D = -\mathbb E_x \mathbb E_{q_\varphi(z \mid x)} \log p_\theta(x \mid z)$ is the average reconstruction loss. The negative ELBO is $D + R$, and the β-VAE minimizes $D + \beta R$: the Lagrangian of “reconstruct as well as possible with at most a given rate”. Every $\beta$ picks one point of the trade-off; $\beta = 1$ is the point that best bounds the likelihood.

Worked example: one latent dimension#

To see the trade-off with numbers, take one latent dimension, a standard normal prior $p(z) = \mathcal N(0, 1)$ and a fixed linear decoder $p(x \mid z) = \mathcal N(w z, \sigma^2)$ with $w = 1$ and $\sigma^2 = 0.1$. The data come from this very model, so $x \sim \mathcal N(0, w^2 + \sigma^2) = \mathcal N(0, 1.1)$. Holding the decoder fixed isolates what $\beta$ does to the encoder. Let the encoder be $q(z \mid x) = \mathcal N(m, v)$. For one $x$ the β-objective is, up to constants,

$$J(m, v) = -\frac{(x - w m)^2 + w^2 v}{2\sigma^2} - \frac{\beta}{2}\big(v + m^2 - 1 - \ln v\big),$$

using $\mathbb E_q (x - w z)^2 = (x - wm)^2 + w^2 v$ and the closed-form KL to a standard normal from the VAE note. Setting the derivatives to zero,

$$\frac{\partial J}{\partial v} = -\frac{w^2}{2\sigma^2} - \frac{\beta}{2} + \frac{\beta}{2v} = 0, \qquad \frac{\partial J}{\partial m} = \frac{w(x - wm)}{\sigma^2} - \beta m = 0,$$

gives the optimal encoder

$$m = \frac{w}{w^2 + \beta\sigma^2}\, x, \qquad v = \frac{\beta\sigma^2}{w^2 + \beta\sigma^2} .$$

At $\beta = 1$ these are the mean and variance of the exact posterior $p(z \mid x)$, so the ELBO is tight. Averaging over $x \sim \mathcal N(0, 1.1)$ gives the rate and the squared-error distortion $\mathbb E(x - wz)^2$:

$\beta$	$m / x$	$v$	rate $R$, nats	distortion $\mathbb E(x - wz)^2$
0.25	0.976	0.024	1.89	0.025
1	0.909	0.091	1.20	0.100
4	0.714	0.286	0.55	0.376
16	0.385	0.615	0.13	1.032

Raising $\beta$ shrinks the code toward the prior: the mean is pulled to 0 and the variance grows toward 1, so the code says less about $x$ and the reconstruction gets worse. At $\beta = 16$ the distortion, 1.03, is close to the variance of the data, 1.1: the code carries almost nothing. Lowering $\beta$ does the opposite and buys sharper reconstructions with a higher rate.

A high rate does not by itself mean a code far from the prior. The rate splits into the mutual information $I(X; Z)$ between images and codes plus the KL from the aggregate posterior $q(z) = \mathbb E_x\, q(z \mid x)$ to the prior. Here $q(z) = \mathcal N\big(0, (m/x)^2 \cdot 1.1 + v\big)$: at $\beta = 0.25$ its variance is 1.071, and the mismatch $D_{\mathrm{KL}}\big(q(z)\,\|\,p(z)\big)$ is 0.001 of the 1.89 nats, so almost all of the rate is information about $x$. Prior samples are poor when the aggregate posterior leaves the prior, or when the decoder is wrong where the prior puts mass. In this toy the decoder is fixed and equal to the true model, so $\beta$ changes only the encoder and samples decoded from the prior stay exact. In a jointly trained VAE a small $\beta$ also weakens the penalty on the mismatch, and that is how it can make prior samples worse.

Figure

β moves a VAE along the rate–distortion curve

What β buys and what it does not#

Higgins et al. (2017) report that $\beta > 1$ makes individual latent coordinates track individual factors of variation (position, scale, rotation of a sprite), which they call disentanglement. A disentangled representation is one in which a single latent unit is sensitive to changes in a single generative factor while being relatively invariant to changes in the others (Section 1). Both halves matter: the unit must follow its factor, and it must change little when another factor changes. On CelebA faces with $\beta = 250$, single units turn the head, change the expression toward a smile or change the hair style while the rest of the face stays; in a plain VAE ($\beta = 1$) the unit that turns the head also changes the expression, glasses and gender (Figure 1). Such a code would let us edit one feature at a time, reuse the factors for new tasks, reason about unseen combinations of known factors and notice inputs that fit no known combination (Section 1).

The mechanism is the pressure to spend little rate: a factorized prior is cheapest to match with a posterior whose coordinates are independent. $\beta < 1$ weakens that pressure: it buys sharper reconstructions, and the objective favours a disentangled code even less. The effect of $\beta > 1$ is not guaranteed either. Locatello et al. (2019) prove that without supervision or inductive biases disentanglement cannot be identified from data, and in a study of more than 12,000 trained models find that random seeds and hyperparameters matter more than the choice among disentangling objectives. Burgess et al. (2018) replace the fixed weight by a target rate $C$ that grows during training, $\mathbb E_q \log p_\theta(x \mid z) - \gamma\,\lvert D_{\mathrm{KL}} - C \rvert$, so that the code is allowed one more nat at a time. The useful reading of β-VAE is the one above: a dial on the rate, not a switch for interpretability.

VQ-VAE: a discrete latent code#

The VQ-VAE (van den Oord, Vinyals and Kavukcuoglu, 2017) gives up the continuous code. The encoder still outputs a map of vectors, $z_e(x) \in \mathbb R^{h \times w \times d}$, but each vector is replaced by the nearest entry of a learned codebook before it reaches the decoder:

$$k^*(i) = \arg\min_k \big\lVert z_e(x)_i - e_k \big\rVert^2, \qquad z_q(x)_i = e_{k^*(i)} .$$

The code of an image is the grid of indices $k^*(i)$: for the $32 \times 32$ images below, 64 integers between 1 and 512 at the bottom level. In ELBO terms, the posterior $q(z = k \mid x)$ is one-hot at $k^*$, and with a uniform prior over the $K$ codes its KL is $\log K$ per position, a constant. The KL term drops out of training, and the VQ-VAE trains as a plain autoencoder with a quantized bottleneck.

Figure

Vector quantization: the nearest codebook vector

Training through the argmin#

The argmin has no useful gradient: a small change of $z_e$ either keeps the same $e_k$ or jumps to another. The VQ-VAE uses the straight-through estimator (Bengio, Léonard and Courville, 2013): in the forward pass the decoder receives $z_q$, in the backward pass the gradient with respect to $z_q$ is copied to $z_e$ unchanged. In code it is one line, z_q = z_e + (z_q - z_e).detach(): the value is $z_q$, the derivative is that of $z_e$. The copy is a surrogate, not the derivative of quantization. Inside a cell the quantizer is constant, so its true Jacobian is zero; on a cell boundary it jumps. The straight-through estimator replaces that Jacobian by the identity, and this stays a substitution even when $z_e$ sits exactly on $e_{k^*}$, where the true derivative is still zero. Closeness only makes the heuristic plausible: the decoder’s gradient, evaluated at $e_{k^*}$, is then a reasonable guess for the direction in which to move $z_e$. The loss has three terms:

$$\mathcal L = -\log p_\theta\big(x \mid z_q(x)\big) + \big\lVert \mathrm{sg}[z_e(x)] - e \big\rVert^2 + \beta_c \big\lVert z_e(x) - \mathrm{sg}[e] \big\rVert^2 ,$$

summed over positions, with $e = e_{k^*}$. The reconstruction term trains the decoder and, through the straight-through copy, the encoder. The codebook term moves $e$ toward the encoder output; it does not touch the encoder. The commitment term moves the encoder output toward $e$; it does not touch the codebook. Without it nothing stops the encoder outputs from drifting away from the codebook faster than the codebook can follow. The paper uses $\beta_c = 0.25$ and reports that results are robust to it.

Worked example: one quantization step#

Take $d = 2$, $K = 3$, the codebook $e_1 = (0, 0)$, $e_2 = (2, 0)$, $e_3 = (0, 2)$ and one encoder output $z_e = (1.4, 0.3)$, the red point of the figure. The squared distances are

$$\lVert z_e - e_1\rVert^2 = 1.96 + 0.09 = 2.05, \quad \lVert z_e - e_2\rVert^2 = 0.36 + 0.09 = 0.45, \quad \lVert z_e - e_3\rVert^2 = 1.96 + 2.89 = 4.85 .$$

The nearest vector is $e_2$, so the decoder receives $z_q = (2, 0)$ and the code is the index 2. The codebook term is $0.45$ and the commitment term is $0.25 \cdot 0.45 = 0.1125$: the same distance, with gradients sent in opposite directions. If the decoder’s loss has gradient $g$ with respect to $z_q$, the encoder receives exactly $g$ at $z_e$, plus the commitment gradient $2\beta_c(z_e - e_2) = 0.5 \cdot (-0.6, 0.3) = (-0.3, 0.15)$, which points from $e_2$ toward $z_e$, so a descent step moves $z_e$ toward $e_2$.

The codebook as online k-means#

The codebook term moves each $e_k$ toward the encoder outputs assigned to it, which is the update of k-means. The appendix of the VQ-VAE paper replaces it by exponential moving averages. For each code keep a count $N_k$ and a sum $m_k$, update them with the $n_k$ outputs assigned to code $k$ in the batch,

$$N_k \leftarrow \gamma N_k + (1 - \gamma)\, n_k, \qquad m_k \leftarrow \gamma\, m_k + (1 - \gamma) \sum_{i:\, k^*(i) = k} z_e(x)_i, \qquad e_k = \frac{m_k}{N_k},$$

with $\gamma = 0.99$, and drop the codebook term from the loss. Each $e_k$ is then a running mean of its cell, exactly the k-means centroid computed online. With $N_2 = 10$, $m_2 = (20, 0)$ (so $e_2 = (2, 0)$) and a batch of two outputs $(1.4, 0.3)$ and $(1.8, -0.1)$ assigned to it, the update gives $N_2 = 9.92$, $m_2 = (19.832, 0.002)$ and $e_2 = (1.9992, 0.0002)$: a small step toward the batch mean $(1.6, 0.1)$, as a moving average with $\gamma = 0.99$ should take.

The weak point of this scheme is the one k-means has: a code that no output selects receives no update and stays where it was initialized, a dead code. With many dead codes the effective codebook is small. Practical fixes restart unused codes at recent encoder outputs, as in Jukebox (Dhariwal et al., 2020), or look codes up in a low-dimensional, normalized space, as in ViT-VQGAN (Yu et al., 2022). Implementations often add a small constant to the counts (Laplace smoothing) to avoid dividing by zero for unused codes; this does not revive them.

Where the prior went#

Because the KL term is constant, training leaves the prior uniform, and sampling the codes independently and uniformly would give an image of random patches. The VQ-VAE therefore fits the prior after the autoencoder: encode the training set into grids of indices and train an autoregressive model on them, a PixelCNN in the paper, which predicts each index from the ones before it. Generation samples a grid from this prior and decodes it. The two-stage structure is the point of the method. The autoencoder handles the pixels, and the prior only has to model a grid that is tens of times smaller and made of symbols, where autoregressive models are strong; how they work is the subject of the autoregressive models note.

VQ-VAE-2: two levels of codes#

A single grid has to carry both the global layout and the fine texture. VQ-VAE-2 (Razavi, van den Oord and Vinyals, 2019) uses two: a coarse top grid for the layout and a finer bottom grid for detail that is quantized after seeing the top. For ImageNet at $256 \times 256$ the grids are $32 \times 32$ and $64 \times 64$ with 512 codes each; FFHQ at $1024 \times 1024$ uses three levels. The prior is hierarchical too: an autoregressive model over the top codes, with causal self-attention interleaved with its gated convolutions, and a conditional one over the bottom codes given the top, without attention to keep the larger grid affordable. The paper also filters samples with a pretrained ImageNet classifier, a rejection step that trades diversity for quality.

Figure

A small VQ-VAE-2: two code maps

How many bits#

The figure shows a small version for $32 \times 32$ CIFAR-10 images: a $4 \times 4$ top grid and an $8 \times 8$ bottom grid, 512 codes each. An image is $32 \cdot 32 \cdot 3 = 3072$ bytes, or 24,576 bits. Its code is $16 + 64 = 80$ indices of $\log_2 512 = 9$ bits, or 720 bits, about 34 times fewer. Sampling cost scales the same way: an autoregressive model over raw subpixels needs 3072 sequential steps per image, the two priors need $16 + 64 = 80$. At the paper’s $256 \times 256$ scale the gap is $196{,}608$ subpixels against $32^2 + 64^2 = 5120$ codes.

VAPNEV: a flow likelihood instead of pixel-wise reconstruction#

Optional depth: a flow decoder with an exact likelihood, and its conditional version. The rest of the note does not depend on it.

The changes above leave the decoder’s likelihood alone. With a factorized Gaussian decoder, $-\log p_\theta(x \mid z)$ is a squared error summed over pixels, and the best decoder output for a code is the average of all images that fit it: the blur of the VAE. VAPNEV (Agrawal and Dukkipati, 2016), an anagram of VAE-NVP, keeps the Gaussian code and the ELBO and replaces this term by an exact likelihood computed with a real NVP flow (Dinh, Sohl-Dickstein and Bengio, 2017) that is conditioned on the code. Flows are treated in full in the normalizing flows note later in the course; this section needs only the change-of-variables formula and the coupling layer, both restated here.

Notation. $y \in \mathbb R^D$ is an intermediate variable of the same size as the image. For every code $z$, $f_z: y \mapsto x$ is an invertible map, a flow written in the course’s latent-to-data direction, with inverse $f_z^{-1}: x \mapsto y$. Agrawal and Dukkipati write the map in the data-to-latent direction.

The decoder likelihood as a change of variables#

The decoder network does not output an image. It outputs the mean and the variances of a Gaussian over $y$, and the flow turns $y$ into the image:

$$y \mid z \sim \mathcal N\big(\mu_\theta(z), \sigma_\theta^2(z)\big), \qquad x = f_z(y).$$

By the change-of-variables formula, the likelihood of an image is exact (Section 3.1, Eqs. 10–11):

$$\begin{aligned} \log p_\theta(x \mid z) ={}& \log\mathcal N\big(f_z^{-1}(x);\, \mu_\theta(z), \sigma_\theta^2(z)\big) \\ &+ \log\left|\det\frac{\partial f_z^{-1}(x)}{\partial x}\right| . \end{aligned}$$

The ELBO keeps its form, with this term in place of the squared error:

$$\begin{aligned} \mathcal L(x) ={}& \mathbb E_{q_\varphi(z \mid x)}\Big[\log\mathcal N\big(f_z^{-1}(x);\, \mu_\theta(z), \sigma_\theta^2(z)\big) \\ &\qquad + \log\left|\det\frac{\partial f_z^{-1}(x)}{\partial x}\right|\Big] \\ &- D_{\mathrm{KL}}\big(q_\varphi(z \mid x)\,\|\,p(z)\big). \end{aligned}$$

The Gaussian now lives in the space of $y$, not of pixels. The flow mixes pixels, so the model can represent dependencies between pixels for a given code, which a factorized Gaussian decoder cannot: it treats them as independent noise around one mean image. With $f_z$ the identity, the model is the ordinary VAE with a Gaussian decoder of learned variances. The formula also holds when the flow ignores $z$; the code then acts only through $\mu_\theta$ and $\sigma_\theta$ (Section 3.1).

Conditional coupling layers#

The inverse map $f_z^{-1}$ is a stack of conditional coupling layers (Section 3.2). As in an affine coupling layer, the input of a layer is split into two parts $x_a$ and $x_b$; the first passes unchanged, and the second is scaled and shifted by functions of the first. Here those functions also see the code (Eqs. 12–13):

$$y_a = x_a, \qquad y_b = x_b \odot \exp\big(s_z(x_a)\big) + t_z(x_a).$$

For a fixed $z$ the Jacobian is still triangular, with $\log\lvert\det\rvert = \sum_i s_z(x_a)_i$, and the inverse is $x_b = \big(y_b - t_z(y_a)\big) \odot \exp\big(-s_z(y_a)\big)$ (Eqs. 14–15). Inverting needs the code, and the code is always known: from the encoder during training and from the prior during generation.

The scale combines a network $s_1$ of the unchanged part (two residual blocks in the paper) with a small transposed-convolution network $s_2$ that turns the low-dimensional code into a map of the same shape (Eq. 16):

$$\begin{aligned} s_z(x_a) ={}& \kappa \odot s_1(x_a) \odot s_2(z) + \beta_1 \odot s_1(x_a) \\ &+ \beta_2 \odot s_2(z) + b . \end{aligned}$$

The vectors $\kappa$, $\beta_1$, $\beta_2$ and $b$ are learned per channel and broadcast over positions, and the shift $t_z$ is built the same way from its own pair of networks. The product term lets the code switch features of $x_a$ on and off, and the direct term gives every layer a shortcut to the code. The authors credit these shortcuts with strong conditioning and with good results from a flow with fewer layers.

Worked example: one exact likelihood evaluation#

A toy case with $D = 2$ and a scalar code $z = 0.5$. One conditional coupling layer uses $s_z(x_a) = x_a z$ and $t_z(x_a) = x_a - z$, and the decoder outputs $\mu_\theta(z) = (z, 2z) = (0.5, 1.0)$ and $\sigma_\theta = (1, 0.5)$. Take the image $x = (0.4, 1.0)$.

1. Map to $y$. $y_a = 0.4$. The layer computes $s = 0.4 \cdot 0.5 = 0.2$ and $t = 0.4 - 0.5 = -0.1$, so $y_b = 1.0 \cdot e^{0.2} - 0.1 = 1.1214$.
2. Gaussian term. $\log\mathcal N(0.4;\, 0.5, 1^2) + \log\mathcal N(1.1214;\, 1.0, 0.5^2) = -0.9239 - 0.2553 = -1.1792$.
3. Log-determinant. $s = 0.2$: the layer stretches the $x_b$ direction by $e^{0.2}$, so the density of $x$ is $e^{0.2}$ times the density of $y$.
4. Total. $\log p_\theta(x \mid z) = -1.1792 + 0.2 = -0.9792$ nats.

Going back, $x_b = (1.1214 + 0.1)\, e^{-0.2} = 1.0$: the layer is inverted exactly once $z$ is known.

Training, sampling and results#

• Training. The model maximizes the ELBO with Adam. The authors observed the KL term going to zero early in training, the posterior collapse of a decoder strong enough to model the images without the code, and countered it with the KL annealing of Bowman et al. (2016), Section 3.1: a weight on the KL term that starts at zero and grows to one during training (Section 3.3).
• Continuous data. A flow needs continuous inputs. The 8-bit images are dequantized with uniform noise (see the discrete-data callout), scaled to $[0, 1]$ and mapped by $x' = \operatorname{logit}\big(\alpha + (1 - \alpha)\,x\big)$ with $\alpha = 0.05$, and the bound is corrected by the log-Jacobian of this preprocessing so that it refers to the images themselves (Section 4.1). How flows are applied to images is covered in flow architectures.
• Sampling. Draw $z \sim \mathcal N(0, I)$, then $y \sim \mathcal N(\mu_\theta(z), \sigma_\theta^2(z))$, and return $x = f_z(y)$ (Section 3.4). A VAE’s decoder is a distribution too, but its samples for one code differ only by independent pixel noise, so usually only its mean is shown. Here the noise in $y$ passes through the flow, which mixes pixels, so one code gives different images; $y = \mu_\theta(z)$ gives one image per code. Reconstruction does the same with $z$ from the encoder (Section 3.5).
• Architecture (Section 4.2). An 8-layer convolutional encoder and an 8-layer transposed-convolution decoder, a 256-dimensional code, and a multi-scale real NVP with 2 scales. Each scale has 3 conditional coupling layers with checkerboard masks and 3 with channel-wise masks, with 64 filters in the residual blocks of the first scale and 128 in the second.

Test results in bits per dimension (Tables 1–2; lower is better, and $<$ marks an upper bound from an ELBO):

Model	CIFAR-10	CelebA
PixelRNN	3.00	—
VAE with inverse autoregressive flow	$< 3.28$	—
Real NVP	3.49	3.02
Convolutional DRAW	$< 3.59$	—
VAPNEV	$< 3.55$	$< 2.8$

On CIFAR-10 a VAE with one stochastic layer and this likelihood is roughly level with real NVP, which uses 8 residual blocks per coupling layer against 2, and its bound is slightly below that of convolutional DRAW, a VAE with several stochastic layers and recurrence. It stays behind autoregressive models and behind a VAE with a flow posterior. On CelebA it beats real NVP with 2 scales instead of 4 (Section 4.4). The reconstructions keep high-level content such as pose, background colour, hair colour and expression (Figures 3–4). The bits per dimension are those of the autoregressive models note.

Conditional VAPNEV#

Agrawal and Dukkipati suggest conditional density modelling as an extension but do not implement it (Section 5). The construction is that of the conditional VAE. For an image $x$ with a condition $c$, for CelebA its 40 binary attributes, condition the encoder, $q_\varphi(z \mid x, c)$, the decoder’s Gaussian, $\mathcal N\big(\mu_\theta(z, c), \sigma_\theta^2(z, c)\big)$, and the coupling layers, whose code branch $s_2$ (and its shift counterpart) receives $z$ and $c$ together. Write the resulting flow $f_{z,c}$ and keep the prior $p(z) = \mathcal N(0, I)$. The bound becomes

$$\begin{aligned} \log p_\theta(x \mid c) \ge{}& \mathbb E_{q_\varphi(z \mid x, c)}\Big[\log\mathcal N\big(f_{z,c}^{-1}(x);\, \mu_\theta(z, c), \sigma_\theta^2(z, c)\big) \\ &\qquad + \log\left|\det\frac{\partial f_{z,c}^{-1}(x)}{\partial x}\right|\Big] \\ &- D_{\mathrm{KL}}\big(q_\varphi(z \mid x, c)\,\|\,p(z)\big). \end{aligned}$$

Sampling fixes the attributes, draws $z$ from the prior, draws $y$ from the decoder’s Gaussian $\mathcal N\big(\mu_\theta(z, c), \sigma_\theta^2(z, c)\big)$ and returns $x = f_{z,c}(y)$. Changing one attribute while keeping $z$ asks for the same face with that attribute changed, to the extent that the code has not absorbed the attribute itself.

What changed#

	VAE	β-VAE	VQ-VAE / VQ-VAE-2	VAPNEV
Code	continuous, Gaussian	continuous, Gaussian	grid of indices into a codebook	continuous, Gaussian
Objective	ELBO	reconstruction − β · KL	reconstruction + codebook + commitment	ELBO with an exact flow likelihood
Role of the prior	fixed $\mathcal N(0, I)$, enforced by the KL	the same, with weight $\beta$	uniform during training; learned afterwards by an autoregressive model	fixed $\mathcal N(0, I)$, enforced by the KL
Gradient through the code	reparameterization	reparameterization	straight-through copy	reparameterization
Sampling	decode $z \sim \mathcal N(0, I)$	the same	sample codes autoregressively, then decode	decode $z$ to a Gaussian over $y$, sample $y$, apply the flow
Typical trouble	blur, prior holes	the rate–distortion trade-off itself	dead codes, slow sequential prior	KL collapse without annealing, a costly full-resolution decoder

Limitations#

• β moves along a trade-off; it does not escape it. Lower distortion costs rate. Rate is mostly information about the image, not necessarily a mismatch with the prior; poor prior samples come from the aggregate posterior leaving the prior, which a small $\beta$ penalizes less. Disentanglement is a possible side effect, not a property of the objective.
• The straight-through gradient is a surrogate. It uses an identity Jacobian where the quantizer’s true derivative is zero, so it is never the exact gradient; the commitment term keeps $z_e$ near its codebook vector, which makes the substitution reasonable, not exact.
• Codebook collapse. Without restarts or a normalized lookup, a large share of the codes can go unused.
• Two-stage training. The prior is fitted to codes of a frozen autoencoder; errors of the first stage, such as a code grid that is hard to predict, cannot be fixed by the second.
• A flow decoder is strong and expensive. VAPNEV’s decoder can model the images without the code, so the KL term collapses unless it is annealed, and its flow works at full image resolution, which costs far more than a transposed-convolution decoder.
• Sequential sampling. The autoregressive prior samples one index at a time. Smaller than pixels, the grid is still sequential; later work replaces the prior by transformers over the same codes or by diffusion in a continuous latent space (latent diffusion).

Summary#

• β-VAE scales the KL term by $\beta$ and moves the model along the rate–distortion curve: $\beta > 1$ gives a lower rate and worse reconstructions, $\beta < 1$ the opposite. $\beta = 1$ recovers the standard ELBO, $\beta > 1$ maximizes a looser lower bound, and for $\beta < 1$ the objective need not bound the likelihood.
• VQ-VAE replaces the Gaussian code by the nearest codebook vector. The KL becomes a constant, the argmin is crossed by the straight-through estimator, and the loss adds codebook and commitment terms, or an EMA update that makes the codebook an online k-means.
• The prior is learned after the autoencoder by an autoregressive model over the code grid. VQ-VAE-2 uses two grids, a coarse one for layout and a fine one for detail, and a prior for each.
• VAPNEV keeps the Gaussian code and the ELBO and replaces the pixel-wise reconstruction by the exact likelihood of a real NVP flow whose base Gaussian and coupling layers depend on the code. It needs KL annealing to keep the code in use. Its conditional version adds a label to the encoder, the decoder and the coupling layers.

Notebook#

The VQ-VAE part follows the VQ-VAE-2 notebook of the HSE-LAMBDA Deep Generative Models repository (cell numbers count every cell from 1). The notebook has no saved outputs, so it documents the procedure rather than results.

• Cell 6: Quantize computes squared distances to all 512 codebook vectors at once with the expansion $\lVert a - e\rVert^2 = \lVert a\rVert^2 - 2a^\top e + \lVert e\rVert^2$, a single matrix product, picks the nearest, updates the codebook with the EMA above (decay 0.99), returns the commitment distance diff = mean((quantize.detach() - input)²) and applies the straight-through estimator input + (quantize - input).detach(). The codebook is a buffer drawn from a standard normal, not a parameter, so Adam never touches it; it moves only through the EMA, and only in training mode. The counts are smoothed as $N_k \leftarrow \frac{N_k + \epsilon}{n + K\epsilon}\, n$ with $n = \sum_k N_k$ and $\epsilon = 10^{-5}$, which keeps their total and makes every count positive. One issue: the counts start at $0$ while the sums $m_k$ start at the random codebook. After the first update, a code that the first batch did not select is therefore divided by a smoothed count of about $\epsilon$, and its norm grows about $10^5$ times. It shrinks back only as $0.99^t$, which takes about $\ln 10^5 / {-\ln 0.99} \approx 1150$ steps. In our own toy run of this update with standard-normal encoder outputs (not a notebook output), only 2 to 3 of the 512 codes were selected between steps 2 and 1000. Starting the counts at $1$, so that the first quotient $m_k / N_k$ is the random codebook itself, avoids the jump.
• Cells 7–10: architecture. ResBlock is ReLU, a $3 \times 3$ convolution from 128 to 32 channels, ReLU and a $1 \times 1$ convolution back to 128, added to its input. The bottom encoder applies two $4 \times 4$ stride-2 convolutions and a $3 \times 3$ convolution, the top encoder one $4 \times 4$ stride-2 convolution and a $3 \times 3$ convolution, each followed by residual blocks; the decoders mirror this with a $3 \times 3$ convolution, residual blocks and $4 \times 4$ stride-2 transposed convolutions. The settings are channel=128, n_res_block=2, n_res_channel=32, and codebooks of 512 vectors of dimension 64. The bottom encoder downsamples by 4 (to $8 \times 8$ for CIFAR-10), the top encoder by 2 more (to $4 \times 4$), and $1 \times 1$ convolutions project the features to the 64-dimensional code vectors before quantization. The top code reaches the bottom level along two separate paths: dec_t, a small decoder, brings it back to $8 \times 8$ to be concatenated with the bottom features and quantized as the bottom code, while upsample_t, a single transposed convolution, enlarges it for the final decoder, which receives it together with the bottom code.
• Cell 12: training minimizes the mean squared reconstruction error plus $0.25 \cdot$ diff, the commitment term; there is no codebook term because the EMA replaces it. Every sample_interval steps (200 in cell 20) it saves a grid of the batch above its reconstructions, computed in evaluation mode.
• Cells 13 and 22: infer_vqvae reconstructs two batches of 64 CIFAR-10 test images in evaluation mode and saves the 128 reconstructions to a file.
• Cells 14–20: CIFAR-10 at $32 \times 32$ normalized to $[-1, 1]$ (the comment in cell 15 still says $64 \times 64$), batch size 32, Adam at $3 \cdot 10^{-4}$, five epochs.
• Cells 25–27: the priors. Masked convolutions of types A and B, and two PixelCNN priors with seven $7 \times 7$ masked layers of 64 channels and plain ReLUs between them, no gating. The input is an embedding of the code indices, Embedding(512, 64), and a $1 \times 1$ convolution gives 512 logits per position. The unconditional prior models the top codes. The conditional one models the bottom codes: it embeds the nearest-neighbour upsampled top codes with Embedding(512, 16), concatenates them with the embedding of the bottom codes and mixes the two with a $1 \times 1$ convolution. Two issues. This plain stack has the blind spot and lacks the gated layers and the attention of the paper’s priors. And because the condition is mixed into the same channels before the first type-A masked layer, the mask also hides the condition at the position being predicted. Each bottom index sees the top codes only through earlier positions, so the top-left index of every $2 \times 2$ block, whose top code appears at no earlier position, is predicted without it. We checked this by differentiating a re-implementation of the network: the logits at a position have zero gradient with respect to the condition at that position. The condition should enter without a mask, added in every layer as in the conditional Gated PixelCNN.
• Cells 29–32: datasets that encode every image on the fly into its code grids, and cross-entropy training of the two priors with Adam at learning rate $10^{-3}$, batch size 32 and 10 epochs each. Encoding should run with the VQ-VAE in evaluation mode, which cell 22 leaves it in; in training mode the quantizer would keep updating its codebook while the codes are being extracted.
• Cells 34–35: sampling: 16 sequential steps for the top grid, 64 for the bottom grid given the upsampled top, and decoding. Each step reruns the PixelCNN on the whole grid.
• Implementation note. utils.save_image(..., range=(-1, 1)) in cells 12, 22 and 35 uses the argument name that torchvision renamed to value_range; with current versions these calls need the new name.

The β-VAE has no trained example in the course notebooks: the VAE notebook defines the β-weighted loss in cell 24 but trains only with $\beta = 1$.

VAPNEV has no implementation in the course repository. The exercise notebook 03_hw_vapnev.ipynb of the same repository states the tasks on CelebA and leaves the code cells empty. Its only saved outputs are the device (mps), the shape of one image, $3 \times 64 \times 64$, and a grid of 25 random faces.

• Cells 3–9: data. CelebA through the course’s dataset class, with Resize(64) and CenterCrop(64), which keep the aspect ratio, and pixels mapped from $[0, 1]$ to $[-1, 1]$ by $2x - 1$.
• Cells 10–11: a normalizing flow on these faces, evaluated with FID and the negative log-likelihood. Any layers are allowed, coupling, autoregressive or linear, either convolutional or on the flattened image; the layers are described in flow architectures. A likelihood quoted for images must account for the dequantization and for the Jacobian of the preprocessing, as Agrawal and Dukkipati do for their logit transform.
• Cells 12–13: a VAE on the same faces, evaluated with FID, as in the VAE notebook.
• Cells 14–15: VAPNEV, following the paper, and a conditional VAPNEV. The notebook leaves the details of the condition to the task description; the natural choice is the 40 CelebA attributes, which the course’s dataset class returns with every image and the wrapper of cell 6 discards.
• Cells 16–17: an analysis of the latent spaces: which model interpolates better, encoding a face from outside CelebA and finding the nearest codes of training images, and the most interesting sample.

Next#

The VAE family trains its generator through a likelihood bound, with or without a discrete code, and pays for the bound with blur, with a separate prior or with a costly flow decoder. The next note takes the opposite route: a generator trained without any likelihood, judged only by a second network that tries to tell its samples from real data.

References#

• I. Higgins, L. Matthey, A. Pal, C. Burgess, X. Glorot, M. Botvinick, S. Mohamed, A. Lerchner. β-VAE: Learning basic visual concepts with a constrained variational framework. ICLR, 2017. Section 1: the definition of disentanglement; Figure 1: CelebA latent traversals.
• A. Alemi, B. Poole, I. Fischer, J. Dillon, R. A. Saurous, K. Murphy. Fixing a Broken ELBO. ICML, 2018. Rate and distortion.
• C. P. Burgess, I. Higgins, A. Pal, L. Matthey, N. Watters, G. Desjardins, A. Lerchner. Understanding disentangling in β-VAE. arXiv, 2018.
• F. Locatello, S. Bauer, M. Lucic, G. Rätsch, S. Gelly, B. Schölkopf, O. Bachem. Challenging Common Assumptions in the Unsupervised Learning of Disentangled Representations. ICML, 2019.
• A. van den Oord, O. Vinyals, K. Kavukcuoglu. Neural Discrete Representation Learning. NeurIPS, 2017. VQ-VAE; the EMA codebook update is in Appendix A.1.
• A. Razavi, A. van den Oord, O. Vinyals. Generating Diverse High-Fidelity Images with VQ-VAE-2. NeurIPS, 2019.
• Y. Bengio, N. Léonard, A. Courville. Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation. arXiv, 2013. The straight-through estimator.
• P. Dhariwal, H. Jun, C. Payne, J. W. Kim, A. Radford, I. Sutskever. Jukebox: A Generative Model for Music. arXiv, 2020. Random restarts of unused codes.
• J. Yu, X. Li, J. Y. Koh, H. Zhang, R. Pang, J. Qin, A. Ku, Y. Xu, J. Baldridge, Y. Wu. Vector-quantized Image Modeling with Improved VQGAN. ICLR, 2022. Factorized, normalized codebook lookup.
• S. Agrawal, A. Dukkipati. Deep Variational Inference Without Pixel-Wise Reconstruction. arXiv:1611.05209, 2016. VAPNEV; Section 3: the conditional likelihood and the conditional coupling layer; Section 4: preprocessing, architecture and results.
• L. Dinh, J. Sohl-Dickstein, S. Bengio. Density Estimation Using Real NVP. ICLR, 2017.
• S. R. Bowman, L. Vilnis, O. Vinyals, A. M. Dai, R. Jozefowicz, S. Bengio. Generating Sentences from a Continuous Space. CoNLL, 2016. Section 3.1: KL cost annealing.