Discrete Latent VAEs

This note records the training-specific considerations that arise when the latent variable is discrete rather than continuous.

Related notes:

• CVAE model components
• CVAE objective and loss terms
• CVAE parameterization

A discrete-state VAE uses categorical latent variables instead of Gaussian latents.

Model and parameterization

Given data $x$ or conditional context $c$:

• prior: $p(z)$, often uniform categorical, or conditional prior $p_\phi(z \mid c)$
• encoder / inference model: logits $\ell_\psi(x) \in \mathbb{R}^K$ with $$q_\psi(z = k \mid x) = \operatorname{softmax}(\ell_\psi(x))_k$$
• decoder: $p_\theta(x \mid z)$ or $p_\theta(y \mid x, z)$, with $z$ provided through an embedding or other conditioning interface

ELBO objective

The negative ELBO is

$$\mathcal{L}(\theta, \psi) = - \mathbb{E}_{z \sim q_\psi(\cdot \mid x)}[\log p_\theta(x \mid z)] + \beta \, \mathrm{KL}\!\left(q_\psi(z \mid x)\,\|\,p(z)\right).$$

In the conditional case, replace $p(z)$ with a conditional prior such as $p_\phi(z \mid c)$ and condition the decoder on $c$ as well.

KL for categorical latents

For categorical $q$ and $p$ with probabilities $q_k$ and $p_k$,

$$\mathrm{KL}(q, p) = \sum_{k=1}^{K} q_k \log \frac{q_k}{p_k}.$$

This term is analytic, so no sampling is needed for the KL itself.

Reconstruction expectation

The main practical issue is the reconstruction expectation

$$\mathbb{E}_{z \sim q_\psi(\cdot \mid x)}[\log p_\theta(x \mid z)].$$

Two main regimes matter:

1. Exact marginalization for small $K$: $$\mathbb{E}_{z}[\log p_\theta(x \mid z)] = \sum_{k=1}^{K} q_k \log p_\theta(x \mid z = k).$$ This is low variance, unbiased, and differentiable.
2. Monte Carlo for large $K$: $$\mathbb{E}_{z}[\log p_\theta(x \mid z)] \approx \frac{1}{S} \sum_{s=1}^{S} \log p_\theta(x \mid z^{(s)}).$$ This is straightforward for decoder parameters but raises the usual discrete-gradient issue for the encoder.

Differentiating through discrete sampling

Sampling a categorical latent is not pathwise differentiable. Common options are:

• REINFORCE / score-function estimators: $$\nabla_\psi \mathbb{E}_{z}[f(z)] = \mathbb{E}_{z}\!\left[f(z)\nabla_\psi \log q_\psi(z \mid x)\right]$$ often with a baseline for variance reduction
• Gumbel-Softmax / Concrete relaxations: use a soft one-hot approximation in the backward pass
• control-variate estimators such as VIMCO, RELAX, or REBAR: more complex, but potentially lower variance

Why this note matters here

The main methodology cluster already covers the discrete-latent interfaces. What this note adds is the training-specific reminder that discrete latents change the reconstruction expectation and gradient-estimation story, even when the high-level CVAE factorization looks the same.