A S AMPLING -F REE A PPROXIMATION OF G AUSSIAN V ARIATIONAL A UTO -E NCODERS
Anonymous authors
Paper under double-blind review
A BSTRACT
We propose a sampling-free approximate formulation of Gaussian variational autoencoders.
Instead of computing the loss via stochastic sampling we propagate the
Gaussian distributions from the latent space into the output space.
As computing the exact probability is intractable we propose to locally approximate the decoder network by its Taylor series.
We demonstrate that this approximation allows us to approximate the Gaussian variational auto-encoder training objective in closed form.
We evaluate the proposed method on the MNIST, the Omniglot, the CelebA and the 3D Chairs data sets.
We ﬁnd that our sampling-free approximation performs better than its sampling counterpart on the Frechet inception distance and on par on the estimated marginal likelihood.
I NTRODUCTION
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 efﬁciently 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 afﬁne (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, speciﬁcally, 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 ﬁnd 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.
R ELATED W ORK
Balestriero et al.
[12] present an approach to compute the exact likelihood for small yet deep neural networks using the continuous piecewise afﬁnity 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 efﬁcient 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 predeﬁned 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 ﬂow.
Nielsen et al.
[29] bridge the gap between normalizing ﬂows [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 ﬂow 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] deﬁned the ( ǫ, δ ) − posterior collapse in order to reliably quantify posterior collapse in VAEs.
M ETHOD
We begin this section by introducing the conventional data generating process of VAEs.
After that, we brieﬂy motivate and derive the evidence lower bound (ELBO) and present our sampling-free local approximation of the ELBO including an efﬁcient 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. G AUSSIAN V ARIATIONAL A UTO -E NCODERS
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,
[Equation]
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 I m 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:
[Equation]
where q φ ( z | x ) =
N ( z ; g φ ( x ) , h φ ( x )) .
[Equation]
The training objective is to maximize the log-likelihood of our data log p θ ( x ) by optimizing θ and φ :
[Equation]
The last term is not computable, because p θ ( z | x ) is intractable, and hence the log-likelihood log p θ ( x ) cannot be computed exactly.
However, since a KL-divergence is non-negative, the ﬁrst 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.
[Equation]
The KL-Divergence D KL ( q φ ( z | x ) , p ( z )) can be computed in closed form.
However, the computation of E
[Equation]
θ ( 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 ) , c I ) 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
[Equation]
which is, apart from the normalization constant, a mean squared error [1].
3. S AMPLING -F REE A PPROXIMATION
At this point, we introduce the proposed sampling-free approach.
In the sampling-free approach, a data point x is propagated through g φ and h φ , and then the obtained distribution N ( z ; g φ ( x ) , h φ ( x )) is (without sampling) approximately propagated through f θ , which yields the output distribution
N ( x ; f θ ( z ) , c I ) where z ∼ N ( z ; g φ ( x ) , h φ ( x )) .
The problem is that computing the exact transformation of the distribution N ( z ; g φ ( x ) , h φ ( x )) by f θ is intractable (for non-trivial f θ ) because it can have an exponential complexity.
The largest analytical sampling-free approach so far was able to operate on 24 hidden neurons [12].
This is because, for the exact computation, every linear region of f θ has to be computed separately.
To simplify notation, in the following, we drop subscripts θ and φ for f, g, h .
To make this problem tractable, we approximate the neural network decoder via a local Taylor series expansion, which simpliﬁes computing the output distribution.
For that, we deﬁne the Taylor series of the decoder f around g ( x ) for each data point x as ˆ f g ( x ) .
We display the Taylor series of an example neural network function in Figure A.1 in the supplementary material.
Notably, the function and its Taylor series are equal near this point g ( x ) , which is also the point of largest density of the distribution
N ( z ; g ( x ) , h ( x )) .
Due to the continuous piece-wise afﬁnity of f , we only need to compute the ﬁrst two terms of the Taylor expansion, as the remaining terms are zero.
Thus,
[Equation]
where J f ( g ( x )) is the Jacobian of f at point g ( x ) , in the following abbreviated as J .
This approximation is exact if ǫ is small enough, such that it stays inside the linear region of f in which g ( x ) is located.
In Section 4.1, we empirically demonstrate that the Taylor series delivers a good approximation for the neural network decoders in VAEs.
For neural networks that are not piece-wise afﬁne, like the tanh networks used in the experiments in Section 4.3, we also use Equation 10, which in this case is the ﬁrst-order Taylor expansion.
Using the Taylor series for each point ˆ f g ( x ) we can approximate the output distribution as follows:
[Equation] [Equation]
As the family of Gaussian distributions is closed under afﬁne transformations and summation [39]:
[Equation]
where z ∼ N ( z ; g ( x ) , h ( x ))
[Equation]
12) and including the observation uncertainty
[Equation]
where o ∼ N
[Equation]
o ;
[Equation] [Equation]
This leads to a sampling-free approximation of the log-likelihood
[Equation]
At this point, it is necessary to resolve Σ − 1 x and det( Σ x ) .
Lemma 1 (Woodbury matrix identity) .
The inverse of a rank- m correction of a matrix can be computed in terms of the inverse of the original matrix as
[Equation]
Since Σ x = c I + J h
( x ) J ⊤ , and h ( x ) is a diagonal matrix, we can set U =
J p h ( x ) such that
[Equation]
As the computation of the covariance via the Jacobian directly produces a low-rank decomposition of the covariance matrix and because m ≪ n , this allows for a fast and space-efﬁcient computation.
Thus,
[Equation]
which only requires inverting an m × m matrix (innermost parentheses).
Lemma 2 (Matrix determinant lemma) .
The determinant of a rank- m correction of a matrix can be computed in terms of the inverse and determinant of the original matrix as
[Equation]
With U deﬁned as above, we can see that
[Equation]
As we can now compute Σ − 1 x and det( Σ x ) very efﬁciently, we can extend the log-likelihood as
E
[Equation]
Here, as the inverse of the covariance matrix ( Σ − 1 x of size n × n ) does not need to be computed explicitly and the factorization of the covariance can be multiplied with the error, only operations on small matrices need to be performed.
Speciﬁcally, it requires one multiplication of a vector of size n with a matrix of size m × n , the inversion of an m × m matrix, and the multiplication of vectors from both sides with an m × m matrix.
This signiﬁcantly reduces the computation cost.
Based on these insights
[Equation]
we can write the ELBO as
[Equation]
Since the training objective is to maximize the ELBO, the corresponding loss is L = − ELBO .
3. S TABILIZING T RAINING AND R EDUCING B IAS BY R EGULARIZATION
In practice, for covariance matrices with large variances, the ratio between the largest and smallest eigenvalue can become extremely large even with Tikhonov regularization (e.g. 10 8 ).
Thus, training is not feasible because the gradients in some or most directions are vanishingly small.
To account for this, we propose a regularization that works independently of the ratio between eigenvalues.
For this, we use a mixture between the approximated log-likelihood as described above and the log-likelihood of a conventional auto-encoder (i.e., only considering observation-noise).
Speciﬁcally, we replace Σ −x by γ · Σ − 1 x + (1 − γ ) · ( c I ) − 1 .
This regularization allows for an efﬁcient and stable training, even if training with Σ − 1 x alone would be vanishingly slow or diverges.
This means, we calculate Σ
[Equation] [Equation]
U
[Equation]
I + U
[Equation]
U
[Equation]
instead of Σ − 1 x = 1 c
I − 1 c U
( c I + U ⊤ U ) − 1 U ⊤ (21)
[Figure]
e. γ ) reduces the approximation bias and increases the point of convergence, which improves the model
The model is trained with the biased sampling-free
VAE objective on the CelebA data set and evaluated using the unbiased sampling ELBO estimate.
where γ quantiﬁes the degree to which the variation produced by the decoder is considered.
In our experiments, we found that γ
≈
.avoids vanishingly slow training and divergence due to numerical instabilities.
Empirically, we found that reducing γ can reduce the bias of our sampling-free approximation.
In Figure 1, we train a model using the proposed sampling-free objective with γ and evaluate it using the unbiased sampling ELBO.
We ﬁnd that with γ =0 .
95 , the model converges to an ELBO of − 8000 .
By reducing γ , i.e., increasing the regularization, we can increase the ELBO that the model converges to, which means that we improve the model.
Setting γ =0 .
67 , the ELBO already converges to − 7500 , and after setting γ =0 .
5 it converges to − 7400 .
We ﬁnd
that γ = 0 .
5 works well overall and, thus, we use it as a canonical choice in our experiments.
E XPERIMENTS 4. Q UALITY OF THE T AYLOR S ERIES A PPROXIMATION
[Table]
To evaluate the quality of the
Taylor series approximation in
VAEs, we consider the precision to which variances and covariances are estimated and compare this to different numbers of samples.
For this, we use sampling
VAEs trained on the CelebA and 3D Chairs data sets as described in more detail in the next section and with a latent dimensionality of 16.
We consider two metrics:
The Pearson correlation coefﬁcient of the (pixel-wise) output variances and the Frobenius norms of output covariance matrices.
Here, we compute the true output (co-)variances via a large number of samples ( 15 000 ) and compare it to the (co-)variances produced by our approximation or the (co-)variances produced by a certain number of samples.
Results are shown in Table 1.
We ﬁnd that, even in large deep neural networks, we achieve Pearson correlations of 96 . 7 − 97 .
5% for estimating the variances, which would require 70 − 100 samples in the latent space when relying on sampling.
On the Frobenius norm of covariance matrices, our approximation is (on average) equivalent to more than 235 stochastic samples for the CelebA data set.
Because computing the true covariance matrix is intractable, we estimate it to a precision of 1% of the reported
Frobenius norms.
For the 3D Chairs data set, we omit estimating the true covariance matrices due to computational feasibility.
In addition, we analyzed the ratio between the variances produced by our sampling-free propagation and the sampling approach and found that across data sets this ratio is around 1 . 2 .
This means that the variances are over-estimated by a factor of 1 . 2 by our approximation.
We attribute this to the fact that non-linearities such as ReLU can become constant which is not anticipated by the Taylor approximation.
4. MNIST, C ELEB A, AND 3D C HAIRS
The traversals are in the range of [ − 3 , 3] .
An exhaustive display of all latent dimensions as well as a comparison to sampling VAEs and AEs is provided in the supplementary material.
We evaluate the sampling-free VAE on the CelebA data set [14], on the 3D Chairs data set [15], as well as on the MNIST data set [16].
For the
CelebA and 3D Chairs experiments, we use the same network architecture as Higgins et al.
[3] and train it with the Adam optimizer [42] at a learning rate of 10 − 4 for 10 6 iterations.
We use 12 −latent dimensions for CelebA and 8 − 32 latent dimensions for 3D Chairs.
As in previous works on VAEs, we set c = 0 .
5 which corresponds to a standard deviation of the observation uncertainty of σ = √ 0 .
5 for all methods.
We use a model for the encoder and decoder with 5 convolutional and 1 fully connected layers each.
This is also the architecture that was used in many recent works on sampling VAEs [3, 11] as well as similar to most public implementations.
Further details can be found in the supplementary material.
Qualitative Evaluation In Figure 2, we show the manifold of a sampling-free VAE in comparison to a sampling VAE.
Figure 3 displays qualitative results for CelebA and 3D Chairs.
Here, we display reconstruction results as well as traversals in the latent space for 3 images and 2 latent dimensions.
In the supplementary material, we qualitatively compare the sampling-free VAE to the sampling VAE, the (sampling) β − VAE, as well as to a conventional auto-encoder.
Figure 3 shows that the sampling-free VAE can successfully reconstruct and that latent traversals produce meaningful images.
For the CelebA data set, we show the traversal of the latent dimensions corresponding to, ‘smiling‘, ‘color temperature’, ‘brightness of background’, and ‘hair color’.
For the 3D Chairs data set, the latent dimensions correspond to ‘chair ←→ arm chair’ and ‘color’, ‘wheels’, and ‘thickness of legs and back’.
For the autoencoder on the 3D Chairs data set, even among a variety of hyper-parameters, training always converges to the plain white image as can be seen in the supplementary material.
Quantitative Evaluation To quantitatively evaluate the sampling-free approximate VAE, we compare the Frechet Inception Distance (FID) as well as the estimated marginal likelihood (ML) between a conventional sampling VAE and a sampling-free VAE as shown in Figure 4.
The FID quantiﬁes how close the images produced by the VAE’s data generating process are to the original data set as well as the quality of the produced images.
For the FID, the distribution of images sampled from the VAE is compared to the distribution of images in the data set by utilizing embeddings produced by an Inception neural network.
The ML quantiﬁes the quality of a VAE by quantifying how close
[Figure]
posteriors are to the priors as well as considering the reconstruction error under sampling.
For the
ML, for each image, the marginal likelihood is estimated via many samples ( 50 000 ).
We found that for CelebA, for each size of the latent space, the FID of the sampling-free VAE is signiﬁcantly better than the FID of the sampling VAE.
Furthermore, we found that the ML of the sampling-free VAEs is very close to the ML of the sampling VAEs, for some latent dimensionalities, the sampling-free VAE is better, while for others the sampling VAE is better.
Plots for the 3D Chairs data set can be found in the supplementary material.
We emphasize that we evaluated the sampling-free VAE also via sampling such that there is no bias in the estimates.
[Figure]
[Table]
[Figure]
[Table]
To evaluate the models based on computational budget, we additionally report results for the sampling-free VAE to match the computational budget of the sampling
VAE’s full training.
Speciﬁcally, we evaluate the sampling-free variant after 650training steps, while the sampling counterpart trains for 1 000 000 steps.
The results are shown in Table 2.
Even in this setting, the sampling-free formulation achieves a better on FID than the sampling VAE on the CelebA data set, however, the estimated log marginal likelihood suffers.
4. C OMPARISON TO OTHER G RADIENT E STIMATE V ARIANCE R EDUCTION M ETHODS
We compare our method to the Sticking-the-Landing (Stick) method by Roeder et al.
[26] and to the importance weighted auto-encoders (IWAE) by Burda et al.
For this, we train a VAE on the
MNIST [16] and Omniglot [43] data sets and for the sampling methods (vanilla sampling VAE, Stick,
IWAE)
we vary the number of samples.
The encoder and decoder of the architecture each have two hidden layers with tanh activation functions and 200 neurons per layer.
We train the model at a batch size of 20 using the Adam optimizer [42].
This follows the setup by Roeder et al.
However, a difference is that we model the image space with Gaussians in contrast to the Bernoulli outputs used by Roeder et al.
[26]; as the method ([26]) readily supports Gaussian outputs, we use it with Gaussian outputs.
As tanh is not piece-wise linear, we use the ﬁrst-order Taylor approximation in this case.
[Table] [Table]
We report the results in Table 3.
We can see that on both data sets, the sampling-free method achieves competitive performance.
The sampling VAE and the Sticking-the-Landing method do not achieve the same NLL even with 200 samples.
Only the IWAE can outperform the sampling-free method on
MNIST with at least 20 samples.
On Omniglot, our method also performs better than IWAE withsamples.
Overall, the results show that the sampling-free approximation performs very well, even when compared to training with large numbers of samples.
4. P OSTERIOR C OLLAPSE
We analyze the effect of a sampling-free approximate formulation of deep VAEs on the problem of posterior collapse.
For this, we use the notion of an ( ǫ, δ ) − collapse of a dimension introduced by
Lucas et al.
It is deﬁned as follows:
A latent dimension i has ( ǫ, δ ) − collapsed if
[Equation]
We follow Lucas et al.
[13] in setting δ = 0 .
We replicate their setting of training a VAE with 2 hidden layers and 200 latent dimensions for 1 000 steps with full batch training on a subset of 1 000 MNIST images.
We report the fraction of dimensions that have ( ǫ, δ ) − collapsed where ǫ = 0 .
δ = 0 .
01 , the reconstruction error, as well as the KL divergence after 100 and 1 000 training steps.
We report these results averaged over 10 runs in Table 4.
We ﬁnd that, for the sampling-free VAE, the reconstruction error is smaller than the reconstruction error for the sampling VAE.
While the sampling VAE has a posterior collapse in 60% of the latent dimensions, the sampling-free VAE suffers from posterior collapse in only 0 .
6% of the latent dimensions.
Simultaneously, the KL-Divergence is lower for the sampling-free VAE.
These results indicate that the sampling-free VAE is more robust against ( ǫ, δ ) − collapse than the sampling VAE.
4. R UNTIME A NALYSIS
Finally, we discuss the runtime of the proposed approach and compare it to its sampling counterpart.
For this analysis, we compare the training times of both methods on the CelebA data set for 10steps at a batch size of 64 and 16 latent dimensions on a single Quadro RTX 6000 GPU.
For the sampling VAE we measure an average training time of 9 .
7 hours, while the proposed sampling-free formulation takes on average 14 . 5 hours on the same GPU.
Note that, in case of the sampling-free
VAE, the computationally most expensive component is computing the Jacobian.
Thus, if the hidden dimension would be very large, the computational overhead would scale linearly ( O ( m ) ) with the dimensionality of the hidden space m .
Note that the overhead of our method is small in comparison to the (in the number of neurons) exponential overhead of computing the exact sampling-free ELBO.
D ISCUSSION & C ONCLUSION
We proposed a novel approximate sampling-free approximation of Gaussian VAEs.
Our formulation is reasonably fast to compute, i.e., it can be computed in linear time in the number of latent dimensions, contrasting prior art which had an exponential complexity for computing the exact objective.
We found that the proposed formulation achieves a signiﬁcantly lower variance in estimating the objectives than the sampling approach.
However, while the sampling approach is unbiased, our lower-variance approximation is not unbiased.
Nevertheless, our approximation learns better Gaussian VAEs than its sampling counterpart and can even compete when the computational budget is ﬁxed.
Furthermore, we observe that the sampling-free VAE is more robust against posterior collapse than the sampling VAE.
We hope to have paved the way for future research to analyze and understand the behavior of VAEs.
R EPRODUCIBILITY
In this work we only use public data sets.
We implemented all experiments in PyTorch [44] and our implementation will be made publicly available upon publication.
Network architectures and training hyperparameters are speciﬁed in the supplementary material
B I MPLEMENTATION D ETAILS
In the following, we discuss implementation details including network architectures.
Throughout the experiments in this work, we use c = 0 .
5 , which coincides with many public implementations (which omit the prefactor of 12 c / set this to 1 ).
B. C ELEB A AND 3D C HAIRS
Network Architecture For the CelebA and 3D Chairs experiments, we use the architecture by
Higgins et al.
The encoder consists of ﬁve convolutional layers with kernel size 4 × 4 , a stride of 2 , and (32, 32, 64, 64, 256) latent dimensions with ReLU activations.
Following that, a fully-connected layer mapping to the latent space.
The explicit architecture is: Conv 32x4x4 (stride 2), ReLU,
Conv 32x4x4 (stride 2), ReLU, Conv 64x4x4 (stride 2), ReLU, Conv 64x4x (stride 2), ReLU, Conv 256x4x4 (stride 1), ReLU, FC 256 → latent_dim .
The decoder is the deconvolutional reverse of the encoder.
In the qualitative evaluation, we use a latent dimension of 32 for CelebA, and a latent dimension of 16 for 3D Chairs.
In the quantitative evaluation, we use the range of 12 − 64 for CelebA, and the range of 8 − 16 for 3D Chairs.
Training We trained the CelebA and 3D Chairs models for 10 6 iterations with the Adam optimizer [42] at a learning rate of 10 − 4 .
B. MNIST
Network Architecture Here, we use a fully connected network architecture.
The encoder consists of three fully connected layers with hidden dimensions ( 512 , 256 ) and with ReLU activations.
The decoder is the reverse of the encoder.
For the 2D traversal experiment, we use a latent dimension ofand for the experiments examining posterior collapse a latent dimension of 200 .
Training We trained the MNIST model for the 2D traversal for 5 · 10 4 iterations with the Adam optimizer [42] at a learning rate of 10 − 4 .
For the experiments examining posterior collapse, we trained for 10 3 iterations / epochs at full batch training on a reduced data set of 1 000 random samples and with a learning rate of 10 − 3 .
C A DDITIONAL R ESULTS C. Q UANTITATIVE E VALUATION FOR 3D C HAIRS
See Figure C.1 for the quantitative performance of the sampling-free VAE compared to a sampling
VAE for different latent dimensionalities.
16 # of latent dimensions
F I D
Sampling-Free VAE Sampling VAE        16 # of latent dimensions 7079.
l o g M L
[Figure]
C. S TANDARD D EVIATIONS
[Table]
C. G AP BETWEEN UNBIASED ELBO AND BIASED ELBO
[Figure]
[Figure]
| {z } γ
[Figure]
| {z } γ =0 .| {z } γ =0 .|
{z } γ =0 .
Sampling-Free ELBO
Sampling ELBO
C. F URTHER E XAMPLES , F ULL L ATENT T RAVERSALS , AND C OMPARISON TO O THER M ETHODS
In Figures C.3–C.4, we show reconstructions and full latent traversals for 3D Chairs for the samplingfree VAE, the sampling VAE, the sampling β − VAE, and an auto-encodere.
In Figures C.5–C.8, we show the same for CelebA.
For the auto-encoder on the 3D Chairs data set, even among a variety of hyper-parameters, training always converges to the plain white image.