S AMPLING WITH M OLLIFIED I NTERACTION E NERGY D ESCENT
A BSTRACT
Sampling from a target measure whose density is only known up to a normalization constant is a fundamental problem in computational statistics and machine learning.
In this paper, we present a new optimization-based method for sampling called mollified interaction energy descent (MIED).
MIED minimizes a new class of energies on probability measures called mollified interaction energies (MIEs).
These energies rely on mollifier functions—smooth approximations of the Dirac delta originated from PDE theory.
We show that as the mollifier approaches the Dirac delta, the MIE converges to the chi-square divergence with respect to the target measure and the minimizers of MIE converge to the target measure.
Optimizing this energy with proper discretization yields a practical firstorder particle-based algorithm for sampling in both unconstrained and constrained domains.
We show experimentally that for unconstrained sampling problems, our algorithm performs on par with existing particle-based algorithms like SVGD, while for constrained sampling problems our method readily incorporates constrained optimization techniques to handle more flexible constraints with strong performance compared to alternatives.
I NTRODUCTION
Sampling from an unnormalized probability density is a ubiquitous task in statistics, mathematical physics, and machine learning.
While Markov chain Monte Carlo (MCMC) methods (Brooks et al., 2011) provide a way to obtain unbiased samples at the price of potentially long mixing times, variational inference (VI) methods (Blei et al., 2017) approximate the target measure with simpler (e.g., parametric) distributions at a lower computational cost.
In this work, we focus on a particular class of VI methods that approximate the target measure using a collection of interacting particles.
A primary example is Stein variational gradient descent (SVGD) proposed by Liu & Wang (2016), which iteratively applies deterministic updates to a set of particles to decrease the KL divergence to the target distribution.
While MCMC and VI methods have found great success in sampling from unconstrained distributions, they often break down for distributions supported in a constrained domain.
Constrained sampling is needed when the target density is undefined outside a given domain (e.g., the Dirichlet distribution), when the target density is not integrable in the entire Euclidean space (e.g., the uniform distribution), or when we only want samples that satisfy certain inequalities (e.g., fairness constraints in Bayesian inference (Liu et al., 2021)).
A few recent approaches (Brubaker et al., 2012;
Byrne & Girolami, 2013; Liu & Zhu, 2018; Shi et al., 2021) extend classical sampling methods like
Hamiltonian Monte Carlo (HMC) or SVGD to constrained domains.
These extensions, however, typically contain expensive numerical subroutines like solving nonlinear systems of equations and require explicit formulas for quantities such as Riemannian metric tensors or mirror maps to be derived on a case-by-case basis from the constraints.
We present an optimization-based method called mollified interaction energy descent (MIED) that minimizes mollified interaction energies (MIEs) for both unconstrained and constrained sampling.
An MIE takes the form of a double integral of the quotient of a mollifier —smooth approximation of δ 0 , the Dirac delta at the origin—over the target density properly scaled.
Intuitively, minimizing an MIE balances two types of forces: attractive forces that drive the current measure towards the target density, and repulsive forces from the mollifier that prevents collapsing.
We show that as the mollifier converges to δ 0 , the MIE converges to the χ 2 divergence to the target measure up to an additive constant (Theorem 3.3).
Moreover, the MIE Γ -converges to χ 2 divergence (Theorem 3.6), so that minimizers of MIEs converge to the target measure, providing a theoretical basis for sampling by minimizing MIE.
While mollifiers can be interpreted as kernels with diminishing bandwidths, our analysis is fundamentally different from that of SVGD where a fixed-bandwidth kernel is used to define a reproducing kernel Hilbert space (RKHS) on which the Stein discrepancy has a closed-form (Gorham & Mackey, 2017).
Deriving a version of the Stein discrepancy for constrained domains is far from trivial and requires special treatment (Shi et al., 2021; Xu, 2021).
In contrast, our energy has a unified form for constrained and unconstrained domains and approximates the χ 2 divergence as long as the bandwidth is sufficiently small so that short-range interaction dominates the energy: this idea of using diminishing bandwidths in sampling is under-explored for methods like SVGD.
Algorithmically, we use first-order optimization to minimize MIEs discretized using particles.
We introduce a log-sum-exp trick to neutralize the effect of arbitrary scaling of the mollifiers and the target density; this form also improves numerical stability significantly.
Since we turn sampling into optimization, we can readily apply existing constrained sampling techniques such as reparameterization using a differentiable (not necessarily bijective) map or the dynamic barrier method by Gong & Liu (2021) to handle generic differentiable inequality constraints.
Our method is effective as it only uses first-order derivatives of both the target density and the inequality constraint (or the reparameterization map), enabling large-scale applications in machine learning (see e.g. Figure 4).
For unconstrained sampling problems, we show MIED achieves comparable performance to particlebased algorithms like SVGD, while for constrained sampling problems, MIED demonstrates strong performance compared to alternatives while being more flexible with constraint handling.
R ELATED W ORKS
KL gradient flow and its discretization for unconstrained sampling.
The Wasserstein gradient flow of the Kullback-Leibler (KL) divergence has been extensively studied, and many popular sampling algorithms can be viewed as discretizations of the KL-divergence gradient flow.
Two primary examples are Langevin Monte Carlo (LMC) and Stein variational gradient descent (SVGD).
LMC simulates Langevin diffusion and can be viewed as a forward-flow splitting scheme for the
KL-divergence gradient flow (Wibisono, 2018).
At each iteration of LMC, particles are pulled along −∇ log p where p is the target density, while random Gaussian noise is injected, although a Metropolis adjusting step is typically needed for unbiased sampling.
In contrast with LMC, SVGD is a deterministic algorithm that updates a collection of particles using a combination of an attractive force involving −∇ log p and a repulsive force among the particles; it can be viewed as a kernelized gradient flow of the KL divergence (Liu, 2017) or of the χ 2 divergence (Chewi et al., 2020).
The connection to the continuous gradient flow in the Wasserstein space is fruitful for deriving sharp convergence guarantees for these sampling algorithms (Durmus et al., 2019;
Balasubramanian et al., 2022; Korba et al., 2020; Salim et al., 2022).
Sampling in constrained domains.
Sampling in constrained domains is more challenging compared to the unconstrained setting.
Typical solutions are rejection sampling and reparameterization to an unconstrained domain.
However, rejection sampling can have a high rejection rate when the constrained domain is small, while reparameterization maps need to be chosen on a case-by-case basis with a determinant-of-Jacobian term that can be costly to evaluate.
Brubaker et al.
(2012) propose a constrained version of HMC for sampling on implicit submanifolds, but their algorithm is expensive as they need to solve a nonlinear system of equations for every integration step in each step of HMC.
Byrne & Girolami (2013) propose geodesic Hamiltonian Monte Carlo for sampling on embedded manifolds, but they require explicit geodesic formulae.
Zhang et al.
Ahn & Chewi (2021) propose discretizations of the mirror-Langevin diffusion for constrained sampling provided that a mirror map is given to capture the constraint.
Similarly, Shi et al.
(2021) propose mirror SVGD that evolves particles using SVGD-like updates in the dual space defined by a mirror map.
While these two methods have strong theoretical convergence guarantees, they are limited by the availability of mirror maps that capture the constraints.
Liu et al.
(2021) extend SVGD to incorporate a population constraint over the samples obtained; their setting is different from ours where the constraint is applied to every sample.
Pairwise interaction energies on particles.
Pairwise interaction energies on particles take the form
E ( { x i } Ni =1 ) =
(cid:80) i ̸
=
j W ( x i , x j ) for a kernel W ( x, y ) .
The gradient flow of E on N particles can give rise to phenomena like flocking and swarming, and a long line of mathematical research studies mean-field convergence of such particle gradient flow to continuous solutions of certain PDEs as
N → ∞ (Carrillo et al., 2014; Serfaty, 2020).
Of particular interest to sampling, a separate line of works summarized by Borodachov et al.
(2019) demonstrates that for the hypersingular Riesz kernel W ( x, y ) =
∥
x − y ∥ − s 2 with sufficiently big s , minimizing E over a compact set yields uniform samples as N → ∞ .
Moreover, W can depend on p to obtain non-uniform samples distributed according to p .
As the hypersingular Riesz kernel is not integrable, their analysis is entirely based on geometric measure theory and avoids any variational techniques.
In contrast, we take a different approach by using integrable kernels in MIEs through mollifiers.
This enables us to apply variational analysis to establish interesting connections between MIEs and χ 2 divergence.
We compare our formulation with theirs further in Appendix B. Joseph et al.
(2019) propose to minimize the maximum of pairwise interaction akin to the interaction in Borodachov et al.
(2019) without gradient information using a greedy algorithm.
Recently, Korba et al.
propose sampling via descending on kernel Stein discrepancy which can also be viewed as a type of interaction energy, but like SVGD, it is limited to unconstrained sampling and can be slow as higher-order derivatives of log p are required.
Craig et al.
(2022) propose approximating the χ 2 divergence with a functional different from ours that also involves a mollifier.
The resulting algorithm needs to evaluate an integral containing the target density over the whole domain at each step which can be costly.
Consequently, their method is only applied to special target densities for which the integral has a closed form.
In comparison, our method can handle general target densities with cheap updates in each step.
M OLLIFIED INTERACTION ENERGY
Notation.
Let X ⊂ R n be the domain we work in.
We use L n to denote the Lebesgue measure on R n and write (cid:82) f ( x ) d x to denote integration with respect to L n .
We will assume all measures are Borel.
We use H d to denote the d -dimensional Hausdorff measure.
We use ω N to denote a set of points { x 1 , . . .
, x N } ⊂ R n .
For x ∈ R n , let δ x be the Dirac delta measure at x and δ ω N be the empirical measure 1 N (cid:80)
Nk
=1 δ x k .
We denote L p ( X ) :
= { f :
X → R Lebesgue measurable with (cid:82) X | f ( x ) | p d x < ∞} , and for f ∈
L p ( X ) we let ∥ f ∥ p := (cid:0)(cid:82) X | f ( x ) | p d x (cid:1) 1 /p .
We use ∥ f ∥ ∞ to denote the essential supremum of | f | .
For f, g : R n → R , we define convolution ( f ∗ g )( x ) :=
(cid:82) f ( y ) g
( x − y )
d y provided this integral exists.
We use C k ( X ) to denote continuously k -differentiable functions and C ∞ ( X ) to indicate the smooth functions.
The space of continuous k -differentiable functions with compact support on X is C kc
( X ) .
We denote by P ( X ) the set of probability measures, and P 2 ( X ) the set of probability measures with bounded second moments.
We denote by M sign ( X ) the set of finite signed measures.
Given a Lebesgue measurable map T : X → X and µ ∈ P 2 ( X ) , T # µ is the pushforward measure of µ by T .
Setup.
The domain X ⊂ R n where we constrain samples is assumed to be closed and fulldimensional, i.e., L n ( X ) > 0 ; the unconstrained case corresponds to X = R n .
The target density, always denoted as p , is assumed to satisfy p ∈ C 1 ( X ) , p ( x ) > 0 for all x ∈ X , and 0 < (cid:82) X p ( x ) d x < ∞ .
Let µ ∗ be the target probability measure µ ∗ ( B ) :
= (cid:82) B p ( x )d x / (cid:82) X p ( x )d x for
Borel B ⊂ X .
Our goal is to sample from µ ∗ .
3. M OLLIFIERS
Our framework relies on the notion of mollifiers originally introduced in Friedrichs (1944).
Definition 3.1 (family of mollifiers) .
We say { ϕ ϵ } ϵ> 0 ⊂ C 1 ( R n ) is a family of mollifiers if it satisfies
(a) For any ϵ > 0 , x ∈ R n , ϕ ϵ ( x ) ≥ 0 and ϕ ϵ ( x ) =
ϕ ϵ ( − x ) .
(b) For any ϵ > 0 , ∥ ϕ ϵ ∥ 1 = 1 and sup x ∈ R n ϕ ϵ ( x ) < ∞ .
(c) For any δ > 0 , p ∈ { 1 , ∞} , lim ϵ
→ 0 (cid:13)(cid:13) 1 R n \ B δ (0) ϕ
ϵ (cid:13)(cid:13) p = 0 .
As ϵ → 0 , the distribution with density ϕ ϵ converges to δ 0 , and if f ∈ L p ( R n ) and is continuous, then convolution ϕ
ϵ ∗ f converges to f both pointwise (Proposition A.4)
and in L p (Proposition A.7).
Our definition is different from the one used in PDE theory (H¨ormander, 2015), where some ϕ ∈
C ∞ c ( R n ) is fixed and the family of mollifiers is generated by ϕ ϵ ( x ) :
= ϵ − n ϕ ( x/ϵ ) .
While in
PDE theory mollifiers are typically used to improve the regularity of non-smooth functions, in our framework they are used to construct interaction energies that approximate the χ 2 divergence.
We will use the following mollifiers.
In the following definitions, we include normalizing constants
Z ϵ that ensure ∥ ϕ ϵ ∥ = 1 .
We do not write out Z ϵ explicitly as they might not admit clean forms and they are only relevant in theory but not in our practical algorithm.
For s > n , the s -Riesz family of mollifiers is defined as ϕ sϵ ( x ) :
= ( ∥ x ∥ 22 + ϵ 2 ) − s/ 2 /Z sϵ .
The Gaussian family of mollifiers is defined as ϕ gϵ ( x ) :=
exp (cid:0) − ∥ x ∥ 22 / 2 ϵ 2 (cid:1) /Z
gϵ .
The Laplace family of mollifiers is defined as ϕ lϵ ( x ) :=
exp ( − ∥ x ∥ 2 / ϵ ) /Z
lϵ .
Since { ϕ gϵ } and { ϕ lϵ } correspond to the densities of centered Gaussian and Laplace random variables, they satisfy Definition 3.1.
Proposition A.1 proves that { ϕ s ϵ } also satisfies Definition 3.1.
3. M OLLIFIED I NTERACTION E NERGIES
Definition 3.2 (mollified interaction energy) .
Given a family of mollifiers { ϕ ϵ } ϵ> 0 satisfying Defi- nition 3.1, for each ϵ > 0 , define a symmetric kernel W ϵ : R n
[Equation]
n
[Equation]
[Equation]
] by
[Equation]
Define the mollified interaction energy (MIE), E ϵ : P ( R n ) → [0 , ∞ ] , to be
[Equation]
Intuitively, minimizing E ϵ ( µ ) balances two forces: a repulsive force from ϕ ϵ ( x − y ) makes sure µ has a good spread, and an attractive force from ( p ( x ) p ( y )) − 1 / 2 helps µ concentrate on highdensity regions.
The exponent − 1 / 2 is chosen to balance the two forces so that, as we will show in
Theorem 3.3, E ϵ ( µ ) approximates the χ 2 divergence with respect to µ ∗ for small ϵ .
3. C ONVERGENCE TO χ 2 - DIVERGENCE
Recall that the χ 2 -divergence between probability measures P, Q is defined by
[Equation]
otherwise, (3) where Q ≪ P denotes absolute continuity of Q with respect to P with density d Q / d P .
Theorem 3.3.
Suppose µ ∈ P ( R n ) satisfies χ 2 ( µ ∥ µ ∗ ) < ∞ .
Then, E ϵ ( µ ) < ∞ for any ϵ > 0 .
Furthermore,
[Equation]
We provide a succinct proof of Theorem 3.3 in Appendix A.2 using the theory of mollifiers developed in Appendix A.1.
In Remark A.10 we discuss an extension of Theorem 3.3 to cases when X is not full-dimensional; in particular Theorem 3.3 still holds when X is “flat”, i.e., has Hausdorff dimension d < n and is contained in a d -dimensional linear subspace.
3. C ONVEXITY AND Γ - CONVERGENCE
We next study properties of the minimizers of MIEs: Does min µ ∈P ( X ) E ϵ ( µ ) admit a unique minimum?
If so, do minima of {E ϵ } ϵ> 0 converge to µ ∗ as ϵ → 0 ?
In order to answer affirmatively to these questions, we will need the associated kernel k ϵ ( x, y ) :
= ϕ ϵ ( x − y ) to satisfy the following property.
Definition 3.4 (i.s.p.d. kernel) .
A symmetric lower semicontinuous (l.s.c.) kernel K on X × X is integrally strictly positive definite (i.s.p.d.)
if for every finite signed Borel measure ν on X , the energy E K ( ν ) :
= (cid:82)(cid:82) K ( x, y ) d( ν × ν )( x, y ) is well-defined (i.e., the integrals over the negative and positive parts of ν are not both infinite), and E K ( ν ) ≥ 0 where the equality holds only if ν = 0 on all Borel sets of X .
For the mollifiers we consider, the associated kernel k ϵ ( x, y ) =
ϕ ϵ ( x − y ) is i.s.p.d.
on any compact set (Pronzato & Zhigljavsky (2021, Example 1.2)).
Proposition 3.5.
Suppose X is compact and k ϵ ( x, y ) :
= ϕ ϵ ( x − y ) is i.s.p.d.
on X .
Then,
(a) The kernel W ϵ ( x, y ) defined in (1) also i.s.p.d on X .
(b) The functional E ϵ is strictly convex on M sign ( X ) and attains a unique minimum on P ( X ) .
We next show the convergence of minima of {E ϵ } to µ ∗ as a consequence of Γ -convergence of the sequence {E ϵ } .
Theorem 3.6.
Suppose k ϵ ( x, y ) :
= ϕ ϵ ( x − y ) is i.s.p.d.
on compact sets for every ϵ > 0 .
Then we have Γ -convergence
(Definition A.13) E ϵ Γ → χ 2 ( · ∥ µ ∗ )+1 as ϵ → 0 .
In particular, if X is compact, if we denote µ ∗ ϵ := arg min µ ∈P ( X ) E ϵ ( µ ) , then µ ∗ ϵ → µ ∗ weakly and lim ϵ
→ 0 E ϵ ( µ ∗ ϵ ) = 1 .
We prove Theorem 3.6 in Appendix A.3 using Fourier transforms and Bochner’s theorem.
Theorem 3.6 provides basis for minimizing E ϵ for a small ϵ , since its unique minimum will be a good approximation of µ ∗ .
3. D IFFERENTIAL CALCULUS OF E ϵ IN P 2 ( R n )
We next study the gradient flow of E ϵ in Wasserstein space P 2 ( R n ) for X = R n .
Understanding the gradient flow of a functional often provides insights into the convergence of algorithms that simulates gradient flow with time discretization (Ambrosio et al., 2005, Chapter 11) or spatial discretization (Chizat, 2022).
Proofs of this section are given in Appendix A.4.
Let F : P 2 ( R n ) → R be a functional.
The Wasserstein gradient flow of F (Ambrosio et al., 2005,
Definition 11.1.1) is defined as the solution { µ t } t ≥ 0 of the PDE: ∂µ t ∂t = ∇ · ( µ t w F ,µ t ) where w F ,µ ∈ L 2 ( µ ; R n ) is a Frech´et subdifferential of F at µ .
Intuitively, gradient flows capture the evolution of the variable being optimized if we were to do gradient descent 1 on F when the step sizes go to 0 .
We next show that the gradient flow of E ϵ agrees with that of χ 2 ( · ∥ µ ∗ ) in the sense that their subdifferentials coincide as ϵ → 0 .
Proposition 3.7.
Assume µ ∈ P 2 ( R n ) has density q ∈ C 1 c ( R n ) 2 .
Then any strong Frech ´ et subdifferential (Definition A.20) of E ϵ at µ takes the form
[Equation]
Moreover, if for sufficiently small ϵ , ϕ ϵ has compact support independent of ϵ , then
[Equation]
where w χ 2 ,µ is a strong Frech ´ et subdifferential of χ 2 ( · ∥ µ ∗ ) .
While simulating χ 2 divergence gradient flow is often intractable (Trillos & Sanz-Alonso, 2020,
Section 3.3), our MIE admits a straightforward practical algorithm (Section 4).
We next show that E ϵ is displacement convex at µ ∗ ∈ P 2 ( R n ) as ϵ → 0 , obtaining a result similar to Korba et al.
(2021, Corollary 4).
This hints that gradient flow initialized near µ ∗ will have fast convergence.
Proposition 3.8.
Suppose p ∈ C 2 c ( R n ) .
Suppose that for sufficiently small ϵ , ϕ ϵ has compact support independent of ϵ .
Assume k ϵ
( x, y ) :
= ϕ ϵ ( x − y ) is i.s.p.d.
Let ξ ∈ C ∞ c ( R n ; R n ) and µ t :
= ( I + t ξ ) # µ ∗ .
Then lim ϵ
[-→ 0 d 2 d t 2 (cid:12)(cid:12)(cid:12)(cid:12) t =
 E ϵ ( µ t ) ≥ 0 .-][Equation]
A PRACTICAL SAMPLING ALGORITHM
We now present a first-order particle-based algorithm for constrained and unconstrained sampling by minimizing a discrete version of (2).
Substituting an empirical distribution for µ in (2) gives the discrete mollified interaction energy, for N particles ω N = { x 1 , . .
, x N } ,
[Equation]
Denote ω ∗ N,ϵ ∈ arg
min ω N ⊂ X E ϵ
( ω N ) and µ ∗ ϵ = arg min µ ∈P 2 ( X ) E ϵ ( µ ) ,
If X is compact, by Borodachov et al.
(2019, Corollary 4.2.9), we have weak convergence δ ω ∗ N,ϵ
→ µ ∗ ϵ as N → ∞ .
If in addition k ϵ ( x, y ) :
= ϕ ϵ ( x − y ) is i.s.p.d., then by Theorem 3.6, we have weak convergence µ ∗ ϵ → µ ∗ .
This shows that minimizing (6) with a large N and a small ϵ will result in an empirical distribution of particles that approximates µ ∗ .
Our sampling method, mollified interaction energy descent (MIED), is simply running gradient descent on (6) over particles (Algorithm 1) with an update that resembles the one in SVGD (see the discussion in Appendix C.1).
Below we address a few practical concerns.
Optimization in the logarithmic domain.
In practical applications, we only have access to the unnormalized target density p .
The normalizing constant of the mollifier ϕ
ϵ can also be hard to compute (e.g., for the Riesz family).
While these normalization constants do not affect the minima of (6), they can still affect gradient step sizes during optimization.
Moreover, ϕ ϵ can be very large when ϵ is small, and in many Bayesian applications p can be tiny and only log p is numerically significant.
To address these issues, we optimize the logarithm of (6) using the log-sum-exp trick (Blanchard et al., 2021) to improve numerical stability and to get rid of the arbitrary scaling of the normalizing constants:
[Equation]
Special treatment of the diagonal terms.
Since ϕ ϵ goes to the Dirac delta as ϵ → 0 , the discretization of (2) on a neighborhood of the diagonal { ( x, y ) : x = y } needs to be handled with extra care.
In (6) the diagonal appears as (cid:80) N i =1 ϕ ϵ (0) p ( x i ) − 1 which can dominate the summation when ϵ is small and
then ϕ ϵ (0) becomes too large.
For the mollifiers that we consider, we use a different diagonal term (cid:80) Ni =1 ϕ ϵ
( h i /κ n ) p ( x i ) − 1 where h
i :=
min j ̸ = i ∥ x i − x j ∥ and κ n ≥ 1 is a constant depending only on the dimension n .
Since≤ ϕ ϵ
( h i /κ n ) ≤ ϕ ϵ (0) , the energy obtained will be bounded between (6) and the version of (6) without the diagonal terms.
Hence by the proof of Theorem 4.2.2 of Borodachov et al.
(2019), the discrete minimizers of E ϵ still converge to µ ∗ ϵ as
N → ∞ and ϵ → 0 .
Empirically we found the choice κ n = (1 . 3 n ) 1 /n works well for the Riesz family of mollifiers and we use the Riesz family primarily for our experiments.
Constraint handling.
If X ̸ = R n , we need to minimize (7) subject to constraints ω N = { x 1 , . .
, x N } ⊂ X .
Since the constraint is the same for each particle x
i , we want our algorithm to remain parallelizable across particles.
We consider two types of constraints: (a) there exists a differentiable map f : R n → X , with L n ( X \ f
R n ))
= 0 ; (b) the set X is given by { x ∈ R n : g ( x ) ≤ 0 } for a differentiable g : R n → R .
For (a), we reduce the problem to unconstrained optimization in R n using objective log E ϵ ( f ( ω N )) with ω N = { x 1 , . . .
, x N } ⊂ R n and f ( ω N ) := { f ( x 1 ) , . . .
, f ( x N ) } .
For (b), we apply the dynamic barrier method by Gong & Liu (2021) to particles in parallel; see Appendix C. for details and an extension to handling multiple constraints.
E XPERIMENTS
We compare MIED with recent alternatives on unconstrained and constrained sampling problems.
Unless mentioned otherwise, we choose the s -Riesz family of mollifiers { ϕ sϵ } with s = n +10 − 4 and ϵ = 10 − 8 : we found minimizing the MIE with such mollifiers results in well-separated particles so that we can take ϵ to be very small as our theory recommends.
This is not the case for the Gaussian or the Laplace family as setting ϵ too small can cause numerical issues even when particles are wellseparated.
In Appendix D.5, we compare different choices of s on a constrained mixture distribution.
Unless mentioned otherwise: for SVGD (Liu & Wang, 2016), we use the adaptive Gaussian kernel as in the original implementation (adaptive kernels can be prone to collapse samples—see Appendix D.2); for KSDD (Korba et al., 2021), we use a fixed Gaussian kernel with unit variance.
All methods by default use a learning rate of 0 .
01 with Adam optimizer (Kingma & Ba, 2014).
The source code can be found at https://github.com/lingxiaoli94/MIED .
5. U NCONSTRAINED SAMPLING
Gaussians in varying dimensions.
We first compare MIED with SVGD and KSDD on Gaussians of varying dimensions and with different numbers of particles.
In Figure 1, we see that as the number of dimensions increases, MIED results in best samples in terms of W 2 distance (Wasserstein-2 distance computed using linear programming) with respect to 10 4 i.i.d.
samples, while SVGD yields lower energy distance (Sz ´ ekely & Rizzo, 2013).
We think this is because MIED results in more evenly spaced samples (Figure 5) so the W 2 distance is lower, but it is biased since ϵ >
so for energy distance SVGD performs better.
More details can be found in Appendix D.1.
[Figure]
Product of two Student’s t -distributions.
Next, we consider a 2-dimensional distribution constructed as the product of two independent t -distributions with the degree of freedom ν = 2 composed with a linear transform.
Unlike Gaussians, Student’s t -distributions have heavy tails.
On the left of Figure
[Figure] we visualize the results of each method with 1000 samples after 2000 iterations.
MIED captures the heavy tail while SVGD fails.
Quantitatively, on the right of Figure [Figure] while SVGD captures the distribution in [ − a, a ] 2 better for a ≤ 3 , MIED yields lower metrics for a bigger a
[Figure]
KSDD − 5 − 5 . 0 − 2 .0 .
MIED    10 Edgelength a {+10 − 2 10 −
+}
W 2 D i s t a n ce ( l og s c a l e )
SVGDKSDDMIED    10 Edgelength a 10 −10 −10 − E n e r g y D i s t a n ce ( l og s c a l e ) SVGDKSDDMIED
Bayesian logistic regression.
We compare MIED with SVGD and KSDD for the Bayesian logistic regression setting in Liu & Wang (2016).
We further include a na¨ıve baseline, independent particle descent (IPD), which simply runs gradient descent independently on each particle to maximize the posterior probability.
In addition to using test accuracy as the metric, we include W 2 distance and energy distance between the samples from each method and 10 4 samples from NUTS (Hoffman et al., 2014) after sufficient burn-in steps.
We summarize the results in Table 5.1.
All methods, including the na¨ıve baseline IPD, are comparable in terms of test accuracy.
In other words, accuracy is not a good metric for comparing the quality of posterior samples.
Bayesian inference is typically preferred over maximum a posteriori estimation for its ability to capture uncertainty.
When evaluating the quality of the uncertainty of the samples using distances between distributions, MIED provides the best approximation in terms of the W 2 distance and all methods (except IPD) are comparable in terms of the energy distance.
[Table]
5. C ONSTRAINED SAMPLING
Uniform sampling in 2D. We consider uniform sampling in the square [ − 1 , 1] 2 with 500 particles.
We reparameterize our particles using tanh to eliminate the constraint and show results with various choices of mollifiers—we always choose the smallest ϵ while the optimization remains stable.
We compare our method with mirror LMC
(Zhang et al., 2020) and SVMD/MSVGD by Shi et al.
(2021) using entropic mirror map ϕ ( θ ) = (cid:80) ni =1 ((1 + θ i [-+ log 1 
-]{+) log(
+}
+ θ i ) + (1 − θ i ) log(1 − θ i )) .
To demonstrate that SVGD and KSDD break down in constrained domains, we implement these two methods adapted to the constrained setting using the same reparameterization as our method.
The initial particles are drawn uniformly from [ − 0 .
The left plot of Figure 3 shows that quantita-
[Figure]
[-MIED(Riesz) − 1 − 1 . 0 − 0 .0 . 
-]
tively our method achieves much lower energy distance and W 2 distance (compared to a 5000 i.i.d.
samples uniformly drawn in [ − 1 , 1] 2 .).
On the right of Figure 3, we visualize the samples from each method.
We see that samples from SVGD with an adaptive kernel collapse—we investigate this issue in Appendix D.2.
Samples from SVGD with a fixed kernel size and KSDD are non-uniform (we choose kernel sizes that produce the best result).
While SVMD creates locally clustered artifacts, MSVGD produces good results.
For mirror LMC, the resulting samples are not evenly spaced, resulting in worse W 2 distances.
When using tanh as the reparameterization map, MIED produces decent results for all three families of mollifiers (Riesz, Gaussian, Laplace).
When using the same entropic mirror map, MIED also produces good results.
This highlights the flexibility of our method with constraint handling, whereas, for LMC and SVMD/MSVGD, the mirror map has to be chosen carefully to be the gradient of a strongly convex function while capturing the constraints.
In Appendix D.3, we further test SVMD/MSVGD with a different choice of the mirror map where they break down.
In Appendix D.4, we conduct a similar comparison for sampling from a 20dimensional Dirichlet distribution using the same setup as Shi et al.
In scenarios where a good choice of a mirror map is available, SVMD/MSVGD can obtain better performance compared to MIED.
In Appendix D.6, we conduct additional qualitative experiments for MIED, demonstrating its effectiveness for challenging constraints and multi-modal distributions.
Fairness Bayesian neural networks.
We train fair Bayesian neural networks to predict whether the annual income of a person is at least $50 , 000 with gender as the protected attribute using the
Adult Income dataset (Kohavi et al., 1996).
We follow the same setup as in Liu et al.
(2021) where the dataset D = { x ( i ) , y ( i ) , z ( i ) } |D| i =1 consists of feature vectors x ( i ) , labels y
( i ) (whether the income is ≥ $50 , 000 ), and genders z ( i ) (protected attribute).
The target density is taken to be the posterior of logistic regression with a two-layer Bayesian neural network ˆ y ( · ; θ ) with weights θ , and we put a standard Gaussian prior on each entry of θ independently.
Given t > 0 , the fairness constraint is g ( θ ) =
( cov ( x,y,z ) ∼D
[ z, ˆ y ( x ; θ )]) 2 − t ≤ 0 .
On the left of Figure 4, we plot the trade-off curve of the result obtained using our method and the methods from Liu et al.
(2021) for t ∈ { 10 − 5 , 10 − 4 , 0 . 0001 , 0 .
Details can be found in Appendix D.7.
Our method recovers a much larger Pareto front compared to the alternatives.
On the right of Figure 4, we visualize the curves of the energy and the covariance versus the number of training iterations: as expected we see a smaller t results in bigger MIE (lower log-likelihood) and smaller covariance between the prediction and the protected attribute.
[Figure]
C ONCLUSION
We present a new sampling method by minimizing MIEs discretized as particles for unconstrained and constrained sampling.
This is motivated by the insight that MIEs converge to χ 2 divergence with respect to the target measure as the mollifier parameter goes to 0.
The proposed method achieves promising results on the sampling problems we consider.
Below we highlight three limitations.
First, as discussed in Remark A.10, our theory only applies when the domain is full-dimensional or flat.
Extending our theory to handle cases where the domain is an arbitrary d -rectifiable set is an important next step as it allows the handling of more complicated constraints such as nonlinear equality constraints.
Secondly, when ϵ > 0 , the minimizer of MIE can be different from the target measure.
Finding a way to debias MIE (e.g., like how Sinkhorn distances are debiased (Feydy et al., 2019)) is an interesting direction.
Lastly, the connection between the minimizers of the discretized MIE (6) and those of the continuous MIE (2) is only established in the limit as N → ∞ .
We hope to investigate how well the empirical distribution of particles minimizing (6) approximates the target measure when N is finite as in Xu et al.
Reproducibility statement.
We provide self-contained proofs in Appendix A for the theoretical results stated in the main text.
The source code for all experiments can be found at https:// github.com/lingxiaoli94/MIED . Acknowledgements Qiang Liu would like to acknowledge the support of NSF CAREERand ONR. The MIT Geometric Data Processing group acknowledges the generous support of Army Research Office grants W911NF2010168 and W911NF2110293, of Air Force Office of Scientific
Research award FA9550-19-1-031, of National Science Foundation grants IIS-1838071 and CHS1955697, from the CSAIL Systems that Learn program, from the MIT–IBM Watson AI Laboratory, from the Toyota–CSAIL Joint Research Center, from a gift from Adobe Systems, and from a Google
Research Scholar award.
A D ETAILED A NALYSIS A. P RELIMINARIES ON MOLLIFIERS
Proposition A.1.
For s > n , the s -Riesz family of mollifiers, defined as ϕ sϵ ( x ) :
=
( ∥ x ∥ 22 + ϵ 2 ) − s/ 2 /Z sϵ , satisfies Definition 3.1.
In order to prove Proposition A. we first prove a lemma
Lemma A.2.
Let b ≥ 0 .
Then for any ϵ > 0 ,
[Equation]
[Equation]
where H n − 1 ( S n − 1 ) is the volume of the ( n − 1)
-dimensional sphere.
Furthermore, assuming s > b + n , then for any ϵ > 0 , δ > 0 ,
[Equation]
Proof.
For (8), we compute
[Equation]
where we use substitution r = ϵt on the second line.
If s > b + n , then
[Equation]
where the last integral uses s > b + n .
Proof of Proposition A.1.
It is clear that (a) of Proposition A.1 holds for ϕ sϵ ( x ) .
Since ϕ sϵ ( x ) is bounded we have ϕ sϵ ∈ L ∞ ( R n ) .
For any δ > 0 , we have, for ϵ ≤ δ ,
[Equation]
where we use (8) with b = 0 and C is a constant depending only on n, s .
On the other hand, using (9) with b = 0 ,
[Equation]
(cid:90)
[Equation]
where C ′ depends only on n, s .
With δ =
ϵ , we see that Z sϵ ϕ sϵ ∈ L 1 ( R n ) so (b) is satisfied.
Since δ is fixed and s > n , we see that
[Equation]
Since (cid:82) B δ (0)
ϕ
sϵ ( x ) d x + (cid:82) R n \ B δ (0)
ϕ sϵ ( x ) d x = 1 , we have shown (c).
In the rest of this section, we assume { ϕ ϵ } ϵ> 0 is a family of mollifiers satisfying Definition 3.1.
Lemma A.3.
For any p ∈ [1 , ∞ ] , for any δ > 0 , ϕ ϵ ∈
L p ( R n ) and lim ϵ
→ 0 (cid:13)(cid:13) 1 R n \ B δ (0) ϕ
ϵ (cid:13)(cid:13) p = 0 .
holds.
Proof.
Assume p / ∈ { 1 , ∞} since both cases are covered in the assumptions.
Then for any δ > 0 , by H¨older’s inequality,
[Equation]
Similarly,
[Equation]
Letting ϵ
→
0 and applying (c) gives lim ϵ → 0 (cid:13)(cid:13) 1 R n \ B δ (0)
ϕ
ϵ (cid:13)(cid:13) p = 0 .
Proposition A.4.
Let f ∈
L p ( R n ) , p ∈ [1 , ∞ ] .
Then for every ϵ > 0 , the integral
[Equation]
exists, so that ( f ∗ ϕ ϵ )( x ) is finite.
Moreover, if f is continous at x , then
[Equation]
[Equation]
Proof.
Observe that by H¨older’s inequality, for q such that 1 /p + 1 /q = 1 (allowing infinity),
[Equation]
Hence ( f ∗ ϕ ϵ
)( x ) is integrable and finite.
For any ϵ > 0 , note that
[Equation]
Since (cid:82) ϕ
ϵ ( x ) d x = 1 , we have
[Equation]
Fix t > 0 .
Continuity of f at x implies there exists δ > 0 such that | f ( x − y )
− f ( x ) | < t for all y ∈ B δ (0) .
Then
[Equation]
On the other hand, by H¨older’s inequality, (cid:90)
[Equation]
Hence
[Equation]
By Lemma A.3, since ∥ f ∥ p < ∞ , taking ϵ → 0 we get, for any t > 0 ,
[Equation]
Now let t → 0 we get (10).
Corollary A.5.
Let f ∈ L p ( R n ) , p ∈ [1 , ∞ ] , be a continuous function.
If { f ϵ } is a sequence of functions that converge uniformly to f , then for every x ,
[Equation]
Proof.
Note that
[Equation]
Proposition A.4 shows the second term goes to 0 as ϵ → 0 .
For the first term, we have
[Equation]
by uniform convergence.
Lemma A.6.
For f ∈ L p ( R n ) , p ∈ [1 , ∞ ) , we have
[Equation]
where we use τ a f to denote the translated function
τ a f ( x ) :
= f
( x − a ) .
Proof.
Fix ϵ > 0 .
It is a standard fact that C c ( R n ) is dense in L p ( R n ) .
Hence there exists g ∈ C c ( R n ) such that ∥ f − g
∥
p < ϵ .
Since g is continuous with compact support, it is uniform continuous.
Then there exists δ > 0 with | g ( x − y )
− g ( x ) | < ϵ 1 /p / L n ( K ) if y ∈ B δ ( x ) .
Hence for such y we have ∥ τ y g − g ∥ pp <
ϵ with L n ( K ) < ∞ .
Thus
[Equation]
∥ p ≤ 3 ϵ.
Proposition A.7.
Let f ∈ L p ( R n ) , p ∈ (1 , ∞ ) .
Then for every ϵ > 0 ,
[Equation]
In particular, f ∗ ϕ
ϵ ∈
L p ( R n ) .
Moreover,
[Equation]
Proof.
For every ϵ > 0 , we have
[Equation]
where we use Jensen’s inequality with the observation that ϕ
ϵ ( y )
d y is a probability measure and
Tonelli’s theorem to exchange the order of integration.
To show the first claim, note that,
[Equation]
Hence ∥ f
[Equation]
ϕ
[Equation]
For the second claim, by Lemma A.6, the function y (cid:55)→ ∥
τ y f − f ∥ pp is continuous at y = 0 with limit 0 .
Hence by Proposition A.4, we are done by taking ϵ → 0 .
A. C ONVERGENCE TO χ 2 - DIVERGENCE
We start by giving an alternative formula for χ 2 divergence with respect to the target measure.
Lemma A.8.
Define a functional E : P ( R n ) → [0 , ∞ ] by
[Equation]
if
[Equation]
Then for every µ ∈ P ( R n ) ,
[Equation]
Proof.
If µ is not absolutely continuous with respect to µ ∗ , then both sides are infinity.
Otherwise, µ ( X ) =
and µ has some density q .
We then compute
[Equation]
Proof of Theorem 3.3.
Since χ 2 ( µ ∥ µ ∗ ) < ∞ , we know µ ≪ µ ∗ , so µ ( X ) = 1 .
We let q ( x ) be the density of µ which satisfies q ( x ) = 0 for x ∈ R n \ X .
We compute using Tonelli’s theorem (since our integrand is positive):
[Equation]
where we define g ( x ) :
= q ( x ) / (cid:112) p ( x ) on X and g ( x ) := 0 for x / ∈ X .
Moreover, by Lemma A.8,
[Equation]
The assumption χ 2 ( µ ∥ µ ∗ ) < ∞ then implies g ∈ L 2 ( R n ) .
By Proposition A.7, we have ϕ ϵ ∗ g ∈ L( R n ) .
Next notice that
[Equation]
This shows the first claim.
Finally, the last expression goes to 0 as ϵ → 0 since ∥ ϕ ϵ ∗ g − g ∥ 2 → 0 by Proposition A.7.
Remark A.9.
One may ask if similar results as Theorem 3.3 for f -divergences other than χ 2 divergence can be obtained using similar techniques.
We think the proof of Theorem 3.3 is highly specific to χ 2 divergence because in (13) there are two copies of q coming from the definition of E ϵ as a double integral over µ .
An f -divergence between µ and µ ∗ has the form (cid:82) X f ( q ( x ) /p
( x ))
d p ( x ) , and the only way to have q 2 showing up is to choose f ( x ) =
x 2 corresponding to χ 2 divergence.
Remark A.10.
It is possible to prove versions of Theorem 3.3 when X has Hausdorff dimension d < n : in such cases χ 2 -divergence still makes sense as does (2) .
When X is “flat”, i.e., with Hausdorff dimension d and contained in a d -dimensional linear subspace of R n , e.g., when X is defined by a set of linear inequalities, then Theorem 3.3 follows if we adapt the assumptions of Definition 3.
and calculation in the proof of Theorem 3.3 to be in the subspace.
For a general d -dimensional X , a similar calculation yields E ϵ ( µ ) =
(cid:82) X (cid:16)(cid:82)
X ϕ ϵ
( x − y )
(cid:16) q √ p (cid:17) ( y )
d H d ( y ) (cid:17) (cid:16) q √ p (cid:17) ( x ) d H d ( x ) .
For a similar argument to go through, we will need the normalizing constant (cid:82) X ϕ ϵ ( x − y )
d H d ( y ) to the same for all x .
This is true when X is a d -dimensional sphere, but for a general X the last integral will depend on the base point x .
Proving a version of Theorem 3.3 for X with Hausdorff dimension d < n is an interesting future direction.
A. C ONVEXITY AND Γ - CONVERGENCE
We start by recalling a few definitions regarding functionals in P ( R n ) .
Definition A.11.
We say a functional F : P ( R n ) → ( −∞ , ∞ ] is proper if there exists µ ∈ P ( R n ) such that
F ( µ ) < ∞ , and is lower semicontinuous (l.s.c.)
if for any weakly convergence sequence µ k → µ , we have lim inf k →∞ F ( µ k ) ≥ F ( µ ) .
Lemma A.12.
For any ϵ > 0 , the functional E ϵ : P ( R n ) → ( −∞ , ∞ ] is proper and l.s.c.
Moreover, if X is compact, the minimum min µ ∈P ( X ) E ϵ ( µ ) is attained by some measure in P ( X ) .
Proof.
Taking any x ∈ X , since ϕ ϵ is bounded and p ( x ) > 0 , we see that E ϵ ( δ x ) < ∞ so E ϵ is proper.
Moreover, given weakly convergence µ k → µ , by Portmanteau theorem and the fact that W ϵ is nonnegative and l.s.c.
, we conclude that E ϵ is also l.s.c.
The set of probability distributions P ( X ) ⊂ P ( R n ) is tight by the compactness of X .
It is closed since if { µ k } ⊂ P ( X ) weakly converges to µ , then by Portmanteau theorem, µ ( X ) ≥ lim sup µ ( X ) = {+
+} so that µ ∈
[Equation] Hence by Prokhorov’s theorem (Theorem 5.1.3 (Ambrosio et al., 2005)), P ( X ) is (sequentially) compact.
It is then an elementary result that any l.s.c.
function attains its minimum on a compact set.
We next prove Proposition 3.5 regarding the convexity of E ϵ and the uniqueness of its minimum.
Proof of Proposition 3.5.
Let µ be any finite signed Borel measure.
The compactness assumption and the fact that p ∈ C 1 ( X ) imply p ( x ) > δ for any x ∈ X for some δ > 0 .
Hence p ( x ) − 1 / 2 ≤ δ − 1 / 2 , so the weighted measure ˜ µ defined by d˜ µ ( x ) :
= p − 1 / 2 ( x ) d µ ( x ) is also finite.
By the definition of i.s.p.d.
kernels, we have E W ϵ ( µ ) =
E k ϵ (˜ µ ) >
if ˜ µ is not the zero measure, which is equivalent to µ not being zero since p > 0 .
Thus W ϵ is i.s.p.d.
on X .
Also note that ( p ( x ) p ( y )) − 1 / 2 < δ − 1 for all x, y ∈ X , so E W ϵ is always finite.
By Lemma 1.1 of Pronzato & Zhigljavsky (2021), we conclude that E W ϵ is strictly convex in M sign , the space of finite signed measures, and in particular it is convex on
P ( X ) .
Hence combined with the existence result from
Lemma A.12 we conclude E ϵ attains a unique minimum in P ( X ) .
Our goal in the rest of this section is to prove the Γ -convergence of E ϵ to E : P ( R n ) → [0 , ∞ ] defined in (12).
By Lemma A.8, we have
E ( µ ) =
χ 2 ( µ ∥ µ ∗ ) + 1 .
Hence together we will have proved Theorem 3.6.
Definition A.13 ( Γ -convergence) .
A sequence of functionals F ϵ : P ( R n ) → ( −∞ , ∞ ] is said to Γ -converge to F : P ( R n ) → ( −∞ , ∞ ] , denoted as F ϵ Γ → F , if:
(a) For any sequence µ ϵ ∈ P ( R n ) converging weakly to µ ∈ P ( R n ) , lim inf ϵ → 0 F ϵ ( µ ϵ ) ≥ F ( µ ) ; (b) For any µ ∈ P ( R n ) , there exists a sequence µ ϵ ∈ P ( R n ) converging weakly to µ with lim sup ϵ
→ 0 F ϵ ( µ ϵ ) ≤ F ( µ ) .
We will show E ϵ Γ → E using Fourier transforms and Bochner’s theorem.
Definition A.14 (Fourier transform) .
For f ∈ L 1 ( R n ) , its Fourier transform ˆ f is the complexvalued function defined via
[Equation]
More generally, for a signed finite measure µ ∈ M sign ( R n ) , its Fourier transform ˆ µ is the complexvalued function defined via
[Equation]
This integral is always well-defined and moreover ˆ µ is uniformly continuous; see Borodachov et al.
(2019, Section 1.10).
We will prove the following weak version (under the additional assumption that a mollifier ϕ is integrable) of Bochner’s theorem suitable for our case.
In particular we will need the Fourier inversion formula which is not given in the usual statement of Bochner’s theorem.
On the other hand, we cannot directly use the Fourier inversion formula since it is not obvious how to check the integrability of ˆ ϕ when ϕ is a mollifier.
Lemma A.15.
Suppose ϕ ∈ L 1 ( R n ) is even, bounded, continuous, and k ( x, y ) :
= ϕ
( x − y ) is i.s.p.d.
on any compact sets.
Then its Fourier transform ˆ ϕ is real, nonnegative, and the following inversion formula holds:
[Equation]
Proof.
The proof is adapted from that of Varadhan (2001, Theorem 2.7) and is extended to the multi-dimensional case.
Since ˆ ϕ ( ξ ) := (cid:82) e − 2 πix · ξ ϕ ( x ) d x and ϕ ( x ) = ϕ ( − x ) , with a change of variable x ′ =
− x , we obtain ˆ ϕ ( ξ )
=
ˆ ϕ ( ξ ) , so ˆ ϕ is real.
Next we show that ˆ ϕ is nonnegative.
For T > 0 , we compute, for a fixed ξ ∈ R n ,
[Equation]
where we have used change of variable u = t − s , v = t + s .
Since ϕ ∈ L 1 ( R n ) , by dominated convergence theorem, we have as T → ∞
[Equation]
For t, s ∈ R n , we have
[Equation]
For a fixed T , if we define a finite measure µ as d µ = 1 t ∈ [0 ,T ] n cos [Equation] 2 πt · ξ ) d t , then k being i.s.p.d. implies
[Equation]
and similarly for sin .
Since ˆ ϕ is real, we conclude that ˆ ϕ is nonnegative.
Finally we prove Equation (14).
For σ > 0 , define ˆ ϕ σ ( ξ ) :
=
ˆ ϕ ( ξ ) e − σ 2 ∥ ξ ∥ 22 .
Since ˆ ϕ is bounded (because ϕ ∈ L 1 ( R n ) ), we see that ˆ ϕ σ ∈ L 1 ( R n ) .
We compute, for x ∈ R n , using Fubini’s theorem,
[Equation]
where we use the Fourier transform formula (Borodachov et al., 2019, (4.4.1)) of the Gaussian distribution.
Notice that p σ ( y ) :=
( π/σ 2 ) n/ 2 e − π 2 ∥ x − y ∥
22 /σ 2 is the density of a multivariate Gaussian centered at x with covariance σ 2 / (2 π 2 ) · I .
Hence (cid:82) e 2 πix · ξ ˆ ϕ σ ( ξ ) d ξ = ( p σ ∗ ϕ )(0) .
Since ϕ is bounded by assumption, with x =
we find (cid:82) ˆ ϕ σ ( ξ ) d ξ ≤ ∥ ϕ ∥ ∞ .
Taking σ → 0 , by monotone convergence theorem, we have (cid:82) ˆ ϕ ( ξ ) d ξ ≤ ∥ ϕ ∥ ∞ , so together with the fact that ˆ ϕ ≥
we have ˆ ϕ ∈ L 1 ( R n ) .
Finally, since ˆ ϕ σ ≤ ˆ ϕ , by dominated convergence theorem and Proposition A.4 (note p σ is centered at x ), we have
[Equation]
Proposition A.16.
Given ϕ : R n → R satisfying the assumptions of Lemma A.15, for any ν ∈
M sign ( R n ) , it holds that
[Equation]
Proof.
By Lemma A.15,
[Equation]
where we use Fubini’s theorem (all measures are finite and e − 2 πi · is bounded) to exchange the order of integration.
Proof of Theorem 3.6.
First note that Definition A.13(b) is immediate from Theorem 3.3 with µ ϵ = µ : if E ( µ ) = ∞ , then trivially lim sup ϵ
→ 0 E ϵ ( µ ) ≤ E ( µ ) ; otherwise we apply Theorem 3.3.
So we focus on proving criterion Definition A.13(a).
Fix a sequence µ ϵ ∈ P ( R n ) that converges weakly to µ ∈ P ( R n ) , and our goal is to show lim inf ϵ → 0 E ϵ ( µ ϵ ) ≥ E ( µ ) .
Without loss of generality, we may assume E ϵ ( µ ϵ ) < ∞ for all ϵ > 0 (these terms have no effect in lim inf ), which implies µ ϵ ( X ) = 1 .
By Portmanteau’s theorem and the assumption that X is closed, we have µ ( X ) ≥ lim sup ϵ
→ 0 µ ϵ ( X ) = 1 .
So all of µ ϵ and µ will have support in X .
For m > 0 , ϵ > 0 , define a nonnegative measure ν ϵ,m by
[Equation]
where h m : R n → R is a continuous monotonically decreasing cutoff function satisfying h m ( x ) = 1 if ∥ x ∥ 22 < m and h m ( x ) =
if ∥ x ∥ 22 > m + 1 , and that h m ( x ) ≤ h m ′ ( x ) for m < m ′ .
Then since p > 0 is continuous, it is bounded below on any compact set in X , and hence ν ϵ,m is finite.
Also define, for m > 0 , a nonnegative measure ν m by
[Equation]
which is again finite following the same reasoning.
Then for any m > 0 , denoting d ν ϵ ( x ) :
= p − 1 / 2 ( x ) d µ ϵ ( x ) ,
[Equation]
where we apply Proposition A.16 for the last equality.
On the other hand, note that
[Equation]
Since µ ϵ → µ weakly and the last integrand is a continuous bounded function, we have
[Equation]
On the other hand, ˆ ϕ ϵ ( ξ ) =
(cid:82) e − 2 πiξ · x ϕ ϵ ( x ) d x .
Since by Definition 3.1, ϕ ϵ d x converges to δ 0 in probability (and in particular weakly), we have lim ϵ
→ 0 ˆ ϕ ϵ ( ξ ) = 1 .
Applying Fatou’s lemma to (15), we obtain, for any m > 0 ,
[Equation]
If (cid:82) | ˆ ν m ( ξ ) | 2 d ξ = ∞ for any m > 0 , then we are done since lim inf ϵ → 0 E ϵ ( µ ϵ ) = ∞ .
Otherwise, by K¨uhn (2016, Lemma 2.11), for every m > 0 , ν m has density in L 2 ( R n ) .
This imples µ has density everywhere.
Suppose d µ ( x ) =
q ( x ) d L n ( x ) .
By Plancherel’s theorem and the monotone convergence theorem, we have
[Equation]
This completes the proof of E ϵ Γ → E .
Now suppose X is compact, and let µ ∗ ϵ := arg min µ ∈P ( X ) E ϵ ( µ ) .
To establish µ ∗ ϵ → µ ∗ weakly (resp.
lim
ϵ
→ 0 E ϵ ( µ ∗ ϵ ) =
E ( µ ∗ ) = 1 ), it suffices to show that every subsequence of { µ ∗ ϵ } (resp.
{E ϵ ( µ ∗ ϵ ) } ) has a further convergence subsequence converging to µ ∗ (resp.
E ( µ ∗ ) ).
With a slight abuse of notation, we assume the sequence of ϵ is chosen so that { µ ∗ ϵ } (resp.
{E ϵ ( µ ∗ ϵ ) } ) is already some subsequence of the original sequence.
As argued in the proof of Lemma A.12,
P ( X ) is compact with respect to weak convergence.
Hence { µ ∗ ϵ } has a weakly convergence subsequence µ ∗ ϵ k
→ ν for some ν ∈ P ( X ) .
The Γ -convergence E ϵ Γ → E implies
[Equation]
where we have used, for each inequality, µ ∗ ϵ
k = arg min P ( X ) E ϵ k , Definition A.13(a), the first paragraph of this proof, and again the fact that µ ∗ ϵ k = arg min P ( X ) E ϵ k .
Since lim inf k →∞ E ϵ k ( µ ∗ ) =
E ( µ ∗ ) by Theorem 3.3, we have
[Equation]
where the last inequality follows because µ ∗ is the minimizer of E .
Then lim k →∞ E ϵ k ( µ ∗ ϵ k ) =
E ( µ ∗ ) = 1 .
Moreover, χ 2 ( ν ∥ µ ∗ ) = 0 .
This can only happen if ν and µ ∗ agree on all Borel sets, so ν = µ ∗ .
A. D IFFERENTIAL CALCULUS OF E ϵ IN P 2 ( R n )
Lemma A.17.
Let x 0 ∈ R , h > 0 , and Ω ⊂ R m be a compact set.
Suppose f : ( x 0 − h, x 0 + h ) × Ω
→ R is jointly continuous and the derivative ∂∂x f : ( x 0 − h, x 0 + h ) ×
Ω → R exists and is jointly continuous.
Then (cid:82) Ω f ( x, ω ) d ω is differentiable for x ∈ ( x 0 − h, x 0 + h ) , and
[Equation]
where the integration is with respect to the Lebesgue measure L m on Ω .
Proof.
Fix x ∈ ( x 0 − h, x 0 + h ) and let t > 0 be small enough such that [ x − t, x + t ] ⊂ ( x 0 − h, x 0 + h ) .
Note that (cid:82) Ω f ( x, ω ) d ω is well-defined by the dominated convergence theorem since sup x ∈ [ x − t,x + t ] ,
ω ∈ Ω
f
( x, ω ) < ∞ and L m (Ω) < ∞ .
Define θ ( ω ) :
= sup x ∈ [ x − t,x + t ] (cid:12)(cid:12) ∂∂x
f ( x, ω ) (cid:12)(cid:12) which is finite since ∂∂x f
( x, ω ) is continuous, so that θ ( ω ) ≥ ∂∂x | f ( x, ω ) | for all x ∈ [ x − t, x + t ] and θ is integrable since L m (Ω) < ∞ .
Hence by the differentiation lemma (Klenke, 2013, Theorem 6.28) we are done.
The following lemma is similar to Korba et al.
(2021, Proposition 3) but with different assumptions: we do not put any integrability assumptions on W ϵ , but we do restrict measures to have compact support.
For a differentiable f : R n × R n → R , we use ∇ 1 f ( x, y ) to denote the gradient with respect to x .
We also use H 1 f ( x, y ) to denote the Hessian of f with respect to x .
We similarly denote ∇ 2 f ( x, y ) and H 2 f ( x, y ) .
Lemma A.18.
Assume µ ∈ P 2 ( R n ) has density q ∈ C 1 c ( R n ) .
Let ξ ∈ C
( R n ; R n ) .
Denote s ( x, t ) :
= x + t ξ ( x ) .
Then for all t > 0 and all ϵ > 0 ,
[Equation]
In particularly,
[Equation]
Moreover, d
[Equation]
[Equation]
Proof.
We compute
[Equation]
Since µ has compact support and ( x, y, t ) (cid:55)→ W ϵ ( s ( x, t ) , s ( y, t ))
q ( x ) q ( y ) is jointly continuous and its derivative with respect to t is also jointly continuous, by Lemma A.17, we can push dd t inside the double integral and we obtain (16).
Another application of Lemma A.17 shows that if we take derivative with respect to t again on (16) and evaluate at 0 we obtain (18).
Lemma A.19.
Let f ∈ C 1 ( R n ) , p ∈ [1 , ∞ ] .
Assume ϕ ϵ has compact support for some ϵ > 0 .
Then for all i = 1 , . . . ,
n , ϕ ϵ ∗ f is differentiable and
[Equation]
Proof.
Since supp ( ϕ ϵ ) is compact and f ∈ C 1 ( R n ) is bounded on any compact set, we know ϕ ϵ ∗ f is well-defined at every x ∈ R n .
Note that
[Equation]
Since supp ( ϕ ϵ ) is compact and ( x, y )
(cid:55)→ f ( x − y ) ϕ
( y ) is C 1 , by Lemma A.17 we justify the existence of the derivative and the exchange of differentiation and integration (?).
A.4. S UBDIFFERENTIALS OF E ϵ
Recall the following notion of a “Wasserstein gradient” in P 2 ( R n ) from Ambrosio et al.
Definition 10.1.1).
Definition A.20.
A vector field w ∈ L 2 ( µ ; R n ) is a strong Fr ´ echet subdifferential of a functional
F : P 2 ( R n ) → ( −∞ , + ∞ ] if for all T ∈
L 2 ( µ ; R n ) , the following holds:
[Equation]
Note that we cannot apply Ambrosio et al.
(2005, Lemma 10.4.1) directly to prove (4) because interaction energies cannot be written in the form of (10.4.1) in their setup.
Proof of Proposition 3.7.
Let ξ ∈ C ∞ c
( R n ; R n ) .
By Lemma A.18, we have
[Equation]
where the last step follows from applying Lemma A.17 since q has compact support.
Now suppose w ∈ L 2 ( µ ; R n ) is a strong Fr´echet subdifferential satisfying (19).
For the sequence { T t } , we have by definition
[Equation]
Hence
[Equation]
The previous calculation shows dd t | t =
E ϵ ( µ t ) exists, and hence it is equal to (cid:82) w ( x ) ⊤ ξ ( x ) d µ ( x ) .
This proves for any ξ ∈ C ∞ c ( R n ; R n ) ,
[Equation]
Hence we have shown (4).
Finally, we show lim ϵ
→ 0 w ϵ ( x ) =
w χ 2 ( x ) for µ -a.e.
x under the additional assumption that ϕ ϵ has compact support.
By Ambrosio et al.
(2005, Lemma 10.4.1) with F ( x, ρ ( x ))
= (cid:16) ρ ( x ) p [Equation] x ) − 1 (cid:17) 2 p ( x ) , we find the strong subdifferential of χ 2 ( · ∥ µ ∗ ) is given by
[Equation]
To show (5), we compute, for µ -a.e.
x ,
[Equation]
where for the last equality we have applied Lemma A.19.
Now taking ϵ → 0 , by Proposition A. using the fact that supp ( ϕ ϵ ) is compact (so that q/ √ p and ∇ ( q/ √ p ) are bounded on x + supp ( ϕ ϵ ) ), we obtain (5).
A.4. D ISPLACEMENT CONVEXITY OF E ϵ AT µ ∗ AS ϵ →
The statement of Proposition 3.8 is similar in form as Korba et al.
(2021, Corollary 4) but in our case we do not have the second term in Equation (18) vanishing and we need to take the limit ϵ
→ 0 .
We also do not resort to RKHS theory in the proof.
Proof of Proposition 3.8.
By (18)
[Equation]
we have
[Equation]
where
[Equation]
We tackle F ϵ first.
Observe that successive application of integration by parts using the fact that ξ has compact support gives
[Equation]
If we view (cid:80) ni =1 ∂∂x i ( ξ i ( x ) p ( x )) as the density of a signed measure (it is integrable since it has compact support), and since W ϵ is i.s.p.d.
on the support of ξ by Proposition 3.5(a), we see that each double integral in the last expression is non-negative.
Hence F ϵ ≥ 0 .
Next we show lim ϵ
→ 0 G ϵ = 0 .
Since µ ∗ has compact support, by Fubini’s theorem,
[Equation]
To expand the integral inside the Hessian operator, we have
[Equation]
First by the chain rule and Lemma A.19, we have
[Equation]
Differentiating again while applying Lemma A.19, we obtain after rearranging terms,
[Equation]
By Proposition A [Equation] 4 and the fact that p ∈ C 2 c ( R n ) [Equation]
we have
[Equation]
Finally
[Equation]
we have
[Equation]
where interchanging the limit and the integral is jusfitied by the dominated convergence theorem and the fact that p and ϕ ϵ have compact support (we need compact support assumption of ϕ ϵ to make sure convolutions appearing in (21) are uniformly bounded when ϵ is sufficiently small)
[Equation]
A.4. A DESCENT LEMMA FOR E ϵ WITH TIME DISCRETIZATION
We now consider a time discretization of the gradient flow of E ϵ given by, for an initial measure µ 0 ∈ P 2 ( R n ) and m ∈ Z > 0 , with a step size γ > 0 ,
[Equation]
Similar to Korba et al.
(2021, Proposition 14), we can show the following descent lemma for iterations (22) with the following additional assumptions:
Assumption A.21.
Suppose the target density p ∈ C 2 ( R n ) satisfies 1 /C p ≤ p ( x ) ≤ C p ,
∥ p ( x ) ∥ 2 ≤ C ′ p , ∥
H p ( x ) ∥ 2 ≤ C ′′ p for some C p , C ′ p , C ′′ p for all x ∈ R n , where we use ∥ A ∥ 2 to indicate the matrix spectral norm (i.e. ∥ A ∥ 2 = σ max ( A ) ).
Suppose the mollifier ϕ ϵ ∈ C 2 ( R n ) satisfies, in addition to Definition 3.1, ϕ ϵ ( x ) ≤ C ϵ , ∥∇ ϕ ϵ ( x ) ∥ 2 ≤ C ′ ϵ , and ∥ H ϕ ϵ
( x ) ∥ 2 ≤ C ′′ ϵ for some C ϵ , C ′ ϵ , C ′′ ϵ .
Lemma A.22.
Under Assumption A.21, there exists a constant L > 0 such that the function ∇ 1 W ϵ :
R n × R n → R n is L -Lipschitz in terms of ∥·∥ 2 in either input.
Proof.
Denote r ( x ) :
= p ( x ) − 1 / 2 .
Then our assumptions imply r ( x ) ≤ C r :=
C 1 / 2 p , ∥∇ r ( x ) ∥ 2 ≤
C ′ r := 12 C 3 / 2 p C ′ p , and ∥ H 1 r ( x ) ∥ 2 ≤ C ′′ r := 34 C 5 / 2 p C ′ 2 p + 12 C 3 / 2 p C ′′ p .
We compute, for x, y ∈ R n ,
[Equation]
[Equation]
Then
[Equation]
and ∇ 2 ∇
[Equation]
Then we have
[Equation]
Hence we conclude that ∇ 1 W ϵ is L -Lipschitz with
[Equation]
Remark A.23.
Compared with Korba et al.
(2021, Lemma 1), to ensure ∇ 1 W ϵ is Lipscthiz, we only require uniform boundedness of p and ϕ ϵ up to second order derivatives instead of up to order 3.
Proposition A.24.
Under Assumption A.21, suppose µ 0 ∈ P 2 ( R n ) has compact support.
Then for any γ ≤ 12 L , with L defined as in Lemma A.22 and { µ m } m> 0 defined as in (22) ,
[Equation]
Proof.
We first show µ m has compact support for all m ∈ Z ≥ 0 .
For µ ∈ P 2 ( R n ) with compact support, the proof of Proposition 3.7 implies (note Proposition 3.7 assumes µ has density but it is not necessary to obtain the following formula using the same proof)
[Equation]
Since x (cid:55)→ ∇ 1 W ϵ ( x, y ) is C 1 by the assumptions and µ has compact support, by Lemma A.17, we see that w ϵ,µ is differentiable with
[Equation]
By induction, suppose µ m has compact support.
Then
[Equation]
Since w ϵ,µ m is continuous, the set on the right-hand side is bounded.
Hence µ m +1 has compact support.
Now fix m ∈ Z >
and we show (23).
Define a path { µ t } t ∈ [0 , 1] defined by µ t
= ( I − γt w ϵ,µ m )
# µ m .
Let f ( t ) :
= E ϵ ( µ t ) .
By Lemma A.18, the continuity of w ϵ,µ m implies that f is differentiable.
Moreover, another application of Lemma A.17 implies that f is twice differentiable, and in particular continuously differentiable.
Hence f is absolutely continuous on the compact interval [0 , 1] .
By the fundamental theorem of calculus, we have
[Equation]
Observe that by Lemma A.18,
[Equation]
where we apply Fubini’s theorem to exchange the order of integration together with the fact that µ m has compact support and the integrand is continuous.
Let s ( x, t ) :
= x − γt w ϵ,µ m ( x ) .
Note that, again by Lemma A.18,
[Equation]
[-ϵ ( x, y ) ∥ 
-]
Note that, by Lemma A.22, we have
[Equation]
Thus
[Equation]
where we have used the Cauchy-Schwartz inequality in the last step
[Equation]
Combining everything [Equation] we have shown
[Equation]
since γ < 12 L .
B W EIGHTED HYPERSINGULAR R IESZ ENERGY
We recall results most relevant to us from Borodachov et al.
Suppose X ⊂ R n is compact, of Hausdorff dimension d , and d -rectifiable, i.e., the image of a bounded set in R d under a Lipschitz mapping.
Given a non-vanishing continuous probability density p on X , define a measure H pd ( B ) :=
(cid:82) B p ( x ) d H d ( x ) for any Borel set B ⊂ X .
The target measure h pd is defined to be h pd ( B ) :=
H pd ( B ) / H pd ( X ) .
The weighted s -Riesz energy is defined as the interaction energy
[Equation]
The following result states that minimizers of the weighted s -Riesz energy approximate the target measure when the s is sufficiently large.
Theorem B.1 (Theorem 11.1.2, Borodachov et al.
Suppose s > d .
For ω ∗ N ∈ arg min E s with ω ∗ N = { x N 1 , . . .
, x NN } , we have weak convergence δ ω ∗ N → h pd as N → ∞ .
A similar result with a slightly different assumption holds for s =
d (Theorem 11.1.3 of Borodachov et al.
Comparison of (24) with discretized MIE (6) .
Our theory is only valid for the full-dimensional case, i.e., n = d (see a few exceptions discussed in Remark A.10).
When this is the case and if the mollifier is taken to be from the s -Riesz family, (6) becomes, for s > n ,
[Equation]
Compared with (24), in our case there is an ϵ in the denominator, so that the continuous version of the energy E ϵ does not blow up all the time (this is in contrast with the continuous version of (24)—see Borodachov et al.
(2019, Theorem 4.3.1)).
Moreover, the exponential scaling on p is different: in our case we use − 1 / 2 whereas in (24) it is − s/ 2 n .
The diagonal i =
j is included in our energy, but this is inconsequential as another valid discretization is to discard the diagonal term (Borodachov et al., 2019, Theorem 4.2.2).
On the other hand, our theory allows a bigger class of mollifiers that are not necessarily of Riesz families.
We also allow X to be non-compact.
An interesting future research direction is to extend our theory to cases where d < n , e.g., when X is an embedded d -dimensional submanifold in R n .
C A LGORITHMIC DETAILS
In this section we provide algorithmic details of MIED and compare with the updates of SVGD (Liu & Wang, 2016).
The negative gradient of (7) with respect to x i is
[Equation]
where
[Equation]
and to get the last equality we used the fact that ∇ ϕ (0) =
thanks to the assumption ϕ ( x ) = ϕ
( − x ) .
Then gradient descent on (7) gives our algorithm (Algorithm 1).
The special treatment of the diagonal terms described in Section 4 amounts to modifying only the diagonal I ii .
C
C OMPARISON WITH SVGD
The update formula in Algorithm 1 is similar to the one in SVGD: if we use ϕ ϵ ( x − y ) in place of the kernel k ( x, y ) in the SVGD update and rewrite:
[Equation]
For both algorithms, the update formula for each particle consists of attraction and repulsion terms and the total time complexity of each update iteration is O ( N 2 ) .
We note the following differences.
First, in our formulation we have scaling factors
[Equation]
i,j e Iij which help stabilizing the optimization (as a by-product of working in the logarithmic domain) and put more weight on nearby particles as well as particles in low-density regions, whereas in (SVGD) the scaling factors are ϕ ϵ ( x j − x i ) which are not adapted to prioritize low-density regions.
Second, in MIED, the attraction force for particle i only comes from ∇ log p ( x i ) , whereas in (SVGD) the attraction force comes from ∇ log p ( x j ) for all j ’s. Third, for each j , in our formulation the repulsive force has an additional factor of 2 in front of ∇ log ϕ ϵ ( x j − x i ) .
Empirically, the additional scaling factors in MIED help produce samples with good separations compared to SVGD, since closer pairs of points will have large weights.
Additionally, since MIED optimizes a finite-dimensional objective (7), we can employ accelerated gradient-based optimizers like Adam (Kingma & Ba, 2014), which we used in our experiments.
In contrast, SVGD does not optimize any finite-dimensional objective.
While practical SVGD implementations also use optimizers like Adam, it is unclear how the resulting particle dynamics is related to the gradient flow of KL divergence.
C. H ANDLING CONSTRAINTS WITH DYNAMIC BARRIER METHOD
The dynamic barrier method is introduced in Gong & Liu (2021) which solves min x f ( x ) subject to g ( x ) ≤ 0 where g is scalar-valued.
Intuitively, their method computes update directions by either decreasing g ( x ) when g ( x ) > 0 , following −∇ f ( x ) if the constraints are satisfied, or balancing both gradient directions.
In order to handle multiple constraints such as in Figure 11, we consider a generalized version of their dynamic barrier method.
In this generalized setting, g : R n → R m is vector-valued and constraints are g ( x ) ≤ 0 where the ≤ sign is interpreted coordinate wise..
Suppose we are at iteration t with current solution x t .
Then the next update direction v ∗ is taken to be the arg min of
[Equation]
(26) where α i > 0 are fixed hyperparameters; in our implementation we simply choose α i = 1 .
Then x t +1 = x
t − ηv ∗ with learning rate η .
In our implementation we use Adam (Kingma & Ba, 2014) that modulates the update directions.
Observe that the optimization problem (26) is the same as projecting a point ∇ f ( x t ) onto the polyhedron formed by the intersection of the halfspaces ∩ mi =1 { x ∈ R n : ∇ g i ( x t ) ⊤ v ≥ α
i g i ( x t ) } .
To solve (26), we use Dykstra’s algorithm (Tibshi- rani, 2017) which can be interpreted as running coordinate descent on the dual of (26).
We use a fixed number of 20 iterations for the Dykstra’s algorithm which we found to be sufficient for our experiments; in the case of a single constraint, we only need to use one iteration.
D E XPERIMENT DETAILS AND ADDITIONAL RESULTS D. G AUSSIANS IN VARYING DIMENSIONS
We generate the n -dimensional Gaussians used to produce Figure 1 as follows.
We generate a matrix √ A ∈ R n × n with i.i.d.
entries uniformly in [ − 1 , 1] .
Then we set A = √ A √ A ⊤ / det (cid:16) √ A (cid:17) .
This way det( A ) = 1 .
We then use A as the covariance matrix for the Gaussian (centered at 0).
We use Adam with learning rate 0 .
01 for all methods for a total of 2000 iterations.
This is enough for
SVGD and MIED to converge, while for KSDD the convergence can be much slower.
We visualize the samples from each method for n = 2 in Figure 5.
We notice that MIED is capable of generating well-separated samples, while for SVGD there is a gap between the inner cluster of samples and a sparse outer ring.
For KSDD we see the artifact where too many samples concentrate on the diagonal. [Figure] [Figure] [Figure]
[Figure]
.d.
[Figure]
. D. C OLLAPSED SAMPLES WHEN THE KERNEL WIDTH IS TOO BIG
In Figure 3, we see that samples from SVGD collapse with an adaptive kernel where the variance is taken to be half of the median of the squared distance among all pairs of points (Liu & Wang, 2016); at termination the median of the squared distance is greater than 1 in that experiment.
Here we investigate this issue further.
In Figure 6, for the same uniform sampling setup, we visualize the results of SVGD with a fixed kernel width and MIED with Gaussian mollifiers with the same kernel width: when the kernel width (i.e. 2 ϵ 2 in ϕ gϵ ( x ) = exp (cid:16) − ∥ x
∥ 22 2 ϵ 2 (cid:17) /Z
gϵ ) is too big, both SVGD and
MIED result in collapsed samples.
This is because since the target density is a constant, the only force in the updates is the repulsive force.
When the kernel width is too large, the repulsive force coming from points in the same collapsed cluster is dominated by the repulsive force coming from points from other clusters—this is evident in the update directions shown in the leftmost column of Figure 6 (SVGD with kernel width 0 . 1 ).
[-1 − 1 .− 0 .0 .0 .
-]
When using the dynamic barrier method Gong & Liu (2021) to enforce the square constraints instead of reparameterization with tanh , we obtain similar results as in Figure 6.
This pathological phenomenon is not only limited to constrained sampling: when sampling the 2D
Gaussian from Figure 5, using too big a kernel width can also result in collapsing (Figure 7).
We emphasize that our theory of MIED suggests that in practice we need to choose the kernel width very small in order to sample from the correct target measure according to Theorem 3.3 and
Theorem 3.6.
In comparison, the theory of SVGD has no such implication.
D. U NIFORM SAMPLING WITH AN ALTERNATIVE MIRROR MAP
In this section, we show that for sampling from a uniform distribution in the square [ − 1 , 1] 2 , the results of SVMD/MSVGD (Shi et al., 2021) depend heavily on the choice of the mirror map. Instead of the entropic mirror map used to produce results in Figure 3, mirror map [Table]
[Figure]
[Figure]
.1 .−
1 − 1 .− 0 .0 .0 .1 .−
1 − 1 .− 0 .0 .0 .1 .−
{+1 − 1 .− 0 .0 .0 .
+}
.s
t e p− 1 − 1 .− 0 .0 .0 .1 .−
1 − 1 .− 0 .0 .0 .1 .−
1 − 1 .− 0 .0 .0 .1 .−
1 − 1 .− 0 .0 .0 .1 .−
1 − 1 .− 0 .0 .0 .1 .
ϕ ( θ ) = (cid:80) ni
{+=1 (cid:16) log 1 1 − θ i + log 1 
+}
+ θ i (cid:17) as in Ahn & Chewi (2021).
The results are shown in Figure 8.
SVMD/MSVGD fail to draw samples near the boundary; we suspect this is because the gradient of the conjugate ∇ ϕ ∗ ( η )
= ( √ 1+ η 2 − 1) /
η (coordinate-wise arithmetic) requires coordinates of η to go to ∞ to land near the boundary.
We verify this phenomenon by using ∇ ϕ ∗ ( η ) as the reparametrization map in MIED (rightmost figure in Figure 8): indeed with such reparameterization MIED also struggles near the boundary.
−
[Figure]
1 − 1 .− 0 .0 .0 .1 .
SVMD − 1 − 1 .− 0 .0 .0 .1 .
MSVGD − 1 − 1 .− 0 .0 .0 .1 .
MIED(Riesz) w/ mirror D. 20- DIMENSIONAL D IRICHLET DISTRIBUTION We sample from the 20-dimensional Dirichlet distribution in the same setup as in Shi et al.
(2021) with 50 particles.
Results and visualization are shown in Figure 9.
We see that unlike sampling from a box (Figure 3), both MSVGD and SVMD by Shi et al.
(2021) perform well in this experiment.
This is due to the fact that the entropic mirror map used here is a well-tested choice for simplex constraint, yet obtaining a good mirror map for a generic constraint, even if it is linear, can be challenging.
Our method does not have such a limitation, as it can easily incorporate existing constrained optimization tools [Figure]
D. E FFECT OF s FOR R IESZ MOLLIFIERS
When we use the s -Riesz families of mollifiers in MIED, we have the freedom of choosing the hyperparameter s so long as s > n .
In this section, we study the effect of s on the resulting samples.
We consider the problem of sampling from a mixture of four 2D Gaussians centered at ( ± 1 , ± 1) , each with diagonal variance 0 .
3 and constrained to the [ − 1 , 1] 2 box.
We vary s in [2 , 10] and the number of particles N in { 100 , 200 , 500 , 1000 , 2000 } .
All runs use a total of 1000 iterations with learning rate 0 .
In the top of Figure 10, we plot the W 2 distance and energy distance as functions of s for each N .
Interestingly, we see the best performing s is in [3 , 5 . 0] and depends on N .
This suggests that our choice of s = n + 10 − 4 in Section 5 may not be optimal and there is room for hyperparameter tuning to further improve the performance of MIED with Riesz kernel.
At the bottom of Figure 10 we visualize the samples of MIED with s = 3 and of SVGD with kernel width 0 .
01 (adaptive kernels would result in collapsed samples and other widths we tested on would result in worse samples).
While SVGD samples form visible artifacts, the samples of MIED are evenly distributed and the four modes of the mixtures emerge as N increases.
[Figure]
D. M ORE CONSTRAINED SAMPLING EXPERIMENTS
In this section we test MIED on more low dimensional constrained sampling problems and qualitatively assess the results.
Note that mirror LMC
(Zhang et al., 2020;
Ahn & Chewi, 2021) or mirror
SVGD (Shi et al., 2021) cannot be applied due to non-convexity of the constraints.
In Figure 11, we consider uniform sampling of a challenging 2D region with initial samples drawn from the topright corner: as the number of iterations increases, MIED gradually propagate samples to fill up the entire region.
In Figure 12, we consider sampling from a von Mises-Fisher distribution on a unit sphere.
Although our theory focuses on sampling from a full-dimensional distribution, as discussed in Remark A.10, we can extend Theorem 3.3 to the case of a sphere due to its symmetry.
We see the samples visualized in Figure 12 capture the two modes that emerge by restricting the Gaussian to a unit sphere.
We use 80% / 20% training/test split as in Liu et al. We use the source code provided by Liu et al.
(2021) with default hyperparameters. The source code provided by the authors of Liu et al. (2021) does not implement the calculation of g ( θ ) faithfully as written in the formula, so we corrected it. All methods use 2000 iterations for training.
For our method we use learning rate 0 . One of their four methods (Control+SVGD) got stuck at initialization (with accuracy around 0 . 75 ), so we omit its result from the plot in Figure 4. − 1 − 1 .− 0 .0 .0 .1 .−
1 − 1 .− 0 .0 .0 .1 .− 1 − 1 .− 0 .0 .0 .1 .− 1 − 1 .− 0 .0 .0 .1 .− 1 − 1 .− 0 .0 .0 .1 .−
1 − 1 .− 0 .0 .0 .1 .− 1 − 1 .− 0 .0 .0 .1 .− 1 − 1 .− 0 .0 .0 .1 .−
1 − 1 .− 0 .0 .0 .1 .− 1 − 1 .− 0 .0 .0 .1 .−
1 − 1 .− 0 .0 .0 .1 .− 1 − 1 .− 0 .0 .0 .1 .−
1 − 1 .− 0 .0 .0 .1 .− 1 − 1 .− 0 .0 .0 .1 .−
1 − 1 .− 0 .0 .0 .1 .− 1 − 1 .− 0 .0 .0 .1 .−
1 − 1 .− 0 .0 .0 .1 .− 1 − 1 .− 0 .0 .0 .1 .