Extended Abstract Track
Optimal Latent Transport
Editor: Sophia Sanborn, Christian Shewmake, Simone Azeglio, Arianna Di Bernardo,
Abstract
It is common to assume that the latent space of a generative model is a lower-dimensional
Euclidean space.
We instead endow the latent space with a Riemannian structure.
Previous work endows this Riemannian structure by pulling back the Euclidean metric of the observation space or the Fisher-Rao metric on the decoder distributions to the latent space.
We instead investigate pulling back the Wasserstein metric tensor on the decoder distributions to the latent space.
We develop an efficient realization of this metric, and, through proof of concept experiments, demonstrate that the approach is viable.
1. Introduction
Deep generative models tackle the common problem of inferring a probability distribution from example data.
Generative adversarial networks (GANs) (Goodfellow et al., 2020) and variational autoencoders (VAEs) (Kingma and Welling, 2013; Rezende et al., 2014) both learn a lower-dimensional latent space and a mapping from that to the data space.
This latent space Z describes the data as a lower-dimensional manifold embedded in the data space, X ⊂ R D .
In light of this description and since X is an embedded submanifold, it makes sense to treat Z as a manifold rather than a Euclidean space (Hauberg, 2018).
The geometry of Z is determined by the properties of the decoder distributions { p θ ( ·| z ) | z ∈
Z} .
Early work along this direction focus on the case where the decoder p θ ( ·| z ) is Gaussian, and pulls the Euclidean metric from data space into the latent space (Arvanitidis et al., 2018).
As classic differential geometry does not lend itself to stochastic models, Eklund and Hauberg (2019) reinterpret the reparametrization trick (Kingma and Welling, 2013) as a random projection, such that Gaussian decoders can be seen as spanning a deterministic manifold in the product space of decoder mean and standard deviation.
This then allows us to define a Riemannian metric on Z .
Empirically, this approach has been demonstrated to work well on a range of tasks (Arvanitidis et al., 2018; Kalatzis et al., 2020; Beik-Mohammadi et al., 2021), but it does not generalize beyond Gaussian decoders.
In practice, however, Gaussian decoders are rarely used as the associated likelihood is not a good model for e.g. image data.
Instead discretized mixtures of logistics and similar likelihood models are used in contemporary
VAEs (Salimans et al., 2017; Maaløe et al., 2019).
To construct a more general framework, Arvantidis et al.
(2022) recently proposed to pull back an information metric from the space of decoder distributions to the latent space.
Specifically, this work pulled the Fisher-Rao metric over decoders into the latent space.
Extended Abstract Track
While this provides a general approach, the Fisher-Rao metric disregards the original data metric, which often comes with significant physical grounding.
For example, the Fisher-Rao metric over distributions of pixel intensities disregards that small intensity fluctuations are less important than large ones.
In this paper , we investigate the use of Wasserstein metrics over the decoder distributions.
This elegantly combines the observation space metric with the stochasticity of the decoder.
We focus on the discretized mixture of logistics likelihood as this is the most popular choice in the VAE literature.
We propose an efficient algorithm for evaluating this
Wasserstein metric and demonstrate the feasibility of the approach.
2. Background
Variational autoencoders (Kingma and Welling, 2013; Rezende et al., 2014).
In the VAE setting, we have a latent variable model for our data.
z ∈ R d denotes the latent variables in a low-dimensional Euclidean space and x ∈ X ⊂ R D denotes our observation data which lies on a low-dimensional submanifold of a high-dimensional Euclidean space.
The VAE encoder maps points sampled from the dataset to an approximate posterior distribution, denoted by q ϕ ( z | x ), on the latent space and the decoder maps points sampled from the latent space to a decoder distribution, denoted by p θ ( x | z ), on the data space.
We also define a prior on the latent space denoted by p ( z ).
The VAE learns its parameters by maximizing the evidence lower bound (ELBO) which is a lower bound of the log-likelihood:
[Equation]
Latent space information geometry (Arvantidis et al., 2022).
The VAE decoder function, h : Z → H , maps each point in the sample space to a parameter space H of a probability distribution p θ ( x | z ) on X .
In particular we have z (cid:55)→ η and the likelihood is given by p θ ( .
| η ).
The function h that maps latents to parameters allows us to write the likelihood as p θ ( x | z ).
For two points z 1 and z 2 that are arbitrarily close, i.e. z 2 =
z
1 + ϵ , we can then define the distance between them to be
[Equation]
We can define length of curves in the latent space, γ : [0 , 1] → Z to be:
[Equation]
Distances can then be defined as the length of the shortest connecting curves, i.e. d 2 ( z 1 , z 2 ) =
inf γ l ( γ ).
It can be shown that this notion of distance is a geodesic distance on a Riemannian manifold.
In particular, if we consider the space of parameters H equipped with the FisherRao metric, it can be shown that this infinitesimally coincides with the KL divergence,
[Equation]
The decoder h : Z → H is a map where H is equipped with the Fisher-Rao metric.
Then we can equip Z by pulling back the Fisher-Rao metric along h given by: M ( z ) =
J ⊤ h ( z ) I H ( h ( z ))
J h ( z ), where J h is the Jacobian of h .
Extended Abstract Track In the sequel, we will attempt to define a similar geometry on the latent space by considering the Wasserstein distance instead of the KL divergence.
We will then compute geodesics on the latent space, similar to Eq. 3, and assess if they are sensible (Figure 1).
3. Wasserstein Distance
Recently the Wasserstein distance, from the field of optimal transport, has been attracting increasing attention.
Optimal transport theory provides a way of defining a Wasserstein metric tensor, which gives the space of probability distributions an infinite-dimensional
Riemannian differential structure (Ambrosio and Gigli, 2013; Li, 2022) and it can be shown that the Wasserstein distance can be seen as the geodesic distance on this manifold.
An attractive property of the Wasserstein distance is that it reflects the underlying ground metric on the sample space.
This gives us reason to think that pulling back the Wasserstein metric to the latent space would better reflect the geometry since it takes into account both the distance between decoder distributions and the distance in the sample space.
The Wasserstein distance between multi-dimensional continuous probability distributions is generally intractable (Peyr´e et al., 2019).
And while there are algorithms to approximate the Wasserstein distance between multi-dimensional discrete distributions, such as the Sinkhorn algorithm (Sinkhorn and Knopp, 1967), they are computationally expensive.
However, a closed-form formula exists for the special case of Wasserstein-1 distance between 1-dimensional discrete distributions.
In this case, the Wasserstein-1 distance is equivalent to the Earth movers distance (Levina and Bickel, 2001).
For two discrete distributions, p and q of length N the distance is
[Equation]
Furthermore, it can be shown that for product measures, we have
[Equation]
Jointly, Eqs.
5 and 6 allow us to compute and backpropagate through the Wasserstein distance in some specific setting.
This in turn allows us to compute geodesics in the latent space under the pull-back Wasserstein metric.
We specifically consider a VAE with a decoder that outputs a discretized logistic distribution to generate images, similarly to Salimans et al.
(2017) and Maaløe et al.
As is common in VAEs, we consider data dimensions to be conditionally independent, i.e. p θ ( x | z )
= (cid:81)
Di =1 p θ ( x i | z ).
Each pixel is modeled as a discrete random variable taking values in one of 256 bins.
Given two images, the Wasserstein distance between corresponding pixels can be computed using Eq. 5, and the Wasserstein distance between the two images can be computed using Eq. 6.
Then we can use an approximation of the length of a curve as in Eq.
3 to compute the shortest path (geodesic) between two Extended Abstract Track points,
[Equation]
We parametrize curves γ as cubic splines and minimize curve energy using gradient-based optimization.
This is realized with the StochMan library (Detlefsen et al., 2021).
We train a VAE with a discretized logistic decoder on the dataset and compute geodesics on the latent space equipped with the pull-back of the Wasserstein metric.
We train these VAEs on a subset of the MNIST dataset composed of only the digits with some specific label and Z ⊂ R 2 .
Figure 2 shows the geodesics on the latent space learned by VAE computed using Eq.
We see the geodesics lie inside the latent space.
We also train the VAE on the entire MNIST dataset with Z ⊂ R 30 .
We compute a geodesic on this latent space and decode various points lying on this geodesic and we compare it with linear interpolation in the space of images.
As illustrated in Figure 3, the images on the geodesic are plausible members of the dataset, whereas the images on the linear interpolation don’t always lie on the dataset(See Appendix Section A for more examples).
5. Discussion
We have presented an efficient algorithm for computing latent space geodesics under the pull-back Wasserstein metric associated with discrete decoders.
This both allows for having a deterministic metric in a stochastic model and for incorporating the observation space metric, thereby getting the best parts of existing geometries (Arvanitidis et al., 2018;
Arvantidis et al., 2022).
The presented work is, however, early with several aspects still missing investigation.
Previous studies have, both theoretically and empirically, demonstrated that the uncertainty of the decoder plays a crucial role akin to the topology of the learned Extended Abstract Track manifold (Hauberg, 2018; Detlefsen et al., 2022).
We did not investigate this aspect here.
Furthermore, in the present study, we have only investigated how to compute the geodesic connecting two latent points.
While this is important, many other geometric tools are currently missing.
Perhaps most elementary, we have not provided explicit access to the latent space metric.
These, and more questions, will be investigated in future work.
Acknowledgments.
This work was funded in part by the Novo Nordisk Foundation through the Center for Basic Machine Learning Research in Life Science (NNF20OC0062606). It also received funding from the European Research Council (ERC) under the European Union’s Horizon 2020 research, innovation programme (757360). SH was supported in part by research grants (15334, 42062) from VILLUM FONDEN.
The authors also acknowledge the support of the Pioneer Centre for AI, DNRF grant number P
Appendix A. Comparison of geodesics and linear interpolants
We present various examples of comparisons between images lying on a geodesic and the image lying on the linear interpolant between two images.