Amortized Inference Regularization
The variational autoencoder (VAE) is a popular model for density estimation and representation learning. Canonically, the variational principle suggests to prefer an expressive inference model so that the variational approximation is accurate. However, it is often overlooked that an overly-expressive inference model can be detrimental to the test set performance of both the amortized posterior approximator and, more importantly, the generative density estimator. In this paper, we leverage the fact that VAEs rely on amortized inference and propose techniques for amortized inference regularization (AIR) that control the smoothness of the inference model. We demonstrate that, by applying AIR, it is possible to improve VAE generalization on both inference and generative performance. Our paper challenges the belief that amortized inference is simply a mechanism for approximating maximum likelihood training and illustrates that regularization of the amortization family provides a new direction for understanding and improving generalization in VAEs.
This paper puts forward the idea that we should in certain cases regularize the generative model in VAEs in order to improve generalization properties. Since VAEs perform maximum likelihood estimation, they can in principle exhibit the same overfitting problems as any other maximum likelihood model. This paper argues that we can regularize the generative model by increasing the smoothness of the inference model. The authors consider the Denoising VAE (DVAE) as a means of achieving such regularization. In the special case where the encoder is an exponential family, they show that the optimum natural parameters for any input data can be expressed as a weighted average over the optimum parameters for the data in the training set. Simlarly for a decoder model in the exponential family, the optimum choice of mean parameters for any latent code is a weighted average over sufficient statistics of training data. Experiments consider evaluate the test set log marginal likelihood for MNIST, Omniglot, and Caltech 101 with the DVAE, a convex combination between DVAE and VAE, and a third smoothing strategy based on regularization that bounds the magnitude of encoder weights.  I think that this is a well-written paper that provides some useful insights about the interplay between encoder and decoder expressivity in VAE. Equations (14), (15) and (19) in particular provide a really nice intuition for the behavior of optimal encoders and decoders. The experiments are reasonable. I was not able to find anything amiss with the proofs (albeit on a cursory read). My only complaint about this paper is that there are some issues of clarity in the discussion of experiments approved (see questions below). Overall, this seems like a welcome addition to ongoing efforts to better understand the generalization properties of VAEs. I would be happy for this paper to appear.   Comments and Questions   - I'm not entirely sure what to make of the results in Table 1. While it seems intuitive that regularizing the inference model would improve the test set log marginal likelihood, it is less intuitive that doing so would also decrease Δ_inf, which one would naively expect to *increase*. Could the authors comment on this.  - I was not quite able to follow the definition of the regularization parameter. This should probably be discussed in the main text, rather than the appendix.   - When the authors write  "Our main experimental contribution in this section is the verification that increasing the number of importance samples results in less underfitting when the inference model is over-regularized."   It sounds a bit like they are saying that increasing the number of IWAE sample decreases susceptibility to underfitting. Is this in fact what they mean to say, or do these results simply validate Proposition 3?   - While perhaps not the point per se of this paper, equation (19) provides some intuitions that are probably worth pointing out in expliclity. If we apply (19) to a standard VAE with Bernoulli likelihood, then the mean value of each pixel will be a weighted average over pixel values of examples in the training data encode to similar z values. My read of this is that, for an optimal decoder, "interpolation" really means "averaging over nearest neighbors in the latent space".   - For a normal VAE, the mutual information I(x;z) in the inference model tends to saturate to log n (where n is the size of the training data). This suggests that the weights q(x^(i) | z) would be pretty close to 0 for all (i) except the nearest neighbor for a normal VAE (since each x^(i) encode to non-overlapping regions in z space). In other words, for a normal VAE, an optimal decoder could effectively be doing nearest-neighbor matching in the latent space. As you increase the regularization, presumably the generative model averages over a larger number of neighbors. It would be interesting to look at how the expected entropy of q(x^(i) | z) changes with the amount of regularization (the exponent of this entropy would effectively give you something like "the number of neighors" that an optimal likelihood model would average over).   - The posteriors in Figure 2 of cited ref [10] look pretty significantly non-Gaussian to me. I unfortunately lacked time to read ref [10] in detail, but I was not immediately able to find support for the claim that "it has still been observed empirically that the true posteriors tend to be approximately Gaussian by the end of training as a result of posterior regularization induced by the variational family". Could the authors either clarify, or walk back this claim?   - The notation F(Q) does not make sense to me. In this notation would the domain of F somehow be the set of all possible variational families? Given that Q is a specific family of variational distributions (in practice a parametric family like a Gaussian with diagonal covariance), and F(Q) is a family of functions that map from X to Q, it is not clear how you would even apply the same F to two different families Q and Q'. Writing F_Q, or perhaps just F would be clearer, since F is in practice specific to a family Q. The same argument applies to G(P).  - How are the errors in Table 1 defined? Is this an error over independent restarts? If so, how many restarts? Could the authors figure out a way to reformat the table so the erros are not shown in tiny fonts? I understand that we're all using screens to read papers these days, but I still think that fonts should be legible in print.   - Does Table 1 show results for the regularization parameters listed in Table 7?   - Do Figures 1 a 2 show the same IW-SVI approximation of log p_θ(x) as Table 1?  - Why does the Caltech panel of Figure 1 have error bars, whereas the Omniglot and MNIST panels do not?  - q(x^(i) | z) is never explicitly defined, as far as I can tell.   - The acronym "AIR" results in an unfortunate collision with "Attend, Infer, Repeat", which at this point is pretty well known.