Annealed Stein Variational Gradient Descent
Abstract
Particle based optimization algorithms have recently been developed as sampling methods that iteratively update a set of particles to approximate a target distribution.
In particular
Stein variational gradient descent has gained attention in the approximate inference literature for its ﬂexibility and accuracy.
We empirically explore the ability of this method to sample from multi-modal distributions and focus on two important issues: (i) the inability of the particles to escape from local modes and (ii) the ineﬃcacy in reproducing the density of the diﬀerent regions.
We propose an annealing schedule to solve these issues and show, through various experiments, how this simple solution leads to signiﬁcant improvements in mode coverage, without invalidating any theoretical properties of the original algorithm
1. Introduction
There have been many recent advances on the theoretical properties of sampling algorithms for approximate inference, which changed our interpretation and understanding of them.
Particularly worth mentioning is the work of Jordan et al.
(1998), who reinterpret Markov
Chain Monte Carlo (MCMC) as a gradient ﬂow of the KL divergence over the Wasserstein space of probability measures.
This new formulation not only allowed for a deeper understanding of these methods but also inspired the inception of new and more eﬃcient inference strategies.
Following this direction, Liu and Wang (2016) recently proposed the
Stein Variational Gradient Descent (SVGD) to perform approximate Wasserstein gradient descent.
This method belongs to the more general family of particle optimization variational inference (POVI), where a continuous density p ( x ) is approximated by a set of n particles that evolve over time towards the target.
However, a solid understanding of its behavior in the ﬁnite particle limit beyond the mean ﬁeld convergence analysis (Duncan et al., 2019) remains elusive.
What is more, there is empirical evidence that SVGD suﬀers from a degeneracy that compromises the particle diversity under these conditions, making them collapse to a small number of modes (Zhuo et al., 2018).
In the following, we discuss how an annealing strategy can signiﬁcantly mitigate this issue, encourage exploration of the signiﬁcant modes, and yield better samples from the target density than standard SVGD.
1.1. Related work
Introducing an artiﬁcial temperature parameter is a common practice in many approximate inference methods.
Indeed, annealing approaches have been shown to be beneﬁcial in both sampling and optimization problems for highly non-convex objectives.
In the context of
Markov chain Monte Carlo, sampling at diﬀerent temperatures enhances the mixing times of the chains and thus allows for faster convergence (Marinari and Parisi, 1992; Geyer and
Thompson, 1995).
In Bayesian inference, a temperature parameter has been introduced to anneal the likelihood or the posterior of the model (e.g., Wenzel et al., 2020) to escape from poor local minima.
A similar eﬀect can be obtained in variational inference and its stochastic counterpart by tempering the KL term (Mandt et al., 2016; Huang et al., 2018;
Fu et al., 2019).
However, the impact of similar annealing strategies on the SVGD method have not yet been studied.
In our work, we show that they can be useful to improve the particle diversity and overcome the mode-collapse problem.
2. Background
Stein variational gradient descent (Liu and Wang, 2016) is a technique to perform approximate inference using a set of particles q t ( x ) = 1 n (cid:80) ni =
δ x i ( t ) , with δ x i being the Dirac measure on particle x i , to approximate a positive density function p ( x ) on X ∈ R d .
More precisely, SVGD is an eﬃcient numerical technique to discretize the Wasserstein gradient ﬂow of the Kullback-Leibler (KL) divergence functional on a new metric called the Stein geometry (Duncan et al., 2019).
SVGD considers an incremental transformation given by an inﬁnitesimal perturbation of the identity matrix T ( x ) =
x + (cid:15)φ ( x ) to move the particles from the initialization to the target.
Here, φ ( x ) is the direction of the perturbation and (cid:15) the step size.
The former is chosen to maximally decrease the KL divergence between the discrete density of the particles and the ﬁnal target.
As shown in Liu and Wang (2016), closed-form solutions can be obtained when restricting all perturbations φ to be from the unit ball of a vector valued reproducing kernel Hilbert space (RKHS) H d
= H 0 × ...
× H 0 .
Here, H 0 is a scalar-valued
RKHS associated with a scalar positive deﬁnite kernel k ( x, x (cid:48) ; h ), and h is the set of kernel hyperparameters.
The direction of steepest descent that maximizes the negative gradient of the KL divergence is then given in closed form as:
[Equation]
with A p φ ( x ) =
φ ( x ) ∇ x log p ( x ) (cid:62) + ∇ x φ ( x ) being the Stein operator.
Using this, we can build an iterative procedure that transforms the initial reference distribution q 0 to the target posterior.
Practically, we draw a set of particles { x 0 i } ni =1 with x 0 i ∼ q 0 and subsequently update them using the optimal perturbation in (1):
[Equation]
with
[Equation]
Conceptually the update rule attracts the particles to high density regions of the target via the average score function (driving force in (2)) while the repulsive force pushes them away from each other.
This avoids a collapse to the MAP estimate and allows a certain degree of diversity among the particles to encourage the exploration of multiple modes and a more faithful reﬂection of the variance of the target distribution.
2.1. Mode-collapse in SVGD
Despite its theoretical foundations, it has been shown empirically (Zhuo et al., 2018) and theoretically (Zhang et al., 2020) that the particles in SVGD tend to collapse to a few local modes and that this eﬀect is strongly related with the initial distribution of the particles.
This issue is also clearly visible in our experiments and seems to already be relevant in low-dimensional problems.
Indeed, as shown in Figure 1, it even happens in the case of a one-dimensional mixture of ﬁve Gaussians.
Here, all particles, independent of the choice of the kernel bandwidth (see Appendix B), end up in the mode closest to the initialization without any possibility of escaping.
Additionally, we noticed in our experiments a connection of how a proper reconstruction of the target mass in the diﬀerent modes is related to the initialization, as seen in Figure 2.
In other words, even if the standard SVGD is able to capture diﬀerent modes, the particles are not correctly distributed to faithfully reproduce the original mass.
Instead, the majority of them tend to end up in the mode closest to their initialization.
It is not clear whether this particular issue is a consequence of the modeseeking limitation of all KL-divergence-based inference methods or of the approximation due to a ﬁnite number of particles.
Empirically evident instead is that a deterministic update of the samples, like the one characterizing SVGD, in combination with a random initialization, can lead to a catastrophic convergence.
In comparison, other methods characterized by a stochastic update like SGLD (Welling and Teh, 2011) can instead always rely on a nonzero probability of escaping from a certain mode, given by the injected noise, and consequently mitigate a bad initialization and encourage exploration.
Moreover, it has been shown by Zhang et al.
that injecting noise is also beneﬁcial for the SVGD update to overcome the mode-collapse issues.
Finally, it is well established that distance metrics and corresponding kernel-based methods suﬀer from the curse of dimensionality (Aggarwal et al., 2001; Reddi et al., 2014).
That is why we should not hope for this problem to disappear in higher dimensions, but rather expect it to become even worse.
3. Annealed SVGD
Motivated by related work that mitigates similar mode-collapse issues in MCMC methods by tempering the target density to speed up the mixing time and avoid chains to get stuck in a single mode (Neal, 1996), we propose to introduce an annealing schedule in the SVGD update (A-SVGD).
This modiﬁcation keeps the deterministic nature of the method but is essential to enhance its exploration and mode coverage and ideally would compensate for the limitation of the initialization and the ﬁnite number of particles.
We introduce an annealing parameter γ ( t ) ∈ [0 , 1] depending on the current iteration step and modify the update rule from (2) in the following way:
[Equation]
Intuitively, we can observe two phases in the time evolution of the particles by varying γ in the interval [0 , 1] with an appropriate schedule:
The ﬁrst phase is exploratory with a predominant repulsive force that pushes the particles away from the initialization and thus allows for a good coverage of the target distribution’s support.
The second phase is exploitative, where the driving force takes over and shrinks the distribution of the particles to the area around the diﬀerent modes.
From a statistical perspective, our modiﬁcation corresponds to the introduction of a temperature parameter T ( t ) = 1 γ ( t ) which rescales the target distribution p ( x ) 1 T ( t ) during the evolution of the approximating density.
It is important to notice how the choice of the annealing schedule is fundamental to preserve the convergence properties of SVGD and to keep the ﬁnal target density unchanged.
We ensure this by formulating the annealing schedule in such a way that the ﬁnal iterations are always performed for the true target density, that is, lim t →∞ γ ( t ) = 1.
From this point of view, our alternative method is formally equivalent to a better parametrization of the initial reference distribution q 0 of the particles that depends on the target distribution.
This is due to the fact that even if in the exploratory phase, when the repulsion is dominant, we still have a small component of the driving force that ensures that the particles are not randomly driven far away from the initialization, but are still driven towards high-density regions.
3.1. The annealing schedule
As mentioned before, the choice of the annealing schedule is fundamental for the convergence to multiple modes while yielding samples from the proper target distribution.
For this purpose, we introduced and tested diﬀerent annealing schedules as shown in Figure 3.
The simplest idea is a linear annealing on the interval [0 , 1]; however, our experiments showed that this choice is not optimal.
Indeed, linear tempering of the density leads to slow particle dynamics that are beneﬁcial to neither the exploration nor the convergence.
For this reason we chose to make the transition between the two inference phases steeper, so that the majority of the evolution happens in one of the two phases and not in between.
To do so, we construct the annealing schedule using the hyperbolic tangent: γ ( t ) = tanh (cid:2) (1 . 3 tT ) p (cid:3) , with t being the current time step and T the total number of steps.
Despite its good exploration, this method might encourage modes very far from the initialization due to the ﬂattening of the target density in the initial exploratory phase.
To achieve a tradeoﬀ and have a more “target-guided” exploration, we follow the cyclical annealing schedule idea proposed in Loshchilov and Hutter (2016) and Huang et al.
(2017), which has already been used for
MCMC sampling in Zhang et al.
We adapted this technique to have a sequence of
C cycles of exploratory and converging phases, obtaining the following expression:
[Equation]
with T being the total number of time steps, C the number of cycles and p an exponent determining the speed of the transition between the two phases (as shown in Figure 3).
4. Experiments
We demonstrate the advantages of our method in several synthetic experiments.
In all the experiments we used SVGD with a standard RBF kernel.
Univariate Gaussian mixture.
We ﬁrst assessed the ability of A-SVGD to sample from a multi-modal univariate distribution given by a mixture of ﬁve Gaussians.
The step size was ﬁxed to (cid:15) = 0 .
and we used the hyperbolic annealing schedule.
We see in Figure 1 that our proposed method successfully covers all the modes, while the standard SVGD collapses to just a single mode.
Bivariate regular Gaussian mixture.
Secondly, we tested our method on a 2D mixture of 16 Gaussians with means equally distributed on a 4 × 4 grid and standard deviation σ = 0 .
In this experiment we used the cyclical annealing schedule from (4).
As reported in Figure 4, we observe that the standard SVGD gets trapped in four of the modes, neighboring the initialization.
In contrast, our method is able to ﬁnd and characterize all modes, independently of the initial position.
Bivariate irregular Gaussian mixture.
In our last experiment we studied the ability of SVGD to reproduce the weights of the mixture components of 2D Gaussians.
The cyclical annealing schedule is used for the A-SVGD and for the standard SVGD we show two diﬀerent initializations to illustrate their impact.
We see in Figure 4 that our proposed method not only covers all the modes, but also approximately recovers their mixture weights, while the standard SVGD again collapses to the modes that are closest to its initialization.
5. Conclusion
In this work, we discussed the mode-collapse issue of SVGD for approximate inference and proposed an annealing strategy to overcome these limitations.
We illustrated the impact of the initialization on a deterministic sampling algorithm like SVGD, highlighting two major drawbacks, namely (i) a tendency of the particles to fall into the neighboring modes without any possibility of escape and (ii) the diﬃculty of the particles in reproducing the eﬀective local density of any given mode.
We found that the introduction of a temperature parameter and an annealing schedule can help alleviate these undesirable behaviors, leading to better samples that can eﬀectively capture multi-modal densities.
Appendix A. Additional results on multi-modal data
In extension to the multi-modal synthetic data presented in section 4 we present here the extreme case for which the initialization is exactly in one of the mode of the target distribution.
As shown in Figure A.1, the SVGD with annealing is remarkably able to escape from the initialization, covering all the modes characterizing the target density.
On the other hand, the standard SVGD is not able to model anything besides the mode in which the particles have been initialized.
Appendix B. Diﬀerent RBF kernel bandwidths
We compared the eﬀect of diﬀerent bandwidths for the RBF kernel to show if this parameter aﬀects the issues illustrated in 2.1.
We used the synthetic multi-modal data from Figure 2 to test the following bandwidth values h ∈ { 0 . 001 , 0 . 01 , 0 .
1 , 1 , 10 , 100 , median } where median is the median heuristic.
As illustrated in Figure B.1, none of the bandwidths shows signiﬁcant improvement and, in contrast to our A-SVGD, the inference is still limited to the neighboring modes.
Appendix C. Diﬀerent annealing schedules
We compared the diﬀerent annealing schedule proposed in section 3.1 on the bivariate irregular Gaussian mixture as reported in Figure C.1.
We also computed the maximum mean discrepancy (MMD)
(Gretton et al., 2012) during the evolution of the particles as reported in Figure C.2.