Convergence of Gaussian-smoothed optimal transport distance with sub-gamma distributions and dependent samples

In this paper, we focus on the OT cost $\mathcal{T}_{p}(P,Q)$ associated with the cost function $c_{p}(x,y)=\|x-y\|^{p}$ for $p>0$ and $c_{0}(x,y)=1_{\{x\neq y\}}$. The total variation distance is given by $\mathcal{T}_{0}(P,Q)$ and the $p$-Wasserstein distance is given by $\mathcal{T}_{p}(P,Q)$ if $p\in(0,1]$ and $(\mathcal{T}_{p}(P,Q))^{1/p}$ if $p>1$ Villani (2003).

The minimax convergence rate of $\mathcal{T}_{p}(P,P_{n})$ was established by (Fournier and Guillin, 2015, Theorem 1) who showed that if $P$ has a moment strictly greater than $2p$, then

$\displaystyle\mathbb{E}\mathopen{}\mathclose{{}\left[\mathcal{T}_{p}(P,P_{n})}\right]\asymp\begin{cases}n^{-\frac{1}{2}},&p>d/2\\ n^{-\frac{1}{2}}\log n,&p=d/2\\ n^{-\frac{p}{d}},&p\in(0,d/2)\end{cases}.$

(2)

Unfortunately, this means that the sample complexity increases exponentially with the dimension for $d>2p$.

Recently, Goldfeld et al. (2020b) showed that one way to overcome the curse of dimensionality is to consider the Gaussian-smoothed OT distance, defined as

$\displaystyle\mathcal{T}_{p}^{(\sigma)}(P,Q)\coloneqq\mathcal{T}_{p}(P\ast\mathcal{N}_{\sigma},Q\ast\mathcal{N}_{\sigma}),$

where $\mathcal{N}_{\sigma}$ denotes the iid Gaussian measure with mean zero and variance $\sigma^{2}$. Under the assumption that $P$ is sub-Gaussian with constant $v$, they proved an upper bound on the converge rate that is independent of the dimension:

$\displaystyle\mathbb{E}\mathopen{}\mathclose{{}\left[\mathcal{T}_{p}(P,P_{n})}\right]\leq\frac{C_{d,p,\sigma,v}}{\sqrt{n}},\quad p\in\{0,1,2\}.$

The precise the form of the constant $C_{d,p,\sigma,v}$ is provided for $p\in\{0,1\}$ but not for the case $p=2$ unless $P$ is also assumed to have bounded support. Ensuing work by Goldfeld et al. (2020a) established the same convergence rate for $p=1$ under the relaxed assumption that $P$ has finite moment grater than $2d+2$.

Metric properties of GOT were studied by Goldfeld and Greenewald (2020) who showed that $\mathcal{T}_{1}^{(\sigma)}(P,Q)$ is a metric on the space of probability measures with finite first moment and that the sequence of optimal couplings converges in the $\sigma\to 0$ limit to the optimal coupling for the unsmoothed Wasserstein distance. Their arguments depend only on the pointwise convergence of the characteristic functions under Gaussian smoothing, and thus also apply to the case of general $p$ considered in this paper.

One of the main contributions of this paper is to prove an upper bound on the convergence rate for all orders of $p$ and under more general assumptions on $P$. Specifically, we prove the following result:

This result brings the GOT framework in line with the general setting studied by (Fournier and Guillin, 2015, Theorem 1), and shows that the benefits obtained by smoothing extend beyond the special cases of small $p$ and well-controlled tails. To help interpret this result it is important to keep in mind that for $p>1$, the Wasserstein distance is given by the $p$-th root of the GOT. As for the tightness of the bound, there are two regimes worth considering, namely when $\sigma\to 0$ as $n\to\infty$ and when $\sigma$ is fixed. In the former case, the dependence on $\sigma$ seems to be nearly tight. In Section 4.1, we show that if $\sigma\asymp n^{-\frac{1}{d+2p}}$ then Theorem 1 implies an upper bound on the unsmoothed convergence rate

$\displaystyle\mathbb{E}\mathopen{}\mathclose{{}\left[\mathcal{T}_{p}(P,P_{n})}\right]\leq C_{d,p,s}m^{p}n^{-\frac{p}{d+2p}}.$

(4)

Notice that for $d\ll 2p$ and $d\gg 2p$ this recovers the minimax convergence rate given in (2).

The main technical step in our approach is to establish a novel connection between GOT and a family of kernel MMD distances (Theorem 2). We then show how a particular member of this family, which we call the ‘two-moment’ kernel, defines a metric on the space of probability measures with finite moments strictly greater than $p+d/2$ (Theorem 3).

In addition to Theorem 1, we also provide further results that elucidate the role of the dimension as well the tail behavior of the underlying distribution (Theorem 6). Furthermore, we address the setting of dependent samples and provide an example illustrating how the connection with MMD can be used to go beyond the usual assumptions involving mixing conditions for stationary processes. Finally, we provide some numerical experiments that support our theory.

The convergence of OT distances continues to be an active area of research Singh and Póczos (2018); Niles-Weed and Berthet (2019); Lei et al. (2020). Building upon the the work of work of Cuturi (2013), a recent line of work has focused on entropy regularized OT defined by

$\displaystyle S_{\epsilon}(P,Q)\coloneqq\inf_{\pi\in\Pi(P,Q)}\int c(x,y)\,d\pi(x,y)+\epsilon D(\pi\|P\otimes Q)$

where $D(\mu\|\nu)=\int\log\frac{d\mu}{d\nu}d\mu$ is the relative entropy between probability measures $\mu$ and $\nu$. The addition of the regularization term facilitates the numerical approximation using the Sinkhorn algorithm. The amount of regularization interpolates between OT in the $\epsilon\to 0$ limit and the kernel MMD in $\epsilon\to\infty$ limit; see Feydy et al. (2019). In contrast to the Gaussian-smoothed Wasserstein distance, entropy regularized OT is not a metric since it does not satisfy the triangle inequality. Convergence rates for entropy regularized OT were obtained by Genevay et al. (2019) under the assumption of bounded support and more recently by Mena and Niles-Weed (2019) under the assumption of sub-Gaussian tails. Further properties have been studied by Luise et al. (2018) and Klatt et al. (2020).

There has also been work focusing on the sliced Wasserstein distance, which is obtained by averaging the one-dimensional Wasserstein distance over the unit sphere Rabin et al. (2011); Bonneel et al. (2015). While the sliced Wasserstein distance is equivalent to the Wasserstein distance in the sense that convergence in one metric implies convergence in another, the rates of convergence need not be the same. See Section 1.2 of Goldfeld et al. (2020a) for further discussion.

Going beyond convergence rates for empirical measures, properties of smoothed empirical measures have been studied in a variety of contexts, including the high noise limit Chen and Niles-Weed (2020) and applications to the estimation of mutual information in deep networks Goldfeld et al. (2018). Finally, we note that there has also been some work on convergence with dependent samples by Fournier and Guillin (2015), who focus on OT distance, and also by Young and Dunson (2019), who consider a closely related entropy estimation problem.

Consider the feature map $\psi_{x}\colon\mathbb{R}^{d}\to[0,\infty)$ defined by

$\displaystyle\psi_{x}(z)$

$\displaystyle\coloneqq\frac{\sqrt{\omega_{d}}}{2^{\frac{d+p}{2}}}\frac{\|z\|^{\frac{d-1+2p}{2}}}{\sqrt{f(\|z\|)}}\phi\mathopen{}\mathclose{{}\left(\frac{z}{\sqrt{2}}-\frac{x}{\sigma}}\right),$

(7)

where $\omega_{d}=2\pi^{d/2}/\Gamma(d/2)$ is the volume of the unit sphere in $\mathbb{R}^{d}$, $\phi(u)=(2\pi)^{-d/2}\exp(-\frac{1}{2}\|u\|^{2})$ is the standard Gaussian density on $\mathbb{R}^{d}$, and $f$ is a probability density function on $[0,\infty)$ that satisfies

$\displaystyle f(x)\geq ax^{d+2p-1}\exp(-bx^{2}),$

(8)

for some $a>0$ and $b\in(0,1/2)$. This feature map defines a positive semidefinite kernel $k$ according to

$\displaystyle k(x,y)=\langle\psi_{x},\psi_{y}\rangle.$

After some straightforward manipulations (see Appendix A), one finds that $k(x,y)$ is finite on $\mathbb{R}^{d}\times\mathbb{R}^{d}$ and can be expressed as the product of a Gaussian kernel and a term that depends only on $\|x+y\|$. Specifically,

$\displaystyle k(x,y)$

$\displaystyle=\exp\mathopen{}\mathclose{{}\left(-\frac{\|x-y\|^{2}}{4\sigma^{2}}}\right)I_{f}\mathopen{}\mathclose{{}\left(\frac{\|x+y\|}{\sqrt{2}\sigma}}\right),$

(9)

where

$\displaystyle I_{f}(u)$

$\displaystyle\coloneqq\frac{\omega_{d}}{2^{d+p}(2\pi)^{\frac{d}{2}}}\int_{0}^{\infty}\frac{x^{d-1+2p}}{f(x)}g_{d,u}\mathopen{}\mathclose{{}\left(x}\right)\,dx,$

(10)

and $g_{d,u}(x)$ is the density of the non-central chi-distribution given in (6). Note that this kernel is not shift invariant because of the term $I_{f}(u)$.

The next result shows that the MMD defined by this kernel provides an upper bound on the GOT.

The significance of Theorem 2 is twofold. From the perspective of GOT, it provides a natural connection between the role of Gaussian smoothing and normalization of the kernel. From the perspective of MMD, Theorem 2 describes a family of kernels that metrize convergence in distribution as well as convergence in $p$-th moments.

Similar to the analysis of convergence rates in previous work Goldfeld and Greenewald (2020); Goldfeld et al. (2020b), the proof of Theorem 2 builds upon the fact that $\mathcal{T}_{p}(P,Q)$ can be upper bounded by a weighted total variation distance (Villani, 2008, Theorem 6.13). The novelty of Theorem 2 is that it establishes an explicit relationship with the kernel MMD and also provides a much broader class of upper bounds parameterized by the density $f$.

One potential limitation of Theorem 2 is that for a particular choice of density $f$, the requirement that $\sqrt{k(x,x)}$ is integrable might not be satisfied for probability measures of interest. For example, the convergence rates in Goldfeld and Greenewald (2020) and Goldfeld et al. (2020b) can be obtained as a corollary of Theorem 2 by choosing $f$ to be the density of the generalized gamma distribution. However, the inverse of this density grows faster than exponentially, and as a consequence, the resulting bound can be applied only to the case of sub-Gaussian distributions.

The main idea underlying the approach in this section is that choosing a density with heavier tails leads to an upper bound that holds for a larger class of probability measures. Motivated by the functional inequalities appearing in Reeves (2020), we consider the following density function, which belongs to the family of generalized beta-prime distributions:

$\displaystyle f(x)=\frac{\epsilon}{2\pi x}\mathopen{}\mathclose{{}\left(\mathopen{}\mathclose{{}\left(\frac{x}{\lambda}}\right)^{-\epsilon}+\mathopen{}\mathclose{{}\left(\frac{x}{\lambda}}\right)^{\epsilon}}\right)^{-1}.$

For this special choice, the function $I_{f}(u)$ can be expressed as the weighted sum of two moments of the non-central chi distribution. Starting with (9) and simplifying terms leads to the following:

A useful property of the two-moment kernel is that it satisfies the upper bound

$\displaystyle k(x,y)\leq C_{d,p,\epsilon,\lambda}\mathopen{}\mathclose{{}\left(1+\|x\|^{d+2p+\epsilon}+\|y\|^{d+2p+\epsilon}}\right),$

where the constant depends only on $(d,p,\epsilon,\lambda)$ (see Appendix A for details). As a consequence:

We begin with an upper bound on $\mathbb{E}[k(X,X)]$ that holds whenever $\|X\|$ has a moment greater than $d+2p$.

In view of (15), Theorem 1 follows as an immediate corollary.

Next, we provide a refined bound for distributions satisfying a sub-gamma tail condition.

Properties of sub-Gaussian and sub-gamma distributions have been studied extensively; see e.g,  Boucheron et al. (2013). In particular, if $X_{1}$ and $X_{2}$ are independent sub-gamma random vectors with parameters $(v_{1},b_{1})$ and $(v_{2},b_{2})$, respectively, then $X_{1}+X_{2}$ is sub-gamma with parameters $(v_{1}+v_{2},\max(b_{1},b_{2}))$.

The next result provides an upper bound on the moments of the magnitude of a sub-gamma vector. Although there is a rich literature this topic, we were unable to find a previous statement of this result and so it may be of independent interest.

In view of (15), Theorem 6 gives an upper bound on the convergence rate of $\gamma_{k}(P,P_{n})$ with an explicit dependence on the sub-gamma parameters $(d,p,\sigma,v,b)$. For the special case of a sub-Gaussian distribution ($b=0$) and $p\in\{0,1\}$, this bound recovers the results in Goldfeld and Greenewald (2020). Going beyond the setting of sub-Gaussian distributions (i.e., $b>0$) this bound quantifies the dependence on the dimension $d$ and the scale parameter $b$.

Motivated by applications involving Markov chain Monte Carlo there is significant interest in understanding the rate of convergence when there is dependence in the samples. Within the literature, this question is often address by focusing on stationary sequences satisfying certain mixing conditions Peligrad (1986). The basic idea is that the dependence between $S_{i}$ and $S_{j}$ decays rapidly as $|i-j|$ increases, then the effect of the dependence is negligible.

To go beyond the usual mixing conditions, a particularly useful property of the kernel MMD distance is that the second moment of $\gamma_{k}(P,P_{n})$ depends only on the pairwise correlation in the samples. This perspective is useful for settings where there may not be a natural notion of time.

To make things precise, suppose that $S_{1},\dots,S_{n}\in\mathbb{R}^{d}$ is a collection of (possibly dependent) samples with identical distribution $P$. For each $i\neq j$ let $Q_{ij}$ denote the law of $(S_{i},S_{j})$. Starting with (5), the expectation of the squared MMD can now be decomposed as

$\displaystyle\mathbb{E}\mathopen{}\mathclose{{}\left[\gamma^{2}_{k}(P,P_{n})}\right]$

$\displaystyle=\frac{1}{n}\mathbb{E}\mathopen{}\mathclose{{}\left[\gamma^{2}_{k}(P,P_{1})}\right]+\frac{1}{n^{2}}\sum_{i\neq j}r_{ij},$

(23)

where $r_{ij}\coloneqq\mathbb{E}_{Q_{ij}}[k(S_{i},S_{j})]-\mathbb{E}_{P\otimes P}[k(S_{i},S_{j}^{\prime})]$. Notice that the first term in this decomposition is the second moment under independent samples. The second term is non-negative and depends only on the pairwise dependence, i.e., the difference between $Q_{ij}$ and the probability measure obtained by the product of its marginals.

In some cases, it is natural to argue that only a small number of the terms $r_{ij}$ are nonzero. More generally it is desirable to provide guarantees in terms of measures of dependence that do not rely on the particular choice of kernel. The next result provides such a bound in terms of the Hellinger distance.

The following example is inspired by random feature kernel interpretation of neural networks  Rahimi and Recht (2008).

The sub-gamma condition allow us to address distributions that do not satisfy the sub-Gaussian condition. Upper bounds on the convergence rate for this class of distributions follow from Theorem 3 combined with Theorem 6.

For sub-gamma $X$ as in Example 1, $\mathbb{E}\mathopen{}\mathclose{{}\left[k(X,X)}\right]$ has upper bound (UB) and lower bound (LB) as in (20) and (22) respectively. Here, in addition to the theoretical UB and LB as shown in Figure 1, we also compute the estimate of $\mathbb{E}\mathopen{}\mathclose{{}\left[n\gamma_{k}^{2}(P,P_{n})}\right]$ approximated by the two-sample estimator. Let $P_{n}$ and $P_{n}^{\prime}$ be the empirical measure defined as in (13) for independent iid copies of the sub-gamma random vector as in Example 1, $\{X_{i}\}_{i=1}^{n}$ and $\{X_{i}^{\prime}\}_{i=1}^{n}$, respectively. Define $\hat{\Delta}_{\gamma}^{2}\coloneqq\frac{n}{2}\gamma_{k}^{2}({P}_{n},{P}_{n}^{\prime})$, and then $\mathbb{E}[\hat{\Delta}_{\gamma}^{2}]=\mathbb{E}\mathopen{}\mathclose{{}\left[k(X,X)}\right]-\mathbb{E}\mathopen{}\mathclose{{}\left[k(X,X^{\prime})}\right]=\mathbb{E}\mathopen{}\mathclose{{}\left[n\gamma_{k}^{2}(P,P_{n})}\right]$. Because the empirical kernel MMD (squared) distance $\gamma_{k}^{2}(P_{n},P_{n}^{\prime})$ can be computed numerically from the two samples of sub-gamma random vectors, we can estimate the expectation of $\hat{\Delta}_{\gamma}^{2}$ by empirical average over Monte-Carlo replicas. Detailed numerical techniques to compute the two-moment kernel and experimental setup are provided in the appendix.

The results are shown in Figure 2, where we set the parameter $b$ controlling the shape and scale of $X$ as $b=d^{-\delta}$, $\delta=\{0,0.25,0.5,0.75\}$. The data dimension $d$ takes multiple values, and the kernel bandwidth $\sigma$ takes values 1 and 4. In both cases, the left plot shows the same information as Figure 1 in view of increasing $d$, so as to be compared to the right plot. The right plot focuses on the case of small $d$, where the empirical estimates of $\mathbb{E}[\hat{\Delta}_{\gamma}^{2}]$ observe the UB. Notably, these values also lie between the UB and LB (recall that the LB applies to $\mathbb{E}\mathopen{}\mathclose{{}\left[k(X,X)}\right]$ but not $\mathbb{E}[\hat{\Delta}_{\gamma}^{2}]$) for both cases of $\sigma$, and approach the LB when $\sigma=1$. This shows that our theoretical UB captures an important component of the kernel MMD distances for this example of $X$.

We generate dependent samples following a Gaussian process $\{Z(\alpha)\}_{\alpha\in S^{N}}$ as in Example 2 in Section 4.3, and numerically compute the values of $\gamma_{k}^{2}(P,P_{n})$ by its two-sample version $\gamma_{k}^{2}(P_{n},{P}_{n}^{\prime})$, using the same kernel as in the first experiment. Theoretically, when $n$ is large, $\gamma_{k}^{2}(P,P_{n})$ is expected to concentrate at its mean value as in (24), and then since $\alpha_{i}$ are uniformly sampled on the $N$-sphere, we also expect concentration at the value of (25).

We set $n=\{30,50,70,100\}$, and vary $N$ from 5 to 100. The data are in dimension $d=5$, the kernel parameters are $\sigma=0.5$, $c=1$, $p=1$, and for each value of $n$ and $N$, $\gamma_{k}^{2}(P_{n},P_{n}^{\prime})$ are computed for 100 realizations of the dependent random variables $S_{i}$, conditioned on one realization of the vectors $\alpha_{i}$’s. The results are shown in Figure 3, where for the small values of $N$, the computed approximate values of $\mathbb{E}\mathopen{}\mathclose{{}\left[\gamma_{k}^{2}(P,P_{n})}\right]$ are decreasing because they are dominated by the contribution from the dependence in the samples, namely the part that is upper-bounded by the second term depending on $\alpha$ in (24). As $N$ increases, these values converge to certain positive value which decrease as $n$ increases.

Specifically, we take the average over the mean values of $\gamma_{k}^{2}(P_{n},P_{n}^{\prime})$ for $N\in[80,100]$ as an estimate of the limiting values, and Figure 3 (right) shows that these values decay as $n^{-1}$, as they correspond to the first term in (24). These numerical results show the competing two factors as predicted by the analysis in Section 4.3.