α -VAEs : Optimising variational inference by learning data-dependent divergence skew
Abstract
The skew-geometric Jensen-Shannon divergence (cid:0)
JS G α (cid:1) allows for an intuitive interpolation between forward and reverse Kullback-Leibler (KL) divergence based on the skew parameter α .
While the beneﬁts of the skew in JS G α are clear— balancing forward/reverse KL in a comprehensible manner—the choice of optimal skew remains opaque and requires an expensive grid search.
In this paper we introduce α -VAEs, which extend the JS G α variational autoencoder by allowing for learnable, and therefore data-dependent, skew.
We motivate the use of a parameterised skew in the dual divergence by analysing trends dependent on data complexity in synthetic examples.
We also prove and discuss the dependency of the divergence minimum on the input data and encoder parameters, before empirically demonstrating that this dependency does not reduce to either direction of KL divergence for benchmark datasets.
Finally, we demonstrate that optimised skew values consistently converge across a range of initial values and provide improved denoising and reconstruction properties.
These render αVAEs an efﬁcient and practical modelling choice across a range of tasks, datasets, and domains.
1. Introduction
As variational inference (VI) progresses, state of the art
Variational AutoEncoders (VAEs) increase in complexity (Vahdat & Kautz, 2020; Child, 2020) while continuing to maximise the Evidence Lower BOund (ELBO) (Blei et al., 2017).
Compared to generative adversarial networks (GANs) (Goodfellow et al., 2014), VAEs necessitate less stringent and problem-dependent training regimes and com- pared to autoregressive models (Larochelle & Murray, 2011;
Germain et al., 2015) (that can be interpreted as instances of very deep VAEs (Child, 2020)) are less computationally expensive and more efﬁcient to sample.
VAE learning requires optimisation of an objective which balances the quality of decoded reconstructions from encoded representations, with a regularising divergence term penalising latent space deviations from a prior distribution.
To couple VAEs’ learning with an appropriate regularisation, it is necessary to consider their underlying assumptions. VAEs often assume latent variables to be parameterised by a multivariate Gaussian, p θ ( z )
= N ( µ, σ 2 ) with z, µ, σ ∈ R n , approximated by q φ ( z | x ) with x ∈ R m and n ≤ m .
For instance, in the original VAE (Kingma & Welling, 2014), KL divergence (Kullback & Leibler, 1951) naturally constrained the variational distribution to an isotropic Gaussian unit ball KL ( q φ ( z | x ) (cid:107) N (0 , I )) , despite unfavourable properties (Bishop, 2006), such as unboundedness and asymmetry.
Moreover, KL does not capitalise on the full ﬂexibility of the wider family of exponential distributions, a recent direction which has tightened the ELBO (Brekelmans et al., 2020; Masrani et al., 2019) and rendered VAE divergence regularisation more interpretable in distribution space via skew-geometric JensenShannon (cid:0)
JS G α (cid:1) divergence (Nielsen, 2019; Deasy et al., 2020).
We henceforth refer to VAEs with skewed geometric divergences as α -VAEs, subsuming β -VAEs.
Divergence skew in VAEs balances the contrasting properties of forward and reverse KL (such as zeroavoidance/forcing) and circumvents opaque divergence terms by interpolating between them (see Chapter 10 in (Bishop, 2006)).
However, an expensive grid search over skew values ﬁxed through training is necessary to optimise for tasks such as image reconstruction.
This is particularly problematic as the link between optimal skew and dataset properties is not clear and is not easily resolved.
Moreover, when skew is treated as static, the divergence constraint does not change during training and therefore does not reﬂect improvements in the encoder and its embeddings.
Instead, an improved optimisation would update the divergence skew relative to these factors without using prior knowledge or compromising performance.
[Figure]
[Figure] [Figure] c
[Figure]
n
[Figure]
u c t i o n l o ss [Figure] JS G JS G *
[Figure]
[Figure]
used an expensive grid search to ﬁx divergence skew through training.
Our separate optimisation of α in JS G α (preﬁx t reconstruction
Contributions.
To overcome these issues, we extend the work of Deasy et al. and introduce α -VAEs, by allowing for data- and encoder-dependent skew.
Our ﬁndings indicate trends in ﬁnal skew values which are dependent on proxies for data complexity in synthetic examples.
We demonstrate that optimal divergence skew tends toward a balance of forward and reverse KL as dimensionality increases and, for increasing distribution modes, training skew marginally favours forward KL.
In the higher dimensional setting of standard image benchmark datasets, we then establish that ﬁnal skew values converge across seeds and are consistent for a range of initial values.
We further exhibit that learning skew in α -VAEs has a positive impact on test set reconstruction loss (summarised in Figure 1), reconstructing denoised images from noisy images, and explain both advantages from the rate-distortion perspective (Alemi et al., 2018).
Overall, we show that α -VAEs with learnable α consistently outperform forward and reverse KL independent of dataset, encoder parameters, and initial skew 1 . 2. The JS G α divergence family For distributions P and Q of a continuous random variable
Z =
[ Z 1 , . . . , Z n ] T , the forward KL divergence (Kullback & Leibler, 1951) is deﬁned as [Equation] where p and q are the probability densities of P and Q respectively, z ∈ R n , and reverse KL divergence refers to
KL ( Q (cid:107) P ) .
Reverse KL from a standard normal distribution N 2 (0 , I ) to a diagonal multivariate normal distribution N 1 ( µ, Σ) , µ ∈ R n and Σ ∈ R n × n , is used throughout variational models (Higgins et al., 2017; Kingma & Welling, 2014; Neal, 2012) and is known to enforce zero-avoiding hyperparameters on
N 1 when minimised (Bishop, 2006; Murphy, 2012).
On the other hand, the forward KL divergence is known for its zeroforcing property (Bishop, 2006; Murphy, 2012).
However, there exist well-known drawbacks of the KL divergence, such as no upper bound leading to unstable optimisation and poor approximation (Hensman et al., 2014), as well as its asymmetric property KL ( P (cid:107) Q ) (cid:54) =
KL ( Q (cid:107) P ) .
Underdispersed approximations relative to the exact posterior also produce difﬁculties with light-tailed posteriors when the variational distribution has heavier tails (Dieng et al., 2017).
One attempt at remedying these issues is the well-known symmetrisation, the Jensen-Shannon (JS) divergence (Lin, 1991)
[Equation]
Although the JS divergence is bounded and offers some intuition through symmetry, it includes the problematic mixture distribution p + q 2 .
This term means that no closed-form expression exists for the JS divergence between two multivariate normal distributions using Equation (2).
Recently, (Nielsen, 2019) and (Nishiyama, 2018) have proposed a further generalisation of the JS divergence using abstract means ( quasi-arithmetic means (Niculescu & Persson, 2006), also known as Kolmogorov-Nagumo means ).
By choosing the weighted geometric mean G α ( x, y ) =
x α y 1 − α for α ∈ [0 , 1] , and using the property that the weighted product of exponential family distributions (which includes the multivariate normal) stays in the exponential family (Nielsen
Garcia, 2009)
, we have the divergence
[Equation]
JS G α , the skew-geometric Jensen-Shannon divergence , between two multivariate Gaussians N 1 ( µ 1 , Σ 1 ) and
N 2 ( µ 2 , Σ 2 ) then admits a closed form
[Equation] with the equivalent dual divergence being [Equation] where N α has parameters
[Equation]
In simple terms, these divergences measure a weighted arithmetic mean of divergences from/to the prior/variational distribution to/from a weighted geometric mean of distributions .
We proceed by examining why this skew parameter α should be learnt (replacing ψ in Equation (14)), what this means in divergence space, and where the optimal α lie. 3. α -VAE optimisation 3.1. Learning divergence parameterisations is not constrained optimisation In order to ﬂexibly interchange divergence terms regularising the latent space of VAEs, it is common to formulate VAE training as a constrained optimisation problem (Rezende & Viola, 2018; Higgins et al., 2017).
A suitable objective to maximise is the marginal (log-)likelihood of the observed data x ∈ R m as an expectation over the distribution of latent factors z ∈ R n
[Equation]
The latent representation can be controlled by imposing an isotropic unit Gaussian constraint on the prior p ( z ) =
N (0 , I ) , arriving at the constrained optimisation problem
[Equation]
where ε dictates the strength of the constraint and D is a divergence.
Equation (12) is then rewritten as a Lagrangian under the KKT conditions (Karush, 1939; Kuhn & Tucker, 2014), obtaining
[Equation]
However, here we are interested in learning properties of the divergence itself, rendering Equation (12) and (13) invalid.
In this setting, the constrained optimisation problem is no longer well-posed, as the constraint becomes part of the optimisation.
Instead, we choose to relax these optimisation assumptions while maintaining competing objectives, the log-likelihood and the parameterised divergence
[Equation]
where ψ ( x, z ) parameterises D .
Such a formulation only relates back to Equation (13) in a valid manner when changes in ψ map to a family of divergences and redundant cases are avoided.
For instance, it would be particularly useful if, for a given ψ during optimisation, training could still be understood as a reconstruction loss term and a closed-form divergence regularisation term.
It is, therefore, of import to consider: which divergences allow for such a parameterisation, whether they will have non-trivial minima, and whether the subsequent properties of this optimisation generalise or are useful.
To this end, we consider learning α , the skew , in the JS G α family of divergences (Appendix 2) and proceed by highlighting the connection to forward and reverse KL.
3.2. Skew optimisation
From an information geometric viewpoint, the skew parameter’s inﬂuence on the intermediate distribution can be seen as the weighting of two distributions along a path in the statistical manifold.
In (Masrani et al., 2019), the authors consider the geometric path between the variational distribution q φ ( z | x ) and the model distribution p ( x, z ) .
The expensive grid search therein required integration along this path and was later avoided in (Brekelmans et al., 2020) by using a moment-matching approach to learn the optimal point on the path.
In contrast to these works, here we consider skew between q φ ( z | x ) and the prior p ( z ) as a form of regularisation.
Additionally, the property that JS G α interpolates between forward and reverse KL (Deasy et al., 2020) offers a clear mechanism for trading off forward and reverse KL properties.
In VAE optimisation, as we have generalised the overall objective by extending the constrained optimisation problem to competing objectives, it is possible to consider direct parameterisation of ψ from Equation (14).
As long as ψ maps to [0 , 1] , the output is a valid divergence (a member of the JS G α family) with extremes at forward and reverse
KL, so we can consider optimising ψ via gradient descent.
Before attempting the optimisation, we can derive useful properties of JS G α ( P (cid:107) Q ) which clarify how the optimisation should be carried out.
In particular, it is important to understand the behaviour of JS G α with respect to α , as this will dictate whether convergence is possible.
Proposition 1.
For Gaussian distributions P and Q with probability density functions p ( z ) and q φ ( z | x ) respectively, the derivative of JS G α ( P (cid:107) Q ) with respect to α is
[Equation]
with a stationary point at
[Equation]
which is a global minimum in the α dimension of the optimisation.
Proof in Appendix B.1.
A useful sanity check from Equation (16), is that α ∗ is clearly bounded between 0 and 1 due to the non-negativity of KL divergence.
Secondly, as Equation (16) relies on parameters φ , this expression is clearly data dependent, meaning optimal skew shifts during training and should not be ﬁxed as in (Deasy et al., 2020).
Finally, considered as a standalone 1 D optimisation of α , the global minimum means that training divergence skew is convex and should be simple to minimise in practice, with closed-form convergence.
Proposition 2.
The upper bound on the convergence rate of gradient descent to α ∗ in (16) is
[Equation]
where D
= ||
α 1 − α ∗ || 2 2 , α 1 is the initial skew value, and T is the number of optimisation steps.
Proof in Appendix B.2.
In the next section, we ﬁnd practical optimisation of skew in
JS G α ∗ to be well behaved, suggesting that training skew in the dual divergence is also well-posed and does not reduce to an invalid divergence or one of the KL directions—despite an equivalent proof of convexity via differentiation not being obvious, due to a mixture term being outside of the log.
4. Experiments
As the convexity shown in Section 3.2 suggests simple optimisation of α , we directly train α as another parameter of the model, but use a separate optimiser to the model parameters so that our 1D analysis holds.
4.1. Characterising skew optimisation for JS G α ∗
To better understand how dual α -VAEs, α ∗ -VAEs, will behave in the more complex setting of modern variational inference benchmarks, we ﬁrst highlight their properties on synthetic examples.
In Figure 2, we depict different optimal skew values for a ﬁt of a multivariate Gaussian to an underlying additive mixture of multivariate Gaussians with trained
JS G α ∗ skew.
As the divergence integrals are not tractable, we directly optimise the multivariate Gaussian parameters via samples from the data for all divergences.
In both plots of Figure 2, we depict the emergent lowdimensional trends in optimised skew.
For Figure 2a, we ﬁt a 2D Gaussian ball to a 2D additive mixture of Gaussian balls, with an increasing number of components in the mixture.
Whereas, in Figure 2b, we increase the dimension of the ﬁt and keep the ratio of mixtures to dimensions ﬁxed at 5.
As the number of mixture components increases in
Figure 2a, optimal skew for JS G α ∗ decreases, favouring reverse KL as mass becomes more concentrated.
Similarly, in
Figure 2b, the optimal skew for JS G α ∗ also decreases, underlining the need for data-dependent skewed divergences and suggesting learnt skew in dual
α ∗ -VAEs will be consistent and avoid trivial cases for more complex data.
[Figure]
4.2. Benchmark image dataset performance
Divergence skew convergence.
As the expression for the minimum in Equation (16) is encoder dependent, we begin our benchmark dataset assessment by depicting the alpha landscape.
In Figure 3, we see that optimal α decreases during training, before stabilising in the ﬁnal epochs.
This supports the argument that naively ﬁxing skew is not optimal, leads to inferior solutions, and even a hyperparameter search for ﬁxed skew followed by training is not sufﬁcient.
The divergences converge to a consistent minimum.
This conﬁrms that our 1 D assessment of convexity and separate optimisation of α successfully leads to stable skew values (see Figure 4 in the Appendix).
In addition, we exemplify how skew evolves across 5 different training seeds.
These [Table] plots delineate data and encoder dependency, with consistency across seeds which initialise model parameters and training procedures (e.g. dropout and batching), suggesting learnt skew predominantly derives from, and varies between, datasets (see Table 3).
Robustness to noise.
As learning skew allows for problemdependent measurement of the distance to the prior, which better accommodates more dispersed encoded distributions, we tested how learning skew regularises VAEs in various noise settings.
In Figure 5, we present denoising experiments where we add Gaussian noise, N (0 , σ 2 ) , to the normalised input images and clip to [0 , 1] .
Despite degraded intuition surrounding where the skew balance should lie in the noisy setting, Figure 5 clearly demonstrates that our trainable skew in JS G α , and its dual form, consistently provides lower test set reconstruction loss across noise levels.
As a sanity check, we can also verify expected trends in the other divergences, with the less dispersion-friendly reverse
KL performing poorly at higher noise levels as forward KL becomes more appropriate, as well as the expected robust behaviour of MMD at high noise levels (Zhao et al., 2019).
Improved reconstruction as a rate-distortion trade-off.
We test our model’s reconstruction-loss performance in the standard clean-image setting.
Table 1 demonstrates superior reconstruction loss across multiple datasets, outperforming both the naive choice of ﬁxed α = 0 .
5 and KL divergence— the latter by a substantial margin.
We further detail the relationship with ﬁxed α in Figure 1a, Figure 6, the supplementary ﬁgures in Appendix D, observing low reconstruction loss from an arbitrary starting skew.
To explain this performance gain, we plot distortion (MSE) against rate (Figure 7 in Appendix D), measured using the reverse KL divergence for all divergences.
The shift down and to the right for α
-VAEs
mean that they trade off rate for distortion, even more so when skew is learnt, giving higher quality image reconstruction when starting with unknown divergence scaling λ .
5. Conclusion
We recast VAE optimisation as a multi-objective task with competing reconstruction and divergence regularisation terms.
This allowed us to use gradient descent to directly learn the divergence skew used to control the variational distribution q φ ( z | x ) relative to the prior distribution p ( z ) .
The resulting method, α -VAE, was shown to be well-posed as a 1D optimisation with a dependency on both the encoder and the data, has improved reconstruction and denoising properties over standard regularisation techniques, and avoids an expensive grid search over skew values.
Moreover, as α -VAEs also generalise over the β -VAE family, they are a practical, efﬁcient, and unbiased choice of VAE across a range of tasks and domains.
A. Related work
Since its introduction (Nielsen, 2021), JS G α (with α = 0 . 5 ) has been used to decompose and estimate
multi-modal
ELBO loss as regularisation in multi-modal VAEs.
(Deasy et al., 2020) further investigated the potential of JS G α at different skew values, in turn, demonstrating improved reconstruction performance of VAEs.
In contrast, in this paper we propose data-dependent learnable skew, which retains the beneﬁts of using an optimal skew but avoids an expensive grid-search.
This provides both efﬁcient and practical way of training VAEs, leading to improved denoising and reconstruction performance as well as better lossy-compression rates focussed on in (Huang et al., 2020).
In a similar context, our work is related to (Brekelmans et al., 2020), that introduces an optimal integration schedule via dynamic parameter selection when approximating the Thermodynamic
Variational Objective (Masrani et al., 2019).
α -VAEs extend the standard VAEs (Kingma & Welling, 2014; Rezende et al., 2014) paradigm with regularisation constraint inspired by recent work on closed-form expressions for statistical divergences (Nielsen, 2019; Nishiyama, 2018).
In particular, α -VAEs offer a stable and intuitive regularisation mechanism.
This allows optimal interpolation between forward and reverse KL divergence, therefore combating the issue of posterior collapse (Lucas et al., 2019).
In this regard, our work is related to approaches that address this issue through KL annealing during training (Bozkurt et al., 2021; Huang et al., 2018; Burgess et al., 2018).
In a more general sense, this work is also related to other approaches that utilise various statistical divergences and distances for latent space regularisation as an alternative to the conventional KL divergence (Hensman et al., 2014;
Tolstikhin et al., 2018;
Dieng et al., 2017; Zhang et al., 2019;
Zhao et al., 2019; Li & Turner, 2016).
B. Proofs
B.1.
Proposition 1.
Proof.
First, simplifying the divergence
[Equation]
Then, differentiating (18) with respect to α
[Equation]
19) and rearranging gives Equation (16), before differentiating again
[Equation]
demonstrates the global minimum property as KL divergence is always positive.
B.2. Proposition 2.
Proof.
We ﬁrst deﬁne the convexity strength and smoothness of f ( α ) =
JS G α , using its simpliﬁed form in Equation (18).
A function f ( α ) : R n (cid:55)→ R n is λ -strongly convex if
[Equation]
where (cid:23) is a generalised inequality and, in our 1D case, this expression reduces to the second derivative with respect to α , giving
[Equation]
f ( α ) is also β -smooth if dfdα is β -Lipschitz
[Equation]
and, as our f is doubly differentiable, the Mean Value Theorem (Rudin et al., 1976) with g [Equation]
[Equation]
gives the necessary bound and the same constant
[Equation]
We can now substitute into the result from convex optimisation with gradient descent (Bubeck, 2014) that, for a λ -strongly convex and β -smooth function, the optimal step size γ is
[Equation]
and the upper bound on the convergence rate is
[Equation]
where κ = βλ = 1 is the condition number, completing the proof. [Table] [Table]
[Table]
[Table]
[Table]
[Table] [Table] [Table] Bernoulli.
[Table]
[Table]
[Table]
C. Benchmark datasets for skew exploration
Throughout our experiments we evaluate the reconstruction loss (mean squared error) on four standard benchmark image datasets : MNIST, 28 × 28 black and white images of handwritten digits (LeCun et al., 2010); Fashion-MNIST,
28 × 28 black and white images of clothing (Xiao et al., 2017)
; Chairs, 64 × 64 black and white images of 3D chairs (Aubry et al., 2014); dSprites 64 × 64 black and white images of 2D shapes procedurally generated from 6 ground truth independent latent factors (Matthey et al., 2017); CelebA resampled to 64 × 64 × 3 colour images of celebrity faces (Liu et al., 2015).
For fair comparison, we follow Higgins et al.
(2017) by selecting a common neural architecture across experiments.
For consistent analysis, rather than searching within architecture and hyperparameter spaces for the best performing model by some metric, we standardise comparison and characterise the beneﬁt of learning divergence skew.
In terms of model details, we use the architectures speciﬁed in Table 2 throughout experiments.
We pad 28x28x1 images to 32x32x1 with zeros as we found resizing images negatively affected performance, and we use nearest neighbour interpolation to downsample CelebA to be 64x64x3. We use a learning rate of 1e-4 throughout and use batch sizeand 256 for the two MNIST variants and the other datasets respectively. Where not speciﬁed (e.g. momentum coefﬁcients in Adam (Kingma & Ba, 2014)), we use the default values from PyTorch (Paszke et al., 2019). The only architectural change we make between datasets is an additional convolutional (and transpose convolutional) layer for encoding (and decoding) when inputs are 64x64x N instead of 32x32x1.
We train dSprites and CelebA for 50 epochs, and all other datasets for 100 epochs.
All models were trained on one GPU, with type varying between: NVIDIA Titan X,
{+NVIDIA 
+}
Ti, or NVIDIA A100.
D. Further results
[Figure]
[Figure] [Figure] [Figure] R e c o n s t r u c t
i
o
n l o ss
[Figure]
Divergence
[Table]
[Figure]