REVIEW 3 major objections 5 minor 1 cited by
Distillation of a tractable model from the VQ-VAE
T0 review · 3 major / 5 minor · reviewed 2026-08-05 · deepseek-v4-flash
Pith's one-line read A VQ-VAE can be distilled, without retraining, into a tractable mixture model by keeping only its highest-probability latent codes.
desk verdict A useful but overclaimed distillation: uniform-weight mixtures over high-probability latents don't recover the VQ-VAE unless the prior is uniform; the paper needs to fix that or temper its central claim. read the letter →
The pith
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
The reading
What carries the argument
Index collapse: the VQ-VAE's latent prior assigns almost all its mass to a small subset of latent configurations, so the paper treats this known pathology as a computational asset by using high-probability latent variables as mixture components. The other load-bearing piece is the probabilistic circuit decomposition: the decoder factorizes per-pixel distributions (Assumption 1), making the sum-product structure smooth and decomposable, and beam search is the mechanism that finds the high-probability subset without exhaustive enumeration.
What would settle it
Train a VQ-VAE with a well-utilized codebook (e.g., using codebook resets or EMA updates) on a standard benchmark and measure the cumulative prior mass F(i): if reaching 99% of the mass requires a number of latent configurations N that makes evaluating p(x|z) for each component infeasible, the claimed conversion fails on that model. A second check: replace the factorized decoder with an autoregressive one and verify whether closed-form marginalization still holds; if not, Assumption 1 is doing the load-bearing work.
Extended reading notes
Core claim
The paper claims that because most VQ-VAE priors concentrate probability mass on a small fraction of the latent grid (index collapse), a faithful tractable model can be built by choosing N high-probability latent variables and forming the uniform mixture p_hat(x) = (1/N) * sum_{z in Zbar} p(x|z). Viewed as a probabilistic circuit with a factorized decoder, this distilled model supports exact marginalization and conditioning, so density estimation and image inpainting become closed-form operations. Beam search builds the subset Zbar in O(NHWK) without enumerating K^(HW); experiments on MNIST show the beam-searched mixture nearly matches an exhaustive random-sample mixture and answers conditio
Load-bearing premise
The whole construction stands on the VQ-VAE's prior putting almost all of its probability on a small slice of the latent space; if the codebook is well used, the needed subset grows until the mixture is no longer tractable, and the decoder must also factorize per-pixel, which the paper assumes rather than inherits.
Editorial extensions
If this is right
- A trained VQ-VAE becomes a probabilistic circuit: marginal likelihood, conditionals, expectations, and MAP queries are available in closed form, with no retraining.
- Beam search selects the mixture components in O(NHWK), so the conversion scales to latent grids where exhaustive enumeration is impossible.
- On MNIST, the distilled mixture comes close to exact enumeration, suggesting the approximation loss from dropping low-probability latent variables is modest.
- Because the method converts rather than trains, it keeps the VQ-VAE's expressive decoder while replacing its intractable prior with a tractable mixture.
- The distilled model supports conditional generation such as inpainting, with the paper reporting sample quality at least competitive with trained tractable baselines.
Reading between the lines
- The same distillation recipe should apply to any discrete latent model whose prior is concentrated, not only VQ-VAEs; the derivation only needs a scoreable autoregressive prior and a factorized decoder.
- Since Equation (2) uses uniform mixture weights rather than the prior probabilities p(z), reweighting components by p(z)—or by a diversity-aware score—is a natural extension the paper does not explore; its own latent-correlation appendix suggests such reweighting could improve density estimates.
- The practical window of the method is bounded by codebook utilization: if future VQ-VAEs eliminate index collapse, the subset needed for a given fidelity grows and the distillation stops being tractable—a testable prediction about the trajectory of the field.
- One could test whether selecting components by validation-set likelihood rather than by prior probability improves BPD, since the paper's beam search optimizes p(z) rather than the mixture's actual likelihood.
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. The paper proposes to convert a trained VQ-VAE into a tractable probabilistic circuit, called a distilled model (DM), by selecting a subset of latent configurations and forming an equal-weight mixture of the corresponding decoder distributions. Two selection strategies are considered: random sampling from the class-conditioned PixelCNN prior and beam search. On MNIST, the resulting DMs are evaluated for density estimation (bits per dimension) and image inpainting, with comparisons to continuous mixtures and Einsum networks. The main claim is that the DM preserves the expressiveness of the VQ-VAE while enabling exact probabilistic inference.
Significance. If the core claim holds, the paper offers a practical recipe for turning a trained discrete-latent deep generative model into a tractable probabilistic model without retraining, exploiting the empirically observed concentration of the latent prior. The connection between VQ-VAEs and probabilistic circuits is natural, and the experimental demonstration on MNIST is a useful proof-of-concept. However, the theoretical statement 'preserves expressiveness' is currently not established because of the uniform weighting in Eq. (2), and the paper's own Appendix F indicates that high-p(z) selection yields correlated components, which weakens the mixture interpretation. The work is therefore of interest but needs a substantial revision to make its claims accurate and to clarify the scope of applicability.
major comments (3)
- [Section 3, Eq. (2) and footnote 1] The uniform-mixture definition \hat p(x) = (1/|\bar Z|) \sum_{z \in \bar Z} p(x|z) does not recover the original VQ-VAE even when \bar Z = Z, unless p(z) is uniform. The footnote to Assumption 2 claiming recovery for N = K^{HW} is therefore false for the non-uniform PixelCNN prior used in the paper and shown in Figure 1. Consequently, the central claim that the DM 'preserves expressiveness of the VQ-VAE' does not follow from Eq. (2). The authors should either define prior-weighted components, e.g., \hat p(x) \propto \sum_{z \in \bar Z} p(z) p(x|z), which recovers Eq. (1) in the exhaustive limit, or explicitly reframe the contribution as an approximation and quantify the approximation gap.
- [Section 3, 'Random sampling' and Section 1] The text states that random sampling 'requires enumeration of all latent variables' and calls this exhaustive, yet DMRS results are reported for latent spaces of size 512^16 (DMRS4×4) and 1024^49 (DMRS7×7), where exhaustive enumeration is impossible. Sampling from an autoregressive prior p(z|c) does not require enumerating Z; it can be done sequentially. This inconsistency undermines the stated motivation for the beam-search variant. Please clarify what enumeration is actually needed for, and whether the reported DMRS results were obtained without enumeration.
- [Assumption 1 and Section 4] The tractability of the DM relies on Assumption 1, which requires the decoder to factorize as \prod_i p(x_i|z) with each factor a tractable distribution. Standard VQ-VAEs often use autoregressive decoders (e.g., PixelCNN), which do not satisfy this assumption. The paper trains custom VQ-VAEs with factorized Gaussian decoders. The scope of the distillation claim is therefore limited to VQ-VAEs that are already constrained in this way; the paper should state this limitation explicitly and avoid implying applicability to the general VQ-VAE architecture.
minor comments (5)
- [Section 4, Figure 2 and Figure 7] The text says results are averaged over 5 runs with different random seeds, but no error bars or confidence intervals are shown. Adding variability estimates would strengthen the comparison between DMRS, DMBS, and baseline models.
- [Section 4, 'Settings'] The number of components N used for the distilled models is not specified in the main text; it only appears indirectly through the beam width in Algorithm 1. Please state N (or the beam width and starting samples) for each DM configuration so that the parameter counts in Figure 2 are interpretable.
- [Appendix F] The appendix shows that selecting latents solely by highest p(z) yields correlated and redundant components, which is an important limitation for the beam-search strategy. This observation should be discussed in the main text rather than relegated to an appendix, as it directly affects the interpretation of the DMBS results.
- [Abstract and throughout] The term 'preserves expressiveness' is used repeatedly but never formally defined. Given the issues with Eq. (2), please either define a precise sense in which expressiveness is preserved or use a more cautious formulation such as 'approximates' or 'empirically matches'.
- [General typography] The paper contains inconsistent spacing in 'VQ-V AE' (e.g., in the title, abstract, and body). Please harmonize to 'VQ-VAE'.
Circularity Check
No significant circularity: Eq. (2) is a new uniform mixture evaluated on held-out data; the prior is used only for component selection.
full rationale
The paper's central construction is Eq. (2), a finite uniform mixture of factorized decoder conditionals. This is not derived from the VQ-VAE marginal (1) by algebraic manipulation; it is a new model defined explicitly, with the PixelCNN prior used only to choose which latent configurations enter the component set. Tractable inference follows directly from Assumption 1 (factorized conditionals) and Assumption 2 (small N), not from any equivalence with the original model. Density estimates are computed on held-out MNIST data, so no quantity is fitted and then re-reported as a prediction. There are no load-bearing self-citations: the cited prior work (van den Oord et al., Correia et al., Peharz et al.) is external and provides standard baselines. The one questionable statement is the footnote to Assumption 2 (Section 3) claiming that the original model (1) is recovered when |Zbar| = K^HW; this is mathematically inaccurate unless p(z) is uniform, since Eq. (2) places equal weight 1/|Zbar| on every component while (1) weights by p(z). But this is a correctness/soundness issue about the faithfulness of the distillation, not a circularity: the DM's validity as a PC and its empirical BPD figures do not depend on that recovery claim. Similarly, Appendix F's observation that selecting solely by highest p(z) yields correlated components is an acknowledged limitation of one selection strategy, not a circular argument. The paper is self-contained against external benchmarks, so the circularity score is 0.
Assumptions & free parameters
free parameters (2)
- Mixture size N / beam width B =
N not fixed in experiments; beam width B = N/(s|C|) with s=10
- Starting samples s =
10
assumptions (5)
- domain assumption The latent prior p(z) concentrates most probability on a tiny subset of Z (index collapse).
- domain assumption Per-pixel conditional independence and tractable per-pixel distributions (Assumption 1).
- domain assumption p(z) is autoregressive and can be scored and extended sequentially by beam search.
- ad hoc to paper Uniform mixture weights 1/|Zbar| in (2) are an acceptable approximation of the renormalized prior.
- standard math Standard probabilistic circuit definitions, smoothness and decomposability, are inherited.
Cite this review
Pith. "Pith review of Distillation of a tractable model from the VQ-VAE." pith.science (2026). https://pith.science/paper/XYNGT7GR
@misc{pith2026250901400,
author = {Pith},
title = {Pith review of: Distillation of a tractable model from the VQ-VAE},
year = {2026},
howpublished = {\url{https://pith.science/paper/XYNGT7GR}},
note = {Machine review of arXiv:2509.01400}
}
read the original abstract
Deep generative models with discrete latent space, such as the Vector-Quantized Variational Autoencoder (VQ-VAE), offer excellent data generation capabilities, but, due to the large size of their latent space, their probabilistic inference is deemed intractable. We demonstrate that the VQ-VAE can be distilled into a tractable model by selecting a subset of latent variables with high probabilities. This simple strategy is particularly efficient, especially if the VQ-VAE underutilizes its latent space, which is, indeed, very often the case. We frame the distilled model as a probabilistic circuit, and show that it preserves expressiveness of the VQ-VAE while providing tractable probabilistic inference. Experiments illustrate competitive performance in density estimation and conditional generation tasks, challenging the view of the VQ-VAE as an inherently intractable model.
Figures
Figures from the paper (4 more)
Forward citations
Cited by 1 Pith paper
-
PreScience: A Dataset and Benchmark for Scientific Forecasting
A new benchmark tests whether AI can forecast future scientific papers; frontier LLMs score ~5.6/10 on matching real abstracts, and simulated corpora are measurably less diverse and novel than human science.
Reference graph
Works this paper leans on
-
[1]
Semantic image synthesis with semantically coupled VQ - M odel, 2022
Stephan Alaniz, Thomas Hummel, and Zeynep Akata. Semantic image synthesis with semantically coupled VQ - M odel, 2022
work page 2022
-
[2]
4m-21: An any-to-any vision model for tens of tasks and modalities
Roman Bachmann, O g uzhan Fatih Kar, David Mizrahi, Ali Garjani, Mingfei Gao, David Griffiths, Jiaming Hu, Afshin Dehghan, and Amir Zamir. 4m-21: An any-to-any vision model for tens of tasks and modalities. In A. Globerson, L. Mackey, D. Belgrave, A. Fan, U. Paquet, J. Tomczak, and C. Zhang, editors, Advances in Neural Information Processing Systems, volu...
work page 2024
-
[3]
Probabilistic circuits: A unifying framework for tractable probabilistic models
YooJung Choi, Antonio Vergari, and Guy Van den Broeck. Probabilistic circuits: A unifying framework for tractable probabilistic models. October 2020
work page 2020
-
[4]
Correia, Gennaro Gala, Erik Quaeghebeur, Cassio De Campos, and Robert Peharz
Alvaro H.C. Correia, Gennaro Gala, Erik Quaeghebeur, Cassio De Campos, and Robert Peharz. Continuous mixtures of tractable probabilistic models. Proceedings of the AAAI Conference on Artificial Intelligence, 37 0 (6): 0 7244–7252, June 2023. ISSN 2159-5399. doi:10.1609/aaai.v37i6.25883
-
[5]
Efficient marginalization of discrete and structured latent variables via sparsity
Gon c alo Correia, Vlad Niculae, Wilker Aziz, and Andr\' e Martins. Efficient marginalization of discrete and structured latent variables via sparsity. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 11789--11802. Curran Associates, Inc., 2020
work page 2020
-
[6]
Adnan Darwiche and Pierre Marquis. A knowledge compilation map. Journal of Artificial Intelligence Research, 17: 0 229--264, 2002
work page 2002
-
[7]
Jukebox: A generative model for music, 2020
Prafulla Dhariwal, Heewoo Jun, Christine Payne, Jong Wook Kim, Alec Radford, and Ilya Sutskever. Jukebox: A generative model for music, 2020
work page 2020
-
[8]
Haohan Guo, Fenglong Xie, Dongchao Yang, Hui Lu, Xixin Wu, and Helen Meng. Addressing index collapse of large-codebook speech tokenizer with dual-decoding product-quantized variational auto-encoder, 2024
work page 2024
Show all 36 references
-
[9]
Straightening out the straight-through estimator: Overcoming optimization challenges in vector quantized networks, 2023
Minyoung Huh, Brian Cheung, Pulkit Agrawal, and Phillip Isola. Straightening out the straight-through estimator: Overcoming optimization challenges in vector quantized networks, 2023
2023
-
[10]
Adam: A method for stochastic optimization
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. International Conference on Learning Representations (ICLR), 2015
2015
-
[11]
Auto-encoding variational Bayes , 2022
Diederik P Kingma and Max Welling. Auto-encoding variational Bayes , 2022
2022
-
[12]
Lecun, L
Y. Lecun, L. Bottou, Y. Bengio, and P. Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86 0 (11): 0 2278--2324, 1998. doi:10.1109/5.726791
1998 doi
-
[13]
Scaling up probabilistic circuits by latent variable distillation
Anji Liu, Honghua Zhang, and Guy Van den Broeck. Scaling up probabilistic circuits by latent variable distillation. In 11th International Conference on Learning Representations, 2023
2023
-
[14]
Luong-Ha Nguyen and James-A. Goulet. Analytically tractable inference in deep neural networks, 2021
2021
-
[15]
Pytorch: An imperative style, high-performance deep learning library, 2019
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Yang, Zach DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fan...
2019
-
[16]
On the latent variable interpretation in sum-product networks
Robert Peharz, Robert Gens, Franz Pernkopf, and Pedro Domingos. On the latent variable interpretation in sum-product networks. IEEE transactions on pattern analysis and machine intelligence, 39 0 (10): 0 2030--2044, 2016
-
[17]
Einsum networks: Fast and scalable learning of tractable probabilistic circuits
Robert Peharz, Steven Lang, Antonio Vergari, Karl Stelzner, Alejandro Molina, Martin Trapp, Guy Van Den Broeck, Kristian Kersting, and Zoubin Ghahramani. Einsum networks: Fast and scalable learning of tractable probabilistic circuits. In Hal Daumé III and Aarti Singh, editors,...
2020
-
[18]
Generating diverse structure for image inpainting with hierarchical VQ-VAE , 2021
Jialun Peng, Dong Liu, Songcen Xu, and Houqiang Li. Generating diverse structure for image inpainting with hierarchical VQ-VAE , 2021
2021
-
[19]
Sum-product networks: A new deep architecture, 2012
Hoifung Poon and Pedro Domingos. Sum-product networks: A new deep architecture, 2012
2012
-
[20]
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 Eric P. Xing and Tony Jebara, editors, Proceedings of the 31st International Conference on Machine Learning, volume 32 of Proceedings o...
2014
-
[21]
Theory and experiments on vector quantized autoencoders, 2018
Aurko Roy, Ashish Vaswani, Arvind Neelakantan, and Niki Parmar. Theory and experiments on vector quantized autoencoders, 2018
2018
-
[22]
Generating high-quality and informative conversation responses with sequence-to-sequence models, 2017
Louis Shao, Stephan Gouws, Denny Britz, Anna Goldie, Brian Strope, and Ray Kurzweil. Generating high-quality and informative conversation responses with sequence-to-sequence models, 2017
2017
-
[23]
Probabilistic flow circuits: Towards unified deep models for tractable probabilistic inference
Sahil Sidheekh, Kristian Kersting, and Sriraam Natarajan. Probabilistic flow circuits: Towards unified deep models for tractable probabilistic inference. In Robin J. Evans and Ilya Shpitser, editors, Proceedings of the Thirty-Ninth Conference on Uncertainty in Artificial Intel...
1964
-
[24]
SQ-VAE : Variational Bayes on discrete representation with self-annealed stochastic quantization, 2022
Yuhta Takida, Takashi Shibuya, WeiHsiang Liao, Chieh-Hsin Lai, Junki Ohmura, Toshimitsu Uesaka, Naoki Murata, Shusuke Takahashi, Toshiyuki Kumakura, and Yuki Mitsufuji. SQ-VAE : Variational Bayes on discrete representation with self-annealed stochastic quantization, 2022
2022
-
[25]
A note on the evaluation of generative models, 2016
Lucas Theis, Aäron van den Oord, and Matthias Bethge. A note on the evaluation of generative models, 2016
2016
-
[26]
Visual autoregressive modeling: Scalable image generation via next-scale prediction
Keyu Tian, Yi Jiang, Zehuan Yuan, Bingyue Peng, and Liwei Wang. Visual autoregressive modeling: Scalable image generation via next-scale prediction. In A. Globerson, L. Mackey, D. Belgrave, A. Fan, U. Paquet, J. Tomczak, and C. Zhang, editors, Advances in Neural Information Pr...
2024
-
[27]
Pixel recurrent neural networks, 2016 a
Aaron van den Oord, Nal Kalchbrenner, and Koray Kavukcuoglu. Pixel recurrent neural networks, 2016 a
2016
-
[28]
Conditional image generation with P ixel CNN decoders, 2016 b
Aaron van den Oord, Nal Kalchbrenner, Oriol Vinyals, Lasse Espeholt, Alex Graves, and Koray Kavukcuoglu. Conditional image generation with P ixel CNN decoders, 2016 b
2016
-
[29]
Neural discrete representation learning
Aaron van den Oord, Oriol Vinyals, and Koray Kavukcuoglu. Neural discrete representation learning. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 30. Curran Asso...
2017
-
[30]
Tractable probabilistic models: Representations, algorithms, learning, and applications
A Vergari, Nicola Di Mauro, and G Van den Broeck. Tractable probabilistic models: Representations, algorithms, learning, and applications. Tutorial at the 35th Conference on Uncertainty in Artificial Intelligence (UAI), 2019
2019
-
[31]
Simplifying, regularizing and strengthening sum-product network structure learning
Antonio Vergari, Nicola Di Mauro, and Floriana Esposito. Simplifying, regularizing and strengthening sum-product network structure learning. In Annalisa Appice, Pedro Pereira Rodrigues, V \'i tor Santos Costa, Jo \ a o Gama, Al \'i pio Jorge, and Carlos Soares, editors, Machin...
2015
-
[32]
Learning linear ranking functions for beam search with application to planning
Yuehua Xu, Alan Fern, and Sungwook Yoon. Learning linear ranking functions for beam search with application to planning. Journal of Machine Learning Research, 10 0 (7), 2009
2009
-
[33]
V ideo GPT : Video generation using VQ-VAE and transformers, 2021
Wilson Yan, Yunzhi Zhang, Pieter Abbeel, and Aravind Srinivas. V ideo GPT : Video generation using VQ-VAE and transformers, 2021
2021
-
[34]
Soundstream: An end-to-end neural audio codec, 2021
Neil Zeghidour, Alejandro Luebs, Ahmed Omran, Jan Skoglund, and Marco Tagliasacchi. Soundstream: An end-to-end neural audio codec, 2021
2021
-
[35]
A unified approach for learning the parameters of sum-product networks
Han Zhao, Pascal Poupart, and Geoffrey J Gordon. A unified approach for learning the parameters of sum-product networks. In D. Lee, M. Sugiyama, U. Luxburg, I. Guyon, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 29. Curran Associates, Inc., 2016
2016
-
[36]
Dolfing, Sameer Khurana, Tanel Alumäe, and Antoine Laurent
Adrian Łańcucki, Jan Chorowski, Guillaume Sanchez, Ricard Marxer, Nanxin Chen, Hans J.G.A. Dolfing, Sameer Khurana, Tanel Alumäe, and Antoine Laurent. Robust training of vector quantized bottleneck models. In 2020 International Joint Conference on Neural Networks (IJCNN), page...
2020
Reviewed August 5, 2026 · model on record in the stance chip above.
Discussion (0). Continue with ORCID to comment.