A Sampling-Free Approximation of Gaussian Variational Auto-Encoders

1 INTRODUCTION . Variational Auto-Encoders ( VAE ) allow generating data , mapping data into a latent space , and modifying data by perturbing it in a controlled manner in the latent space [ 1 , 2 ] . This enables them to generate and modify images [ 3 , 4 ] , speech [ 5 ] , and molecular structures [ 6 ] , or detect outliers in distributions [ 7 ] . VAEs extend the concept of regular auto-encoders by modeling the latent space as a distribution and with a distributional constraint that allows for a distributional latent space . Thus , they enforce all samples drawn from the prior to be representative for the data set . For this , VAEs predict a distribution through the encoder , sample from this distribution , and propagate the sample through the decoder . During training , they constrain the predicted distribution with respect to the prior , and require that the reconstructions produced by the decoder are correct . To allow for random sampling in the context of gradient descent , Gaussian VAEs employ the sampling reparameterization trick [ 8 ] , which enables backpropagation through the sampling operation . In fact , the sampling reparameterization trick and VAEs have received a lot of attention in research [ 3 , 9–11 ] . Recently , research has also started to look into techniques to compute the posterior and marginal distribution in an exact analytical solution without the need for sampling . In a sampling-free formulation , during training , the loss is not computed via stochastic sampling but derived by propagating the distributions from the latent space to the output space . However , computing the exact likelihood is in many cases intractable . Two approaches have recently been proposed to address this challenging problem . Balestriero et al . [ 12 ] propose to compute the likelihood for small neural networks by integration over piecewise continuous polytopes . Lucas et al . [ 13 ] formulate an analytical sampling-free linear VAE , which can be efficiently computed in closed form . In this context , the following work proposes a sampling-free approximate formulation for VAEs that allows sampling-free training of VAEs . To this end , we approach the problem of intractability by approximating the probability distributions instead of computing them exactly . This approximation is done by local Taylor series expansion of locally affine ( ReLU-activated ) neural networks . This corresponds to the local linearization of the respective neural network and allows for a local linearization of the decoder of the Gaussian VAE . We propose to use the local linearization to estimate the ( co- ) variances for each data point , while the median of the output distribution is computed exactly . Locally approximating neural networks by their Taylor series around each individual data point enables us to express the covariance matrix of the probability distribution in terms of the Jacobian of the neural network . Furthermore , as computing the covariance matrix explicitly is very expensive , the proposed formulation allows computing the likelihood without explicitly computing the covariance matrix nor its inverse . This allows us to compute the training objective without sampling at a reason- able additional cost for standard architectures , specifically , at the cost of evaluating the Jacobian of the decoder . Note that the local Taylor series expansion or local linearization of the neural network decoder function does neither correspond to a global linearization of the neural network decoder nor a global approximation via its Maclaurin series . We evaluate the proposed method on the CelebA [ 14 ] , the 3D Chairs [ 15 ] , and the MNIST [ 16 ] data sets . When considering the number of samples necessary to estimate a covariance matrix of equivalent quality wrt . the Frobenius norm via stochastic sampling , our approximation is ( on average ) equivalent to more than 235 samples for a deep network on the CelebA data set . Overall , the proposed approach achieves competitive results compared to sampling Gaussian VAEs . Especially on the Frechet Inception Distance ( FID ) [ 17 ] , the sampling-free VAE performs better than the sampling VAE , while the estimated marginal likelihoods between both methods are on par . This demonstrates that local linearization is a good approximation and does not damage performance in comparison to a regular sampling VAE . Following the setup by Lucas et al . [ 13 ] , the sampling-free VAE leads to smaller ( ǫ , δ ) −posterior collapse than its sampling counterpart . We summarize the contributions of this work as follows : • We present a novel sampling-free approximation for training Gaussian VAEs . In contrast to exact computation , it requires only a small overhead compared to its sampling counterpart , and reaches a less noisy approximation of the training objective than sampling . • We show that the sampling-free Gaussian VAE approximation achieves competitive performance on the MNIST , Omniglot , CelebA , and 3D Chairs data sets compared to sampling VAE methods . Also , the sampling-free formulation also achieves a better FID on CelebA when provided with the same computational budget as its sampling counterpart . • We find that sampling-free VAEs are more robust against posterior collapse than sampling VAEs . The full implementation as well as experiments will be publicly available upon publication . 2 RELATED WORK . Balestriero et al . [ 12 ] present an approach to compute the exact likelihood for small yet deep neural networks using the continuous piecewise affinity of respective neural networks to compute the exact likelihood by partitioning the network and integrating over the polytopes . Based on this , an analytical Expectation-Maximization algorithm is derived that enables gradient-free learning . The method is empirically validated by training a three layer deep generative network with 8 and 16 hidden units and with one latent dimension on the number 4 from the MNIST data set , where they demonstrate that an exact sampling-free VAE is superior to the stochastic sampling VAE . However , as the exact analytical formulation of deep neural networks has a computational complexity that is exponential in the size of the neural network , it is intractable for larger networks . Lucas et al . [ 13 ] examine sampling-free linear VAEs , which by the nature of their linearity allow for efficient exact computation of the likelihood and are able to show that the linear VAE recovers probabilistic principal component analysis ( pPCA ) [ 18 ] . They employ the sampling-free linear VAE for understanding posterior collapse in VAEs and apply it to the MNIST data set . Ghosh et al . [ 19 ] propose regularized deterministic auto-encoders , to learn a smooth latent space without sampling via different regularization schemes . They regularize their auto-encoders through weight decay , the gradient norm of the decoder and/or spectral normalization , which all bound the Lipschitz continuity of the decoder . Our work ( implicitly ) also includes a kind of gradient regularization of the decoder . However , in our work , this arises from approximating the marginal log-likelihood via local Taylor approximations of the decoder . Tolstikhin et al . [ 20 ] propose Wasserstein Auto-Encoders ( WAEs ) , which minimize the Wasserstein distance between the model distribution and the target distribution . Kolouri et al . [ 21 ] propose Sliced Wasserstein Auto-Encoders ( SWAEs ) , which regularize the auto-encoders using the sliced / marginal Wasserstein distance between the distribution of the encoded training samples and a predefined sampleable distribution . Note that both WAEs and SWAEs require sampling [ 20–22 ] . Knop et al . [ 22 ] propose Cramer-Wold Auto-Encoders ( CWAEs ) , a sampling-free distributionally regularized Auto-Encoder . They introduce the Cramer-Wold distance , which has a simple analytical formula for computing normality of high-dimensional samples . By regularizing their auto-encoders via the Cramer-Wold distance , they can evaluate their training objective in closed form . Note that CWAEs are not VAEs as , for VAEs , the distance to the true distribution is measured using KL divergence under the latent variable model [ 22 ] . Instead , the CWAE is a different kind of generative model that was designed to have a tractable sampling-free objective . Burda et al . [ 23 ] introduce importance weighted sampling for VAEs by using multiple samples to have a tighter variational lower bound . Huang et al . [ 24 ] propose hierarchical importance weighted auto-encoders reducing the number of redundant samples and introducing a hierarchical structure to induce correlation . Dieng et al . [ 25 ] use importance weighted VAEs using an EM-based algorithm employing moment matching . Roeder et al . [ 26 ] propose a low-variance gradient estimator for VAEs . Park et al . [ 27 ] enrich the posterior by applying the Laplace approximation to VAEs , which enables modeling a more expressive full-covariance Gaussian posterior . Tomczak et al . [ 28 ] also produce an improved posterior distribution through a Householder transformation-based volume-preserving flow . Nielsen et al . [ 29 ] bridge the gap between normalizing flows [ 30 ] and VAEs by preseting SurVAE Flows , which use surjective transformation which are deterministic in one direction and stochastic / sampling in the reverse direction . Morrow et al . [ 31 ] present VAEs with normalizing flow decoders . Taylor series approximations [ 32 ] are a ubiquitous tool in variational inference ( VI ) , e.g. , in the delta or the Laplace method [ 33 , 34 ] . VI methods , such as the delta method or the Laplace method require solving an optimization problem to obtain an approximation of an output distribution . In contrast , our method uses the Taylor approximation to produce a closed-form estimate of the training objective and the only optimization is the training of the VAE . A phenomenon frequently appearing in VAEs is posterior collapse , i.e. , the posterior distribution ( produced by an encoder ) is close to the prior distribution . This leads to a reduced expressivity of the VAE ’ s latent space . To tackle this , Kingma et al . [ 8 ] proposed to constrain the KL divergence such that it is only active if the KL divergence is above a threshold . Other works used KL annealing [ 35 , 36 ] or constrained the posterior to have a minimum distance to the prior [ 37 ] . Lucas et al . [ 13 ] defined the ( ǫ , δ ) −posterior collapse in order to reliably quantify posterior collapse in VAEs . 3 METHOD . We begin this section by introducing the conventional data generating process of VAEs . After that , we briefly motivate and derive the evidence lower bound ( ELBO ) and present our sampling-free local approximation of the ELBO including an efficient way of computing it . Following this , we discuss a method for stabilizing training and investigate the balance between reconstruction quality and variance of reconstructions . 3.1 GAUSSIAN VARIATIONAL AUTO-ENCODERS . We begin by stating the assumptions about the data generating process . Note that this process does not differ from the conventional data generating process of VAEs [ 1 ] . For easier reference , we follow the notation of Kingma and Welling [ 9 ] . The goal of VAEs is to approximate the true distribution of the data p∗ ( x ) with a distribution parameterized via a model pθ ( x ) ≈ p∗ ( x ) . Here , p ( z ) = N ( z ; 0 , I ) ( 1 ) pθ ( x|z ) = N ( x ; fθ ( z ) , cI ) for some c > 0 ( 2 ) pθ ( x ) = ∫ pθ ( z ) pθ ( x|z ) dz . ( 3 ) z denotes an m-dimensional vector referring to the hidden / latent space , x denotes an n-dimensional vector referring to the input / output space , fθ is the decoder / generator network , and c captures the observation uncertainty . N denotes a multivariate Gaussian distribution and Im the m×m identity matrix . To approximate the intractable pθ ( z|x ) , VAEs include an encoder network gφ and a network that computes the encoder ’ s uncertainty hφ ( in the form of a diagonal covariance matrix ) such that : qφ ( z|x ) ≈ pθ ( z|x ) where qφ ( z|x ) = N ( z ; gφ ( x ) , hφ ( x ) ) . ( 4 ) The training objective is to maximize the log-likelihood of our data log pθ ( x ) by optimizing θ and φ : log pθ ( x ) = Eqφ ( z|x ) ( log pθ ( x|z ) p ( z ) qφ ( z|x ) ) + Eqφ ( z|x ) ( log qφ ( z|x ) pθ ( z|x ) ) ( cf . [ 9 ] ) ( 5 ) = Eqφ ( z|x ) ( log pθ ( x|z ) ) − Eqφ ( z|x ) ( log qφ ( z|x ) p ( z ) ) + Eqφ ( z|x ) ( log qφ ( z|x ) pθ ( z|x ) ) ( 6 ) = Eqφ ( z|x ) ( log pθ ( x|z ) ) −DKL ( qφ ( z|x ) , p ( z ) ) ︸ ︷︷ ︸ ELBO +DKL ( qφ ( z|x ) , pθ ( z|x ) ) ︸ ︷︷ ︸ not computable but non-negative . ( 7 ) The last term is not computable , because pθ ( z|x ) is intractable , and hence the log-likelihood log pθ ( x ) can not be computed exactly . However , since a KL-divergence is non-negative , the first two terms yield a lower bound for log pθ ( x ) , the so-called evidence lower bound ( ELBO ) [ 1 , 9 , 38 ] . The ELBO is the substitute objective that allows training VAEs . ELBO = log pθ ( x ) −DKL ( qφ ( z|x ) , pθ ( z|x ) ) ≤ log pθ ( x ) ( 8 ) The KL-Divergence DKL ( qφ ( z|x ) , p ( z ) ) can be computed in closed form . However , the computation of Eqφ ( z|x ) ( log pθ ( x|z ) ) is generally intractable and is therefore usually approximated via stochastic sampling from the latent distribution qφ ( z|x ) . In the stochastic sampling approach , a data point x is propagated through gφ and hφ , an element z is sampled from N ( z ; gφ ( x ) , hφ ( x ) ) , and propagated through fθ , which approximates the output distribution ( including the observation uncertainty ) N ( x ; fθ ( z ) , cI ) where z ∼ N ( z ; gφ ( x ) , hφ ( x ) ) . Since the observation uncertainty is an isotropic normal distribution , the sample log-likelihood log pθ ( x|z ) reduces to log pθ ( x|z ) = − 1 2c ( ‖fθ ( z ) − x‖2 ) − n 2 log ( 2πc ) ( 9 ) which is , apart from the normalization constant , a mean squared error [ 1 ] .

Summary: The paper proposes an approximation that frees the calculation of ELBOs of Gaussian variational autoencoders from sampling. To achieve this, it utilizes Taylor expansion on the decoder networks. The proposed method was evaluated on three typical datasets. The authors also found that the approximation alleviated the issue of posterior collapse to the sampling VAEs.