Title: On the Value of Infinite Gradients in Variational Autoencoder Models
Authors: Bin Dai
Li Kevin Wenliang
David Wipf
Venue: NeurIPS 2021
Keywords: variational autoencoders
sparse representations
latent variable models

Abstract: A number of recent studies of continuous variational autoencoder (VAE) models have noted, either directly or indirectly, the tendency of various parameter gradients to drift towards infinity during training.  Because such gradients could potentially contribute to numerical instabilities, and are often framed as a problematic phenomena to be avoided, it may be tempting to shift to alternative energy functions that guarantee bounded gradients.  But it remains an open question: What might the unintended consequences of such a restriction be?  To address this issue, we examine how unbounded gradients relate to the regularization of a broad class of autoencoder-based architectures, including VAE models, as applied to data lying on or near a low-dimensional manifold (e.g., natural images).  Our main finding is that, if the ultimate goal is to simultaneously avoid over-regularization (high reconstruction errors, sometimes referred to as posterior collapse) and under-regularization (excessive latent dimensions are not pruned from the model), then an autoencoder-based energy function with infinite gradients around optimal representations is provably required per a certain technical sense which we carefully detail.  Given that both over- and under-regularization can directly lead to poor generated sample quality or suboptimal feature selection, this result suggests that heuristic modifications to or constraints on the VAE energy function may at times be ill-advised, and large gradients should be accommodated to the extent possible.

TL;DR: We demonstrate that infinite gradients, although perhaps at times difficult to address in practical, can serve a useful role in pruning the latent space of autoencoder-based models.

Review discussion:
I really like the prose and narrative of this paper, which nicely (and somehow pedagogically) explain an important phenomenon that has been studied in the literature extensively, albeit with a different set of tools.
I must confess I had troubles starting from sec.2 of this work, as my understanding of an optimal representation related to prior works that are more closely related to information theory: "Information Dropout: Learning Optimal Representations Through Noisy Computation", by Achielle, Soatto https://arxiv.org/abs/1611.01353v3, and "The information bottleneck method" by Tishby et al., https://arxiv.org/abs/physics/0004057.
So, condition ii) in Def.1 sounds really strange to me: "the number of latent dimensions of $z$ containing no information about $X$ should be maximal". This is somehow clarified in the comments to the definitions, but honestly I find this phrasing very difficult to digest. My interpretation is: if $k > r$, then we want all of the $k-r$ components to behave as the prior.
Another point of confusion is the way that eq(3) and eq(5) are discussed, below eq(5).
When you plug back eq(5) in eq(3), you do not only consider the $\log \gamma$ term, but there also is a $1/\gamma$ term to take into account.
That being said, here is a question.
* Q1: what about situations in which the likelihood is not Gaussian? I assume similar considerations to those made in this work would somehow hold, but it would be nice to discuss this more general case.
Section 3 reads well, and did not pose any major challenges to be appreciated, especially thanks to sec 3.3 which is helpful (I did not check the full proof in the appendix).
Section 4 is again straightforward to appreciate, and the implications are clear. What would be interesting to understand is in the following question.
* Q2: what is the convergence rate of learning $\gamma$?
That is, how fast does it reaches the optimal solution $\gamma \to 0$ compared to the variational parameters of the encoder, and the parameters of the decoder? Is there any potential instability in the process? Would it be possible (as illustrated in fig.3) that when $\gamma$ reaches small values, the VAE loss gets stuck into a bad local minimum?
In conclusion, I find this a good paper, which nicely extends existing considerations from the literature. It has some practical implications, but an extensive experimental campaign is already contained in previous work from ICLR 2019.
No
Thanks for the positive comments and recognizing the nice theoretical work with practical implications.  We address each question below with a summary of original reviewer statements included where needed.
**Question:**  Further background context regarding the definition of optimal sparse representations
**Response:**  In Definition 1 we formally define the particular notion of optimal sparse representation that is applied throughout our work.  That being said other reasonable definitions can be proposed as well depending on the application scenario being considered.
However, we chose our stated version in part for analysis purposes, and in part because for many types of data lying on or at least very near low-dimensional manifolds (which is essentially all natural images), our definition of an optimal sparse representation directly aligns with the dimensionality of such manifolds.  And Definition 1 also corresponds with a long precedent in the signal and image processing literature involving the search for low-dimensional structure in high-dimensional data.  We can include additional supporting references to a revision, e.g., seminal papers such as "Optimally sparse representation in general (nonorthogonal) dictionaries via $\ell_1$ minimization," "Robust uncertainty principles: Exact signal reconstruction from highly incomplete frequency information," and "Exact matrix completion via convex optimization" among many others.
**Question:**  Another point of confusion is the way that eq(3) and eq(5) are discussed, below eq(5).
When you plug back eq(5) in eq(3), you do not only consider the $\log \gamma$ term, but there also is a $1/\gamma$  term to take into account.
**Response:**  Actually, once we plug in the optimal $\gamma$ value from Eq. (5), the $1/\gamma$ term reduces to an irrelevant constant since the numerator and denominator factors cancel out.
**Question:**  Q1: What about situations in which the likelihood is not Gaussian? I assume similar considerations to those made in this work would somehow hold, but it would be nice to discuss this more general case.
**Response:**  Our work is devoted to VAEs applied to continuous data, for which Gaussian likelihood models are by far the most common assumption.  We are not aware of other continuous data likelihoods being regularly applied to VAEs in practice, but even so, the core conclusions we arrive at would likely still hold.  This is largely because, for data lying on or near a low-dimensional manifold, any possible likelihood model that correctly restricts its probability mass to the manifold will necessarily involve an unbounded density and thus enter the scope of our analysis.
In contrast, for discrete data likelihoods, the infinite gradient issue we explore is not likely to be a factor.
**Question:**  Q2: what is the convergence rate of learning $\gamma$?
That is, how fast does it reaches the optimal solution $\gamma \rightarrow 0$ compared to the variational parameters of the encoder, and the parameters of the decoder? Is there any potential instability in the process? Would it be possible (as illustrated in fig.3) that when  reaches small values, the VAE loss gets stuck into a bad local minimum?
**Response:**  Because of Eq. (5), the value of $\gamma$ will not become small until the other model parameters (in both the encoder and decoder) are capable of producing reasonably small reconstruction errors;
hence they will generally all converge together.
Furthermore, depending on how exactly *bad* is defined, a bad local minimum is indeed possible if some collection of parameters produces a small reconstruction error using a larger-than-necessary number of active latent dimensions.
For example, per the illustration from Figure 1, a 2-dimensional linear subspace could produce small reconstruction errors even though only a 1-dimensional nonlinear alternative is actually adequate. And depending on the model parameterization, it could require increasing the model energy to traverse from the former to the latter.
Thanks for the clarifications, which address all my questions.
### Originality:
- Section 4 of this work is most similar to Dai and Wipf (2019). Dai and Wipf (2019) experimentally showed the benefit of training $\gamma$ (beyond fixing it at a low value).
However, Dai and Wipf (2019) focus more heavily on ensuring a valid ancestral sampling process via a two-stage VAE. This article attempts to explain *why* optimizing  $\gamma$ is better than fixing it at an arbitrarily low value. Figure 3 is incredibly informative at showing this benefit.
Thus, I feel the author(s) are well-enough distinguished from Dai and Wipf (2019).
Although, I might consider moving their citation forward in this section from the last sentence since to the best of my knowledge Dai and Wipf (2019) were the first to recognize the unappreciated benefit of learning $\gamma$, which the author(s) "$\gamma$-dependent smoothing"
### Quality:
- Overall, I think the theoretical findings are well presented.
### Clarity:
- This paper is exceptionally well written.
- Line 47: "The latter requires that we assume specific functional forms for the encoder and decoder distributions." This is only true for the encoder. The decoder distribution can be any distribution (selecting a distribution with reparameterization gradients is unnecessary since you don't differentiate any MC-integrated quantity w.r.t. decoder distribution)
- I think the author(s) need explain $\theta \in \Theta$ better in the main article.
In particular, $\Theta$ is used before it is defined. Can the author(s) please elaborate what is $\Theta$ (or provide an example) somewhere before the supplement?
- Equation 13 drops $\frac{1}{n}$ between lines. This scaling reappears in equation 14.
### Significance:
-  Theorem 4 is significant to those working on disentanglement based on flavors of the $\beta$-VAE.
- This article explains why Dai and Wipf's (2019) proposal to optimize  $\lambda$ is better than fixing it at an arbitrarily low value.
While this explanation (section 4) lacks  the the theoretical scaffolding of the earlier sections, it benefits from well illustrated examples.
No
Thanks for the positive feedback regarding our work, including complementary comments regarding the writing and significance.
In terms of clarity, several point were raised that we address as follows.
**Question:** Line 47: "The latter requires that we assume specific functional forms for the encoder and decoder distributions." This is only true for the encoder. The decoder distribution can be any distribution (selecting a distribution with reparameterization gradients is unnecessary since you don't differentiate any MC-integrated quantity w.r.t. decoder distribution).
**Response:**  Thanks, we can clarify this point.
**Question:** I think the author(s) need explain $\theta \in \Theta$ better in the main article.
In particular,  is used before it is defined. Can the author(s) please elaborate what is  (or provide an example) somewhere before the supplement?
**Response:** The constraint $\theta \in \Theta$ is defined below Eq. (6), which should be the first instance where it is used.
However, we agree that some further clarification could be helpful here.  At a high level, this constraint can be viewed as merely defining the parameter space of a given AE/VAE architecture.
However, it also serves as a mechanism for preventing trivial solutions that can occur within any regularized AE architecture because of scaling ambiguities.
For example, without any constraint on the AE decoder parameters (e.g., bounded magnitudes or related), the latent codes $Z$ can be pushed to arbitrarily small values to minimize any standard shrinkage penalty factor (e.g., L1 norm), while the corresponding decoder weights merely becomes arbitrarily large.  In this way there is effectively no regularization of the AE latent space at all.
**Question:**  Equation 13 drops $\frac{1}{n}$ between lines. This scaling reappears in equation 14.
**Response:** Good catch, we will fix this.
Thank you for your reply. I keep my score and remain in favor of accepting.
Overall, the paper at hand is clear, novel, and provides interesting theoretical and practical insights. The presentation of the paper is good, and notably the introduction is a very natural and well polished exposition of Variational Autoencoders.
I have two mild concerns about the paper:
- I think the way Theorem 4 is stated could be made more formal, which would ease the comprehension and avoid confusions. Currently, given the proof, I think the formal statement that holds is the following:
For any (concave, non decreasing) h with bounded gradients near 0, there exists a dataset X and a class of decoders \mu_x(z, \theta) such that, for an arbitrarily expressible encoder, the global optimum of the loss L_h(\theta, \phi) does not achieve an optimal sparse representation.
With the formulation that is given in the text, I think it could be misinterpreted that the statement holds for all class of decoders, e.g. that the statement would be:
For any (concave, non decreasing) h with bounded gradients near 0, and for all class of decoders \mu_x(z, \theta) there exists a dataset X such that, for an arbitrarily expressible encoder, the global optimum of the loss L_h(\theta, \phi) does not achieve an optimal sparse representation.
- While it is clear, given prior works, that obtaining an optimal sparse representations requires that \gamma should go to 0, it would be useful to make the relationship between \gamma going to zero and gradients becoming unbounded clearer.
Notably, it might be useful to strengthen the fact that when \gamma goes to 0, not only the gradient of the loss w.r.t. to \gamma goes to \infty, but also potentially the gradient of the loss w.r.t. the parameters of the decoder, as the numerator of this gradient will typically scale like \sqrt{\gamma}, while the denominator will scale like \gamma.
While it may not be possible to provide a formal statement on that, I think this would still strengthen the claim on unbounded gradients, as having unbounded gradients only on \gamma could be viewed as an easily fixable problem (e.g. by specifically normalizing this gradient), but having potentially unbounded gradients on the decoder parameters could prove much harder to deal with.
I think those two concerns are easily fixable, and given the interesting insights provided by the paper, I recommend acceptance of the paper. I will increase my ratings if my concerns are alleviated after rebuttal from the authors.
# Post rebuttal edit:
The authors have addressed my two concerns. Given their response, I am raising my score.
No
We appreciate the constructive feedback.  In this regard, the reviewer raised only two mild concerns, suggesting that both are easily fixable.  We agree with this conclusion, and address both issues briefly in turn below.
**Rewording of Theorem 4**
The reviewer's suggested restatement of Theorem 4 is indeed a correct interpretation of our meaning, i.e, for any concave, non-decreasing $h$ with bounded gradients near $0$, there will always exist a dataset $X$ and a class of decoders $\mu_x(z; \theta)$ such that, for an arbitrarily expressible encoder, the global optimum of the loss $L_h(\theta, \phi)$ does not achieve an optimal sparse representation.
While we believe our original phrasing is also technically correct, per the reviewer's comments, we can see how misinterpretation is easily possible.  Fortunately, it is easy to clarify this issue in a revision, thanks for the nice suggestion.
**Further Clarification Regarding Unbounded Gradients**
The reviewer points out that not only does the gradient with respect to $\gamma$ become unbounded, but also the gradient of other decoder parameters.  This is definitely true and can be inferred from the discussion below Eq. (5) whereby $\gamma$ is optimized out of the model, but we agree that it could be beneficial to add some further elaboration of this point.
Indeed, there are many perspectives from which one can see that this will inevitably happen, including the $\gamma$-dependent scaling argument the reviewer alluded to.  More specifically, for any fixed value of the other decoder parameters, the optimal $\gamma$ satisfies Eq. (5).
Hence the squared reconstruction error averaged over the encoder posterior scales proportionally to $\gamma$.
But the gradient of this quadratic term should then scale according to the square-root of $\gamma$, i.e., the raw un-squared error.
Consequently, once we introduce the $1/\gamma$ factor from Eq. (3), we find that the gradient with respect to all decoder parameters should loosely scale as $\sqrt{\gamma}/\gamma = 1/\sqrt{\gamma}$, which will of course become unbounded as $\gamma \rightarrow 0$.  Actually, for simplified VAE models (e.g., affine decoders) it is straightforward to show that this is exactly the case.  Great suggestion, we can add this type of discussion to a revision.

Meta-review:
role: AC
data: comment: Hi

I was hoping the authors could clear up a question for me about the paper: why is Eq (6) considered a generic auto-encoder architecture?  To me, it seems that this is actually relatively particular in that the same h is being applied to both the reconstruction and the regularization (i.e. why does h need to appear in the same manner twice).  Why could one not simply have, say, one of these be linear and the other log?  As such, the claim of generality here seems unfounded to me at present and the results simply apply to a commonly occurring AE formulation, rather than all possible formulations.  If this is the case, this could potentially further imply that some of the conclusions are not fully demonstrated (particularly those relating to "impossibility").

Thanks in advance for your help in clearing this up.
role: Author
data: value: comment: Good question, thanks.  Actually, we originally chose the same function $h$ for both reconstruction and regularization in Eq. (6) because it provides the most natural, direct link with VAE models, e.g., see Lemma 3. (It also represents a form or simplified, balanced regularization akin to the square-root Lasso, and for appropriate choices, a shared $h$ can indeed produce optimal sparse representations.)  Importantly though, this is not at all a requirement for our analysis.  In fact, it is quite straightforward to modify Theorem 4 to accommodate any arbitrary function $g \neq h$ to replace $h$ in the reconstruction term.  To see this, note that if $g$ is chosen such that the global optimum does not enforce zero reconstruction error, then we have not actually achieved an optimal sparse representation per our Definition 1.  And if $g$ is chosen such that zero reconstruction error does occur at the global optimum, then the steps from the proof of Theorem 4 can proceed exactly as before, with $g(0)$ simply replacing $h(0)$ at the appropriate instances.  

In particular, the right-hand side of Eq. (23) in the supplementary becomes $g(0) + \frac{1}{2}h(0) + \frac{1}{2}h(C^2)$ while the right-hand side of Eq. (24) becomes $g(0) + 2h(\epsilon^2)$.  As these quantities are subsequently compared to evaluate which is larger, it makes no difference what the value of $g(0)$ is, and so the same conclusions follow as before.  We had already planned to revise Theorem 4 to account for this generalization, and the AC's question confirms that this change is indeed advisable.
role: AC
data: value: comment: Thanks for clarifying; yes I think this generalization is quite important and should be incorporated for the final version.
role: PC
data: comment: This is a strong paper that all reviewers felt should be accepted to the conference; a viewpoint I concur with myself.  Some constructive criticisms have though been raised and I encourage the authors to incorporate these suggestions for the final version of the work.

A few extra comments from my end to help update the paper for the camera-ready version:
- I'm not sure "sparse" is the right term in Definition 1 (i.e. in the term "optimal sparse representation") and could be a bit misleading.  Sparse representations usually refer to cases where *different* sets of latent dimensions will be "activated" for different inputs, rather than the same latents always being redundant.  I would suggest using some other term like "minimal".
- It should be made clear in Eq 6 that this is now an optimization w.r.t Z.
- Even with the discussed generalization of Theorem 4, I think some care is needed in discussing the generality of the results and impossibility of AE frameworks that can satisfy Definition 1 without infinite gradients as even with this generalization, it is still a particular form of loss function rather than any possible AE framework.  As such, some of the higher level assertions made need slightly more careful quantification to avoid overclaiming / misunderstanding.
- I feel that the introduction could be tightened up a little: though nicely written at a low level, it somewhat meanders at the moment and one has to go through quite a lot of basic material before getting to the point of the paper.  In particular, it is only really in the second half of page 2 that the content becomes specific to the paper; before that it is a very generic overview of the area, which is wasteful.  In other words, though (as a reviewer noted) this is a nice introduction to VAEs, I don't think such a comprehensive introduction to them at a general level is needed and this both slightly slows down the early part of the paper and uses up space that could maybe be better utilized.

Decision: Accept (Spotlight)