Unbiased Learning of Deep Generative Models with Structured Discrete Representations

By composing graphical models with deep learning architectures, we learn generative models with the strengths of both frameworks. The structured variational autoencoder (SVAE) inherits structure and interpretability from graphical models, and flexible likelihoods for high-dimensional data from deep learning, but poses substantial optimization challenges. We propose novel algorithms for learning SVAEs, and are the first to demonstrate the SVAE's ability to handle multimodal uncertainty when data is missing by incorporating discrete latent variables. Our memory-efficient implicit differentiation scheme makes the SVAE tractable to learn via gradient descent, while demonstrating robustness to incomplete optimization. To more rapidly learn accurate graphical model parameters, we derive a method for computing natural gradients without manual derivations, which avoids biases found in prior work. These optimization innovations enable the first comparisons of the SVAE to state-of-the-art time series models, where the SVAE performs competitively while learning interpretable and structured discrete data representations.

Paper

Similar papers

Peer review

Reviewer gyLw4/10 · confidence 4/52023-07-04

Summary

The authors revisit learning algorithms for structured variational autoencoders. This paper proposes using the implicit function theorem to compute the implicit gradients, improving stability and memory efficiency. The authors further revisit the natural gradients studied in the original SVAE paper and propose an unbiased version. The experimentation is carried out on human pose capture and audio spectrogram datasets. The evaluation is based on modelling performances (likelihood), visual quality (FID score) and missing data imputation.

Strengths

1. This paper is impressive by the depth and breadth of its theory (implicit gradients, unbiased natural gradients) 2. A complete and well-detailed appendix featuring all derivations, although I was not in a position to proof-read all of it 3. The experiment section shows a clear improvement of `SVAE-SLDS` (the method introduced in this paper) over the baseline `SVAE-SLDS-Bias` based on likelihood performances and FID

Weaknesses

1. This paper lacks structure and clarity. More work is needed to synthesize and popularize the main contributions. 2. This paper attempts to make too many contributions and therefore fails to present and support each of them sufficiently. 3. The experiment section is not sufficient (only two rather small datasets). 4. Section 4 claims proposing a more memory efficient algorithm: an experiment demonstrating the applicability to larger-scale datasets is missing. 5. Section 3.2 seems to indicate that discrete latent variable models can only be optimized by reparameterizing discrete distributions. This is wrong, score function estimators (e.g. Reinforce) are viable alternatives. ["Monte Carlo Gradient Estimation in Machine Learning", Mohamed et al., 2020](https://jmlr.org/papers/volume21/19-346/19-346.pdf) gives a great overview of gradient estimation methods and should, in my opinion, be discussed in the related work section. Reparameterization-free methods exist and work decently. For instance: - ["Variational Inference for Monte Carlo Objectives", Mnih et al., 2016](http://proceedings.mlr.press/v48/mnihb16.pdf) - ["Revisiting Reweighted Wake-Sleep for Models with Stochastic Control Flow", Anh et al, 2018](https://arxiv.org/abs/1805.10469) - ["The Thermodynamic Variational Objective", Masrani et al., 2019](https://arxiv.org/abs/1907.00031) - ["Optimal Variance Control of the Score Function Gradient Estimator for Importance Weighted Bounds", Lievin et al., 2020](https://arxiv.org/abs/2008.01998) - ["Gradient Estimation with Discrete Stein Operators", Shi et al., 2022](https://arxiv.org/abs/2202.09497)

Questions

1. In the original SVAE paper, the authors explore the case of a simple latent GMM model, why not explore this simple case in this paper? 2. Would your method scale to larger datasets as explored (for instance) in ["Clockwork Variational Autoencoders", Saxena et al., 2021]? 3. Could you please detail how the FID scores are computed?

Rating

4: Borderline reject: Technically solid paper where reasons to reject, e.g., limited evaluation, outweigh reasons to accept, e.g., good evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

2 fair

Presentation

2 fair

Contribution

2 fair

Limitations

Technical limitations are not clearly discussed.

Reviewer FFkU5/10 · confidence 4/52023-07-05

Summary

This work presents a novel implicit optimisation approach for Structured Variational Autoencoders (SVAEs) allowing for the computation of natural gradients in an automatic way. This allows for both an efficient learning scheme that is robust and avoids the bias found in other optimisation approaches of models that comprise Discrete variables, e.g., continuous relaxations. Experimental evaluations of these new optimisation schemes lead to SVAEs with competitive performance when compared to SOTA time series models while resulting in interpretable learned representations.

Strengths

SVAEs constitute a powerful generative approach with potentially rich latent representations. The authors propose a new optimisation scheme that could bypass their training complexity and improve research explorations.

Weaknesses

This work misses comparison against other related methods and some of the claims should be further clarified. 1) The continuous relaxation of the discrete variables do indeed produced a biased approximation, but have been shown to produce good results in practice, even without annealing schedules (line 117). For example see [3,4,5]. These claims should be validated with appropriate experimental evaluations in the considered setting. 2) Since the SVAE framework resorts to maximisation of a surrogate ELBO, doesn't it introduce a different bias? the same applies to when using the Richardson iteration; this approximation does introduce additional bias in the estimation, especially considering the capped implicit gradient estimator that the authors propose. 3) The related work is a bit lacking. The authors should discuss [2,6] and potentially compare against [6]. At the same time, there exists a newly published work by the same authors [7] that was accepted at ICML 2023, published before this submission. Additional discussions and comparisons should be included. 4) In this context, in my view there should be a comparison with other optimisation schemes for discrete latent variables such as the Gumbel-Softmax [1] or the NES strategy [2]. Even though the framework improves over the baseline and performs many times better than the competition, it is important to quantify the impact of the optimisation compared to more recent approaches than the biased estimation of the original SVAE publication. 5) What is the computational complexity of the framework in the various settings that the authors discuss? [1] Maddison, C. J., Mnih, A., and Teh, Y. W. The concrete distribution: A continuous relaxation of discrete random variables, ICLR 2017 [2] Alon Berliner, Guy Rotman,Yossi Adi et al. Learning Discrete Structured Variational Autoencoder Using Natural Evolution Strategies, ICLR 2022. [3] KP Panousis, S Chatzis, S Theodoridis, Nonparametric Bayesian Deep Networks with Local Competition, ICML 2019 [4] K Kalais, S Chatzis, Stochastic deep networks with linear competing units for model-agnostic meta-learning, ICML 2022 [5] Sebastian Jaszczur, Aakanksha Chowdhery, Afroz Mohiuddin, et al., Sparse is Enough in Scaling Transformers, NIPS 2021 [6] Yixiu Zhao, Scott W. Linderman, Revisiting the Structured Variational Autoencoder, NIPS BDL Workshop 2021 [7] Yixiu Zhao, Scott W. Linderman, Revisiting Structured Variational Autoencoders, ICML 2023

Questions

Please see the "Weaknesses" section.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

2 fair

Presentation

3 good

Contribution

2 fair

Limitations

Please see the "Weaknesses" section.

Reviewer 4krc7/10 · confidence 4/52023-07-07

Summary

The paper introduces a novel, efficient learning algorithm for structured variational autoencoder for any graphical model. It uses a block updating algorithm with belief propagation to approximate the posterior and employs implicit differentiation for end-to-end learning that is more memory-efficient than an unrolled alternative. On top of that it derives unbiased natural gradients to improve the convergence of the learning and a few practical tricks such as parallel inference and robust initialization to make the implementation scalable. Overall, the proposed algorithm makes SLDS models competitive to state-of-the-art time series models.

Strengths

### originality The proposed learning algorithm is novel. Although each component seem to be well-established, combining them together is non-trivial. ### quality The paper is solid and the method is sound. ### clarity The paper is well-written and easy to follow. ### significance The proposed method makes learning and inference of a wide range of graphical models practical through SVAEs, which is very important to the filed of probabilistic ML.

Weaknesses

### quality I think the experiments of the paper can be improved to incorporate more graphical models. For now, the paper only tests the learning method on SLDS even though the method can be applied to a wide range of graphical models. ### clarity The arrows for inference and generative networks in Figure 2 are pretty small.

Questions

If the straight-through estimator is used, is the natural gradient still unbiased as a whole? Are the time reported in Table the overall time of the 3-stage training or the number without the initialization (which seems to be expansive)?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Limitations are not discussed in the paper. Some of my questions above could be discussed in a dedicated limitation section/paragraph.

Reviewer JW8q7/10 · confidence 3/52023-07-10

Summary

The paper builds upon a previous work (Structured VAEs) which tries to integrate variational auto encoders with probabilistic graphical models resulting in more expressing latent variable models for data. The authors propose a new unbiased estimator of the gradients for optimization the ELBO (for SVAE) and show via experiments that the superior performance of their algorithm. In particular, they show their structured model outperforms transformers and RNN based models on sequence prediction tasks (time-series predictions).

Strengths

This paper makes learning SVAEs much more feasible by making the training more efficient due to the unbiased natural gradient estimators proposed by the authors. The authors further propose an interesting implicit differentiation method to bypass the need to back propagate through the block coordinate ascent iterations while optimizing q(z) in the inner optimization loop. This would result in tremendous memory gains. The paper and the appendix is well-written for most part (see weaknesses).

Weaknesses

1. The paper is very clear as to its optimization innovations for improving the training to SVAEs but unclear or ambiguous in other aspects. In particular, one of the claimed contributions of this work is the ability to learn discrete latents (it is in the title of the paper). However, only one small section (S3.2) is devoted to this and the contribution is not clear. The way it is written, the solution seems to be simply removing the discrete variables from input to the decoder/inference network. It is not clear why this is not needed in SVAEs? What approximation is being made? 2. The authors claim that continuous relaxations of discrete distributions result in optimization procedures sensitive to the annealing schedule whereas their method is free from any such limitation since they do not need to backpropagate through discrete latents. I think an empirical comparison with VAEs that employ such relaxations like the Concrete VAE and the more popular VQ-VAE (this variant does not use a smooth relaxation of the discrete distribution but instead uses a straight-through estimator) would further strengthen this point. Without a comparison with contemporary methods it is hard to appreciate the benefit of SVAEs over these methods for optimizing over discrete latents.

Questions

See the points made in the weaknesses. If my concerns are adequately addressed in the rebuttal I would be happy to improve my score. Here are few other minor concerns that could improve the readability of the paper. 1. I am not sure I understand the distinction between the true loss and the surrogate loss. Is the point that while optimizing q(z) (the latent variables) we learn an additional inference network \hat{l}(z | x) which parametrizes p(x | z) as an exponential distribution which is a conjugate to p(z)? Is \hat{l}(z | x) replaced by p_{\gamma}(x | z) while optimizing the outer loop of the ELBO with the locally optimal q(z)? Moreover, how is it ensured that \hat{l}(z | x) will be a good approximation of p_{\gamma}(x | z). 2. In Fig.2 and Sec 3.1 what are the disentangled factors and residual edges? 3. Lines 166-177 the capped implicit gradient estimator is not clear to me. If the Richardson iteration is run every step of the forward optimization. Does it mean that in every iteration of the inner loop when the \omega gets updated (by block coordinate updates), the gradient in eq (8) gets computed and the parameter \eta gets updated?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

N/A

Reviewer JW8q2023-08-16

Further Clarification

Thank you for your detailed responses. I had some follow-up questions. 1. The author's response does not answer my question as to how discrete variables are optimized in SVAEs. Could the authors illustrate their point with the simple example of the GMM L124-126? 2. The following line is unclear to me "we simply need to specify our model such that the discrete latent variables are not direct inputs to the decoder network." - Is this a modelling assumption or any given distribution involving discrete random variables can be written this way? 3. I am still not clear what the surrogate loss is. I understand it is still a lower bound, but what exactly is it? It will help to illustrate the surrogate loss vs true loss (which the VAE optimizes) using the standard normal example. This has been described in L81-84, but is not clear.

Authorsrebuttal2023-08-18

Thank you for your questions! To answer **3**: Given a data point $x^{(i)}$, the ELBO-optimal posterior $q(z|x^{(i)})$ depends on both the prior $p(z|\theta)$ and the likelihood $p_\gamma(x^{(i)}|z)$. Because $x^{(i)}$ is a fixed value pulled from the data set, we need to think of the likelihood as a function over $z$, writing $\ell(z|x^{(i)})=p_\gamma(x^{(i)}|z)$. To find the optimal posterior, we must integrate $p(z|\theta)$ (which is easy) and $\ell(z|x^{(i)})$ (which is impossible due to the neural network decoder). The central innovation of the SVAE is that, instead of approximating the posterior directly as in amortized VI, we approximate the likelihood function $\ell(z|x^{(i)})$ and then "combine" it with the prior. By combine, we mean: we imagine a world where our prior is still $p(z|\theta)$ but our likelihood is now $\hat{\ell}_\phi(z|x^{(i)})$, and optimize the posterior in that setting. This is the surrogate objective: the ELBO in this hypothetical world. Let's consider the standard normal case with latent dimension 1. The prior is a $\mathcal{N}(0,1)$. Because we cannot integrate the decoder likelihood $\ell(z|x^{(i)})$ over $z$, we approximate it with a normal distribution whose mean is $\mu$ and precision (i.e. $\frac{1}{\sigma^2}$) is $\tau$. We write $\hat{\ell}_\phi(z|x^{(i)}) = \mathcal{N}(z;\mu,\tau^{-1})$, where $\mu$ and $\tau$ depend on $x^{(i)}$. In this setting, the surrogate loss is the ELBO given our $\mathcal{N}(0,1)$ prior and our $\mathcal{N}(\mu,\tau^{-1})$ (surrogate) likelihood. If the prior and likelihood are both Gaussian, then the optimal posterior (which minimizes the surrogate loss) is also Gaussian, with mean $\frac{\tau}{\tau+1}\mu$ and precision $\tau+1$. Note that in training we would never have to actually evaluate the surrogate loss: it is set up so we can find its optimizer in closed form. We now proceed to the Gaussian mixture model, where we will assume all model parameters ($\pi,\mu,\Sigma$) are point estimates for simplicity. The *full* generative model for an SVAE with a latent Gaussian mixture (and fixed output variances) is: $$k\sim\text{Cat}(\pi),\quad z \sim\mathcal{N}(\mu_k,\Sigma_k),\quad x\sim\mathcal{N}(\mu_\gamma(z),I)$$ Where $\mu_\gamma(z)$ is the *decoder* network. The ELBO for this model is: $$\mathcal{L}= E_{q(k, z)}\left[ \log \frac{p(k)p(z\mid k) p_\gamma(x \mid z)}{q(k, z)}\right]\leq \log p(x)$$ The amortized inference approach would estimate $q(k,z)$ via $q_\phi(k,z\mid x)$. However, in this case computing the loss requires sampling from $q_\phi(k, z\mid x)$, which is problematic with respect to gradients as discussed elsewhere (enumeration over $k$ is possible in the GMM case, but is infeasible for time-series models). An alternative used by several models is to make the (structured) mean-field assumption and assume that $q(k, z)$ factorizes as $q(k, z)=q(k)q(z)$. $$\mathcal{L}=E_{q(k)q(z)}\left[ \log\frac{p(k)p(z\mid k)p_\gamma(x \mid z)}{q(k)q( z)}\right]\leq\log p(x)$$ The advantage of this factorization is that both $E_{q(k)}\left[\log p(k) \right]$ and $E_{q(k)q(z)}\left[\log p(z\mid k) \right]$ can be computed in closed-form and thus don't require sampling, while the reconstruction term: $E_{q(z)}\left[\log p_\gamma(x\mid z) \right]$ *no longer depends on $k$* and thus can be estimated with the standard reparameterization trick. As previously discussed, if the decoder took $k$ as input we would be left with $E_{q(z)q(k)}\left[\log p_\gamma(x\mid z, k) \right]$ and would be back to needing to sample $k$, thus to answer **2**, this is how we avoid sampling $k$ and it can be seen as a modeling choice. In the case of GMM, this choice does not hurt us: the continuous latent variable contains strictly more information than its cluster assignment so should be sufficient as an input to the decoder. If we again use amortized inference for both of these distributions, so that $q(k)$ becomes $q_\phi(k\mid x)$ and $q(z)$ becomes $q_\phi(z\mid x)$, we get the *stochastic inference network* (SIN) approach. This approach performs poorly in empirical tests due to the strong decoupling of $q(k)$ and $q(x)$ ($q(k)$ is also not directly informed by the decoder). The SVAE approach is to consider a world where the likelihood function isn't parameterized by neural networks, but is instead a simple Gaussian (the likelihood is only a function of $z$, not of $k$, so we don't need to define an approximate likelihood on $k$). GMMs with Gaussian likelihoods are well studied: in this case, optimizing the surrogate loss cannot be done in close form but can be done via *coordinate ascent variational inference* (CAVI). To answer **1**, CAVI alternates updates of $q(z)$ and $q(k)$ to optimize both using expectations, not samples. As a result, $q(k)$ can be optimized without sampling and its contributions to the loss can be evaluated without sampling, but it informs reconstruction: the parameters of $q(z)$ depend on the expectations of $q(k)$.

Reviewer JW8q2023-08-21

Update score

Thank you for the response. It addresses my concerns. Please add these details to the appendix since SVAE (in my opinion) is a niche subject and since the original paper does not answer some of these points well (like. ability to model discrete variables), I urge the authors to add this information in the appendix at thee very least.

Program Chairsdecision2023-09-21

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC