Pith. sign in

REVIEW 2 major objections 4 minor 23 references

An Introduction to Discrete Variational Autoencoders

T0 review · 2 major / 4 minor · reviewed 2026-08-15 · deepseek-v4-flash

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗ Review JSON

Pith's one-line read A discrete VAE with independent categorical latents can be trained from first principles by stochastic gradient ascent on the ELBO, using a REINFORCE-style encoder gradient plus an entropy term and a binary cross-entropy decoder gradient.

desk verdict A competent, readable tutorial on discrete VAEs that re-derives standard results; the notation has a BCE argument-order slip and the training recipe is untested, but it is a legitimate pedagogical contribution. read the letter →

arxiv 2505.10344 v1 pith:BQRNQF6B submitted 2025-05-15 cs.LG

Alan Jeffares , Liyuan Liu This is my paper · ORCID
classification cs.LG
keywords discretevariationalautoencodercategoricallatentvariablesevidencelowerboundlog-derivativetrickREINFORCEbinarycross-entropystochasticgradientascentunsupervisedlearning
verification ladder T0 review T1 audit T2 compute T3 formal

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The reading

This paper sets out to show that the canonical discrete variational autoencoder—a latent space made of independent categorical variables—can be derived and trained entirely from first principles, without the Gaussian assumptions that dominate VAE tutorials. The central claim is concrete: maximizing the evidence lower bound yields an encoder gradient that splits into an entropy term plus a one-sample log-derivative (REINFORCE) term, and a decoder gradient that is exactly negative binary cross-entropy. This matters because discrete latents are a natural fit for text, sparse features, and other modalities, yet most pedagogical treatments stop at the continuous case. A sympathetic reader comes away with a self-contained training recipe and a working mental model of why discrete VAEs are harder than Gaussian ones: the sampling step blocks backpropagation, and the log-derivative trick is the workaround.

What carries the argument

The load-bearing mechanism is the log-derivative trick, also called the score-function estimator or REINFORCE: $\nabla_\psi \mathbb{E}_{p_\psi(\mathbf{x})}[f(\mathbf{x})] = \mathbb{E}_{p_\psi(\mathbf{x})}[f(\mathbf{x}) \nabla_\psi \log p_\psi(\mathbf{x})]$. This identity moves the gradient inside the expectation so a single Monte Carlo sample yields an unbiased gradient estimate even though sampling from a categorical distribution is not differentiable. It converts the intractable expectation over discrete latents in Eq. (45) into the concrete term in Eq. (81), and it combines with the simplification of the negative KL term into an entropy plus the constant $-D \log K$.

What would settle it

Take a small random encoder and decoder, compute the ELBO, and compare automatic differentiation of the ELBO against the formulas in Eqs. (84) and (85); any mismatch would refute the derivation.

Watch

Extended reading notes

Core claim

The paper claims that for a discrete VAE in which the latent space $\mathbf{z} \in \{0,1\}^{D \times K}$ consists of $D$ independent categorical variables, each with $K$ categories, a uniform prior $p(\mathbf{z})$, and a Bernoulli likelihood on binarized inputs, the ELBO gradients take a closed, computable form. With $f_\phi$ the encoder and $g_\theta$ the decoder, the decoder gradient is $\hat\nabla_\theta \approx -\nabla_\theta \mathrm{BCE}(g_\theta(\mathbf{z}), \mathbf{x})$, while the encoder gradient is $\hat\nabla_\phi \approx \nabla_\phi \mathrm{Entropy}(f_\phi(\mathbf{x})) - \mathrm{BCE}(g_\theta(\mathbf{z}), \mathbf{x}) \nabla_\phi \sum_{d=1}^{D} \log f_\phi(\mathbf{x})^{(d)}_{k^{(d)}}$, where $k^{(d)}$ is the sampled category for latent dimension $d$. The second term is a one-sample REINFORCE estimate obtained via the log-derivative trick, and the whole recipe maximizes the ELBO, whose per-sample value is estimated as $\mathrm{Entropy}(f_\phi(\mathbf{x})) - D \log K - \mathrm{BCE}(g_\theta(\mathbf{z}), \mathbf{x})$.

Load-bearing premise

The recipe assumes that a single random sample gives a reliable estimate of the encoder gradient, even though that estimate can have high variance without any variance-reduction technique.

Editorial extensions

If this is right

  • A discrete VAE with independent categorical latents and Bernoulli decoder can be trained with standard stochastic gradient ascent; no reparameterization, Gumbel noise, or control variate is required for the central recipe.
  • The same ELBO estimate used for training, $\mathrm{Entropy}(f_\phi(\mathbf{x})) - D \log K - \mathrm{BCE}(g_\theta(\mathbf{z}), \mathbf{x})$, doubles as a monitoring and model-comparison score during training.
  • The constant $-D \log K$ drops out of the encoder gradient but remains in the ELBO, so architectures with different latent dimensionality or category counts can be compared on the same scale only if this term is included.
  • Because every step is derived from first principles, the recipe transfers to any data modality where a Bernoulli or categorical likelihood is appropriate, such as binarized images or token-like discrete observations.
  • The tutorial positions this simple log-derivative recipe as the baseline that more sophisticated discrete gradient estimators—Gumbel-softmax, straight-through, control variates—are trying to improve on.

Reading between the lines

Editorial extensions of the paper, not claims the author makes directly.

  • Beyond the paper, the variance of the one-sample REINFORCE term in Eq. (81) is the natural failure point; adding a simple baseline (such as the mean reconstruction loss) would preserve unbiasedness and likely stabilize training, which the paper leaves implicit.
  • A testable extension is to compare this recipe head-to-head against Gumbel-softmax and straight-through estimators on identical architectures and data, measuring both ELBO and gradient variance; the paper's derivation suggests the simple recipe should win in simplicity but may lose in sample efficiency.
  • The entropy term in Eq. (70) acts as a built-in regularizer that pushes the encoder toward uniform categorical distributions; this implicit connection to rate-distortion tradeoffs in discrete representation learning is not discussed in the paper.
Share X Bluesky LinkedIn Reddit HN

Signed reviews

No signed human review yet.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, and a circularity audit.

Referee Report

2 major / 4 minor

Summary. The manuscript is a tutorial that derives, from first principles, the evidence lower bound (ELBO) for a variational autoencoder with independent categorical latent variables, then shows that the decoder gradient reduces to a binary cross-entropy term and the encoder gradient splits into an entropy-gradient term plus a REINFORCE-style term. It concludes with a concrete stochastic-gradient-ascent training recipe (Algorithm 1) and points to a PyTorch implementation. The derivation is self-contained, starting from probability basics and the log-derivative trick, and it does not rely on any external gradient estimators as an input.

Significance. The paper's pedagogical value is real. It gives a careful, step-by-step derivation that many students and practitioners will find accessible, and it correctly identifies the log-derivative trick as the central technical obstacle for discrete latents. The final gradient expressions in Eqs. (84)-(85) are standard but are consolidated here in one place, which is useful for teaching. The derivation is internally consistent apart from the binary cross-entropy argument-order error discussed below. That error is fixable, but it is load-bearing because it appears in the final gradient updates and in the ELBO estimate.

major comments (2)
  1. [Eq. (8) vs. Eqs. (57)-(58), (74)-(75), (83), (86)] The definition of BCE in Eq. (8) is BCE(a,b) = -a log b - (1-a) log(1-b), i.e., the first argument is the target and the second is the predicted probability. However, the derivation in Eqs. (55)-(56) produces a reconstruction term of the form sum_p [x^{(p)} log g_theta(z)^{(p)} + (1-x^{(p)}) log(1-g_theta(z)^{(p)})], which equals -BCE(x, g_theta(z)) under this definition. The manuscript instead writes -BCE(g_theta(z), x) in Eqs. (57)-(58), (74)-(75), (83), and (86). Since BCE is not symmetric, these expressions are not equivalent; for example, -BCE(g_theta(z), x) contains log x terms and would lead to an incorrect gradient with respect to theta in Eq. (85). This is a load-bearing notational error: a reader implementing Eq. (85) literally with the definition of Eq. (8) would optimize the wrong objective. Please swap the arguments consistently to BCE(x, g_theta(z)), or alternatively redefine BCE in Eq. (8) to use prediction-first order and state this convention explicitly.
  2. [Section 8, Eq. (81) and Algorithm 1] The training recipe relies on a one-sample REINFORCE estimator with no baseline or control variate. The paper itself notes in Eq. (73) that this estimator is 'potentially high variance,' but it provides no empirical evidence that the proposed recipe actually converges. There are no training curves, no final ELBO values, no reconstruction examples, and no comparison on binarized MNIST, despite MNIST being used as the running example throughout. The variance of the estimator scales with the magnitude of the reconstruction reward and with the number of latent dimensions D; for high-dimensional images this can easily overwhelm the deterministic entropy gradient in Eq. (70). Since the abstract and Section 8 claim a 'concrete training recipe,' I ask that the authors either (a) include a small experimental section demonstrating that the recipe trains successfully and reporting the gradient estimator's variance, or (b) explicitly frame the contribution as a derivation-only tutorial and note that the vanilla one-sample REINFORCE recipe may require variance-reduction techniques (e.g., baselines) for practical use.
minor comments (4)
  1. [Section 7, Eq. (54)] In Eq. (54), the second term uses x^{(d)} where the context and surrounding equations require x^{(p)}; this looks like a typo that should be corrected.
  2. [Sections 3 and 6-8] The parameter notation for the encoder and decoder is swapped: Section 3 and Figure 1 use f_theta for the encoder and g_phi for the decoder, while Sections 6 through 8 use f_phi and g_theta. This reversal is confusing and should be normalized throughout.
  3. [Section 7, Eq. (79)] The Iverson bracket [z^{(d)}_k = k] is ill-formed because k is used both as the category index and as the value being compared. It should be written as [z^{(d)} = k] or, equivalently, using the sampled category index k^{(d)}.
  4. [Algorithm 1] The stopping criterion 'while L_ELBO(D_val) not converged' is vague; specifying a tolerance or a maximum number of epochs would make the recipe more actionable.

Circularity Check

0 steps flagged · score 0.0 of 10

No circularity: the discrete VAE derivation is self-contained and does not reduce to its inputs.

full rationale

The paper is a tutorial that derives the discrete VAE ELBO and its gradients from first principles. Eq. (37) follows from the log-likelihood decomposition and non-negativity of the KL divergence; Eq. (43) splits the ELBO into a KL term and a reconstruction term. The decoder gradient in Eq. (58) follows by substituting the Bernoulli probability mass function into the decoder likelihood, Eq. (70) follows from the categorical KL divergence against a uniform prior, and Eq. (81) follows from the log-derivative trick in Eq. (26). No parameter is fitted to a target quantity and then reported as a prediction; the final training recipe in Eqs. (84)-(85) is the algebraic combination of these derived gradients. Self-citations [LDL+23] and [LKW+24] appear only in the introduction as pointers to related work on gradient estimation and mixture-of-experts layers, and they are not load-bearing for any derivation step. The paper explicitly acknowledges that the one-sample REINFORCE estimator is 'potentially high variance' (Eq. 73), and the absence of empirical convergence evidence is a practical risk about whether the recipe trains reliably, not a circularity in the derivation.

Assumptions & free parameters 0 free parameters · 7 assumptions · 0 invented entities

The tutorial contains no fitted free parameters and no invented entities. It relies on standard probability identities, the non-negativity of KL divergence, Monte Carlo sampling, the log-derivative trick, and explicit modeling assumptions about uniform categorical priors and Bernoulli likelihoods. These are reasonable assumptions but are not empirically tested in the paper.

assumptions (7)
  • standard math KL divergence is non-negative and zero iff the two distributions match.
    Used in Section 4 (Eq 35) to turn the exact likelihood decomposition into an ELBO.
  • standard math Law of total probability and conditional probability factorization.
    Used in Sections 2.2 and 4 to derive p_theta(x) = integral p_theta(x,z) dz and the ELBO.
  • standard math Monte Carlo samples provide unbiased gradient estimates.
    Used in Section 2.5 and throughout Section 7 to replace expectations with single samples.
  • standard math Log-derivative identity: p_psi(x) * grad_psi log p_psi(x) = grad_psi p_psi(x).
    Derived in Section 2.6 and used in Eq (72) for the encoder gradient.
  • domain assumption Prior over latents is uniform categorical; posterior approximation is a factorized categorical distribution.
    Stated in Section 6; produces the Entropy(f_phi(x)) - D log K term and the sampled log-probability gradient.
  • domain assumption Binarized input pixels with independent Bernoulli likelihood, using the continuous-Bernoulli caveat.
    Stated in Section 6 with reference [LGC19]; required to write the decoder loss as binary cross-entropy.
  • domain assumption Single-sample Monte Carlo approximation is adequate for training and evaluation.
    Adopted in Section 7 Eq (83) and Algorithm 1; not empirically validated in the paper.

how reviews work

0 comments
Cite this review

Pith. "Pith review of An Introduction to Discrete Variational Autoencoders." pith.science (2026). https://pith.science/paper/BQRNQF6B

@misc{pith2026250510344,
  author       = {Pith},
  title        = {Pith review of: An Introduction to Discrete Variational Autoencoders},
  year         = {2026},
  howpublished = {\url{https://pith.science/paper/BQRNQF6B}},
  note         = {Machine review of arXiv:2505.10344}
}
read the original abstract

Variational Autoencoders (VAEs) are well-established as a principled approach to probabilistic unsupervised learning with neural networks. Typically, an encoder network defines the parameters of a Gaussian distributed latent space from which we can sample and pass realizations to a decoder network. This model is trained to reconstruct its inputs and is optimized through the evidence lower bound. In recent years, discrete latent spaces have grown in popularity, suggesting that they may be a natural choice for many data modalities (e.g. text). In this tutorial, we provide a rigorous, yet practical, introduction to discrete variational autoencoders -- specifically, VAEs in which the latent space is made up of latent variables that follow a categorical distribution. We assume only a basic mathematical background with which we carefully derive each step from first principles. From there, we develop a concrete training recipe and provide an example implementation, hosted at https://github.com/alanjeffares/discreteVAE.

Figures

Figures reproduced from arXiv: 2505.10344 by the authors.

Figure 1
Figure 1. An Autoencoder. The encoder network fθ maps to a latent representation z which is decoded by a second network, gϕ, attempting to reproduce the original input. broad practical applications in a deep learning context (e.g. data generation and anomaly detection) where it has been highly influential. The network is conceptually subdivided into an encoder and decoder around the latent representation, which act as the com… view at source ↗
Figure 2
Figure 2. A Variational Autoencoder. Unlike the vanilla autoencoder, the encoder of this variational version outputs the parameters of a probability distribution from which we sample and decode. The parameters are learned by maximizing the ELBO which provides a lower bound on the log-likelihood. least in theory, we can marginalize over these latent variables to make their contribution explicit. pθ(x) = Z pθ(x, z)dz. (28) In p… view at source ↗
Figure 3
Figure 3. The Discrete Variational Autoencoder. The input is encoded into the parameters of D categorical distributions; then, we sample one of the K categories from each. These samples are passed back through the decoder that attempts to reproduce the original input. model using a neural network fϕ(x) : X P → ∆D×K where X denotes some arbitrary input space and ∆ denotes a space of probabilities such that ∆ := [0, 1]D×K where… view at source ↗

Discussion (0). Continue with ORCID to comment.

Reference graph

Works this paper leans on

23 extracted references · 12 canonical work pages

  1. [1]

    Deep learning: Foundations and concepts

    Christopher M Bishop and Hugh Bishop. Deep learning: Foundations and concepts . Springer Nature, 2023

    work page 2023

  2. [2]

    Estimating or propagating gradients through stochastic neurons for conditional computation

    Yoshua Bengio, Nicholas L \'e onard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv:1308.3432 , 2013

    arXiv 2013

  3. [3]

    The helmholtz machine

    Peter Dayan, Geoffrey E Hinton, Radford M Neal, and Richard S Zemel. The helmholtz machine. Neural computation , 7(5):889--904, 1995

    work page 1995

  4. [4]

    Tutorial on variational autoencoders

    Carl Doersch. Tutorial on variational autoencoders. arXiv preprint arXiv:1606.05908 , 2016

    arXiv 2016

  5. [5]

    Backpropagation through the void: Optimizing control variates for black-box gradient estimation

    Will Grathwohl, Dami Choi, Yuhuai Wu, Geoff Roeder, and David Duvenaud. Backpropagation through the void: Optimizing control variates for black-box gradient estimation. In International Conference on Learning Representations , 2018

    work page 2018

  6. [6]

    Factor analysis, probabilistic principal component analysis, variational inference, and variational autoencoder: Tutorial and survey

    Benyamin Ghojogh, Ali Ghodsi, Fakhri Karray, and Mark Crowley. Factor analysis, probabilistic principal component analysis, variational inference, and variational autoencoder: Tutorial and survey. arXiv preprint arXiv:2101.00734 , 2021

    arXiv 2021

  7. [7]

    Scaling and evaluating sparse autoencoders

    Leo Gao, Tom Dupr \'e la Tour, Henk Tillman, Gabriel Goh, Rajan Troll, Alec Radford, Ilya Sutskever, Jan Leike, and Jeffrey Wu. Scaling and evaluating sparse autoencoders. arXiv preprint arXiv:2406.04093 , 2024

    arXiv 2024

  8. [8]

    Categorical reparameterization with gumbel-softmax

    Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmax. In International Conference on Learning Representations , 2017

    2017

Show all 23 references
  1. [9]

    Nonlinear principal component analysis using autoassociative neural networks

    Mark A Kramer. Nonlinear principal component analysis using autoassociative neural networks. AIChE journal , 37(2):233--243, 1991

    1991

  2. [10]

    Kingma and Max Welling

    Diederik P. Kingma and Max Welling. Auto-encoding variational bayes. In Yoshua Bengio and Yann LeCun, editors, 2nd International Conference on Learning Representations, ICLR 2014, Banff, AB, Canada, April 14-16, 2014, Conference Track Proceedings , 2014

    2014

  3. [11]

    An introduction to variational autoencoders

    Diederik P Kingma, Max Welling, et al. An introduction to variational autoencoders. Foundations and Trends in Machine Learning , 12(4):307--392, 2019

    2019

  4. [12]

    Bridging discrete and backpropagation: Straight-through and beyond

    Liyuan Liu, Chengyu Dong, Xiaodong Liu, Bin Yu, and Jianfeng Gao. Bridging discrete and backpropagation: Straight-through and beyond. In Thirty-seventh Conference on Neural Information Processing Systems , 2023

    2023

  5. [13]

    The continuous bernoulli: fixing a pervasive error in variational autoencoders

    Gabriel Loaiza-Ganem and John P Cunningham. The continuous bernoulli: fixing a pervasive error in variational autoencoders. Advances in Neural Information Processing Systems , 32, 2019

    2019

  6. [14]

    Grin: Gradient-informed moe

    Liyuan Liu, Young Jin Kim, Shuohang Wang, Chen Liang, Yelong Shen, Hao Cheng, Xiaodong Liu, Masahiro Tanaka, Xiaoxia Wu, Wenxiang Hu, et al. Grin: Gradient-informed moe. arXiv preprint arXiv:2409.12136 , 2024

    2024 arXiv

  7. [15]

    Gemma scope: Open sparse autoencoders everywhere all at once on gemma 2

    Tom Lieberum, Senthooran Rajamanoharan, Arthur Conmy, Lewis Smith, Nicolas Sonnerat, Vikrant Varma, J \'a nos Kram \'a r, Anca Dragan, Rohin Shah, and Neel Nanda. Gemma scope: Open sparse autoencoders everywhere all at once on gemma 2. arXiv preprint arXiv:2408.05147 , 2024

    2024 arXiv

  8. [16]

    Probabilistic machine learning: Advanced topics

    Kevin P Murphy. Probabilistic machine learning: Advanced topics . MIT press, 2023

    2023

  9. [17]

    Tutorial: Deriving the standard variational autoencoder (vae) loss function

    Stephen Odaibo. Tutorial: Deriving the standard variational autoencoder (vae) loss function. arXiv preprint arXiv:1907.08956 , 2019

    1907 arXiv

  10. [18]

    Understanding deep learning

    Simon JD Prince. Understanding deep learning . MIT press, 2023

    2023

  11. [19]

    Stochastic backpropagation and approximate inference in deep generative models

    Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In International conference on machine learning , pages 1278--1286. PMLR, 2014

    2014

  12. [20]

    Discrete variational autoencoders

    Jason Tyler Rolfe. Discrete variational autoencoders. In International Conference on Learning Representations , 2017

    2017

  13. [21]

    Generating diverse high-fidelity images with vq-vae-2

    Ali Razavi, Aaron Van den Oord, and Oriol Vinyals. Generating diverse high-fidelity images with vq-vae-2. Advances in neural information processing systems , 32, 2019

    2019

  14. [22]

    Neural discrete representation learning

    Aaron Van Den Oord, Oriol Vinyals, et al. Neural discrete representation learning. Advances in neural information processing systems , 30, 2017

    2017

  15. [23]

    Dvae++: Discrete variational autoencoders with overlapping transformations

    Arash Vahdat, William Macready, Zhengbing Bian, Amir Khoshaman, and Evgeny Andriyash. Dvae++: Discrete variational autoencoders with overlapping transformations. In International conference on machine learning , pages 5035--5044. PMLR, 2018

    2018

Pith tools

Reviewed August 15, 2026 · model on record in the stance chip above.