Sliced Wasserstein Variational Inference
Abstract
Variational Inference approximates an unnormalized distribution via the minimization of
Kullback-Leibler (KL) divergence.
Although this divergence is efficient for computation and has been widely used in applications, it suffers from some unreasonable properties.
For example, it is not a proper metric, i.e., it is non-symmetric and does not preserve the triangle inequality.
On the other hand, optimal transport distances recently have shown some advantages over KL divergence.
To make use of these advantages, we propose a new variational inference method by minimizing sliced Wasserstein distance.
This sliced Wasserstein distance can be approximated simply by running very few MCMC steps without solving any optimization problem.
Our approximation also does not require a tractable density function of variational distributions so that approximating families can be amortized by generators like neural networks.
Experiments on synthetic and real data are illustrated to show the performance of the proposed method.
1. Introduction
Variational inference (VI) is a method that recasts Bayesian inference as an optimization problem where it utilizes Kullback-Leibler (KL) divergence as a measurement to capture the discrepancy of two probability distributions.
Unlike inference methods that utilizes
Monte Carlo Markov Chains which needs to sample from the target probability space, VI is fast and lightweight in terms of computation.
Therefore, it is preferred in many modern machine learning tasks.
Optimal Transport (OT) (Villani, 2009) has recently gained significant attention in the machine learning community.
Compared to KL divergence, OT gives a valid metric that is symmetric and preserves triangular inequality.
It is reported to show good performances in some downstream applications (Arjovsky et al., 2017; Gulrajani et al., 2017).
While OT provides us with a new horizon to some old machine learning scenarios, the original OT problem requires a computationally demanding optimization procedure which impedes the popularity of applying the original optimal transport methods.
To address this difficulty, sliced Wasserstein distance (Bonneel et al., 2015) reduces the computational inefficiency of OT by projecting high dimensional probability distributions into univariate slices where OT has a closed form solution.
Sliced Wasserstein distance is successfully used in many practical tasks (Deshpande et al., 2018; Kolouri et al., 2018b,a) but it has not yet been applied to variational inference tasks.
In this paper, we extend sliced Wasserstein distance to variational inference tasks.
Our methods utilize a small number of steps of MCMC to obtain an estimation of the distance.
The advantage is that by leveraging the sliced Wasserstein distance, our method does not rely on simultaneous adversarial training (Mescheder et al., 2017; Li et al., 2017; Zhang et al., 2020) to estimate the discrepancy but can still perform amortized inference (Gershman and Goodman, 2014), i.e., use a parameterized function as a sampler to capture target distributions. Similar work is contrastive divergence Hinton (2002), in which the gradient approximation is obtained with MCMC. However, our method uses MCMC to estimate the discrepancy between two distributions.
2. Background 2.1. Variational Inference
Given a joint probability distribution p ( x, z ), where z is latent variables and x refers to observations, variational inference is finding a distribution q ϕ ( z ) which approximates the posterior p ( z | x ) as close as possible.
Such approximations can be obtained via minimizing
Kullback-Leibler (KL) divergence
[Equation]
Note that D KL ( q ϕ ( z ) ||
p
( z | x ))
=
if and only if q ϕ ( z ) =
p ( z | x ).
However, Eq(1) is intractable because the posterior p ( z | x ) is unnormalized.
Instead, we can equivalently maximize Evidence Lower Bound (ELBO) L ( ϕ )
[Equation]
Since we observe that the model evidence log p ( x ) is a constant w.r.t.
ϕ
and the above inequality becomes tight if D KL (cid:0) q ϕ ( z ) || p
( z | x ) (cid:1) = 0.
Optimization of ELBO requires the differentiation of the r.h.s.
expectation.
Gradient descent is a standard approach that allows for such optimization.
To obtain a valid estimation of the gradient, a solution is to apply the score function method (Paisley et al., 2012)
(Ranganath et al., 2014).
An alternative solution to obtain the gradient of ELBO is the reparameterization trick (Kingma and Welling, 2014)
(Rezende et al., 2014).
Vanilla VI leverages KL divergence but this can be substituted by any other f -divergence and importance sampling (Jerfel et al., 2021; Wan et al., 2020) can be used to obtain gradient estimation for general f -divergences.
2.2. Wasserstein Distance
Wasserstern distance arises in optimal transport (Villani, 2009) in which a distribution is transformed to another by moving probability mass.
Wasserstern distance measures the cost of such a transformation.
Given two marginal distributions p ( x ) and q ( y ) over domains X and Y , let Π( p, q ) be a set of any coupled joint distributions γ ( x, y ) where (cid:82) γ ( x, y ) dx
= q ( y ) and (cid:82) γ ( x, y ) dy
= p ( x ).
The c − Wasserstein distance is defined as
[Equation]
where ∥ x − y ∥ is a cost function of moving a point from X to Y .
Intuitively, the c − Wasserstein distance aims to find an optimal joint distribution γ ( x, y ) where the expected cost specified by Eq(3) achieves its minimum.
Solving this optimization problem is generally difficult but we can rewrite p − Wasserstein distance in a univariate case as
[Equation]
where F ( · ) is a cumulative distribution function and F − 1 ( · ) is a quantile function of a probability distribution and the composition F − 1 v F u ( · ) defines a transportation map that moves mass from p ( x ) to q ( y ).
Given two empirical distributions, we can simply utilize Eq(4) to estimate c − Wasserstein distance by sorting samples.
2.3. Sliced Wasserstein Distance
Motivated by the computational efficiency of estimating Wasserstein distance with univariate distributions.
We give a brief review of sliced Wasserstein distance (Bonneel et al., 2015).
We first introduce Radon transformation (Beylkin, 1984).
Let h ( · ) be a function h : R d −→ R .
Radon transform is defined as
[Equation]
Eq(5) defines a surface integral on a hyper-plane S : l = ⟨ x, θ ⟩
where l ∈ R and θ ∈ S d −where S
d − 1 is a unit ball embedded in R d .
For any pair of vectors θ and h , we obtain a sliced function h Rθ ( · ).
We note that marginalization of a high dimensional joint probability distribution can be regarded as a special case of the Radon transform with θ = e i , where e i is an all-zero vector with only 1 at the i -th position.
Note that the sliced function yielded by Eq(5) is univariate.
Leveraging this property, we define sliced Wasserstein distance for probability distributions p ( x ) and q ( y ) as the averaged distance of these slices.
[Equation]
Given an empirical distribution described by ˆ p = 1 n (cid:80) ni =
δ x i , it is trivial to write down
its
Radon transformation defined in Eq(5) as ˆ p Rθ = 1 n (cid:80) ni
=1 δ ⟨ x i ,θ ⟩ .
We summarize the procedure of calculating sliced wasserstein distance via empirical samples in Algorithm 1 .
3. Proposed Method
This method is named sliced Wasserstein variational inference (SWVI).
We use the following notation: q ϕ ( z ) as the variational distribution parameterized by ϕ , and p ( z | x ) as the target distribution.
The question we are interested in is finding an optimal parameter ϕ ∗ that minimizes sliced Wasserstein distance between the variational distribution and the target distribution.
[Equation]
Motivated by the intractability of the target distribution p ( z | x ), the main idea of SWVI is to use MCMC methods to estimate the distance between the variational distribution and the target distribution.
Require:
ˆ p = 1 n (cid:80) ni
=
δ x i and ˆ q = 1 n (cid:80) ni
=
δ y i for k = 0 , 1 · · · m
3.1. Estimating Distances with MCMC
Let K ( ·|· ) be a transition kernel of a MCMC with the stationary distribution p ( z | x ), and q ϕ ( z ) be the initial distribution of the corresponding MCMC.
We denote by q t ( z ) the marginal distribution of the Markov chain after applying t times transitions.
[Equation]
where
[Equation]
Given a sufficiently long run, q t ( z ) converges to p ( z | x ) because of the stationary property of
Markov chain.
At the current stage, one can directly evaluate sliced Wasserstein distance
SW c ( p, q ϕ ) via
[Equation]
Unfortunately, running a long enough MCMC chain is time consuming and it might be difficult to diagnose the burn-in period.
To solve this difficulty, we instead evaluate a local distance SW c ( q t , q ϕ ) with a few number of steps t of iterating MCMC algorithms.
Next we update parameters ϕ via gradient descent with the learning rate α
[Equation]
Since q t ( z ) is an improvement of q ϕ , minimizing the sliced Wasserstein distance between them guides the variational distribution q ϕ ( z ) towards the target distribution p ( z | x ).
3.2. Stochastic Optimization
MCMC methods approximate an unnormalized probability distribution in a non-parametric manner, i.e., by particle approximations.
We can thus use Monte Carlo methods to estimate SW c ( q t , q ϕ ), which is described by Algorithm 1 .
The approximation can be done by drawing samples respectively from q ϕ ( z ) and q t ( z ).
Suppose that { z 0 i } i =1 , 2 ··· n ∼ q ϕ
( z ) and { z ti } i =1 , 2 ··· n ∼ q t ( z ).
Sliced Wasserstein is then approximated by
[Equation]
Here we rewrite the approximate distance as a function L ( · , · ) of two sets of samples.
In order to optimize of the parameter of the variational distribution q ϕ ( z ), we still need to reparameterize samples { z 0 i } i =1 , 2 ··· n .
This can be done by an amortized sampler that can be either a parametric probability distribution or a flexible neural network.
The amortized sampler is written as z ( ϕ ) =
g ϕ ( ϵ ) , ϵ ∼ r ( ϵ ), where r ( ϵ ) is a noise distribution and g ϕ is a parametric model.
We can use chain rule to obtain the gradient estimation of Eq(11),
[Equation]
This can be implemented easily via back-propagation.
The overall procedure of SWVI is summarized in Algorithm
Algorithm 2: Sliced Wasserstein Variational Inference (SWVI)
Require: An unnormalized probability distribution p ( z | x ) and learning rate α .
Initialize sampler q ϕ 0 ( z ).
for m = 0 , 1 , 2 · · · s −
Sample { z 0 i } i =1 , 2 ··· n from q ϕ m ( z ) 2.
Run MCMC towards p ( z | x ) with particles initialized at { z 0 i } i
n to get
4. Experiments
For all experiments, we use sliced 1-Wasserstein distance.
4.1. 2D Experiment
In this experiment, we set target distributions as a bivariate Gaussian distribution and a bimodal Gaussian mixture.
The sampler is a Gaussian distribution.
We fit the sampler model to the target distribution with vanilla variational inference under KL divergence and the proposed method SWVI.
In our method, we adopt the random walk Metropolis-Hastings algorithm as our MCMC instance and we run the Markov chain with 5 steps.
The proposal distribution for Metropolis-Hastings algorithm is a standard Gaussian distribution.
In first figure of Figure 1, we use a diagonal Gaussian distribution as the sampler model and it shows that SWVI results in an approximation with a larger variance compared to standard
VI which uses the KL divergence that is mode seeking.
In the second figure, we use a regular Gaussian distribution with full trainable covariance matrix, both VI and SWVI can approximate the target distribution well.
For the 3rd and 4th figures, the target distribution is set as a Gaussian mixture model, VI always fits one mode but SWVI can have different behaviours if we tune the step size of random walk in MCMC (step sizes are 1.0 and 2.5).
4.2. Neural Samplers
In this subsection, we show an experiment where we have amortized SWVI with a neural sampler, where the density function is no longer tractable.
Note that the proposed method in Algorithm 2 does not require a closed form of density function of q ϕ ( z ).
Hence, we can easily adapt a neural sampler that can generate samples from a more flexible distribution.
For comparison, we also implemented amortized Stein variational gradient (SVGD) method (Liu and
Wang, 2016; Feng et al., 2017).
The target distribution is a mixture of two Gaussians but one with a larger variance and another with a smaller variance.
We fit a neural sampler to this distribution with SWVI and amortized SVGD.
This experiment shows that amortized SVGD fails to capture the other mode.
The reason is that the kernel function (RBF) cannot adjust the bandwidth to the two modes with different ranges of variance.
However, with the asymptotic guarantees of MCMC, the neural sampler trained with the proposed method can efficiently capture two different modes and outputs considerably better samples.
4.3. Bayesian Logistic Regression
We apply SWVI to binary classification tasks in the UCI repository (Asuncion and Newman, 2007) with Bayesian logistic regression models.
We use a constant prior distribution for the
Bayesian logistic model and 20 steps Langevin Dynamics (with learning rate 0.0001) (Neal, 2011; Welling and Teh, 2011) as the MCMC instance.
The fitted distributions are both diagonal Gaussians for SWVI and VI.
We present results in Table 1 and it can be seen that the performance of SWVI is on par with the vanilla VI.
[Table]