Estimation of Wasserstein distances in the Spiked Transport Model

Contibutions:

• Estimation of Wasserstein distance in HD for distributions that differ only in a LD subspace.
• Lower bounds.
• A computational-statistical gap.

Introduction:

Spiked Transport Model:

• Fix $\mathcal{U}\subseteq \mathbb{R}^d$ of dim. $k \ll d$.
• R.v. $X^{(1)}, X^{(2)}\in \mathcal{U}$ with arbitrary distributions.
• R.v. $Z\perp (X^{(1)}, X^{(2)})$ supported on $\mathcal{U}^{\perp}$.

Spiked Transport Model:

Question: Given $n$ i.i.d. observations from both $\mu^{(1)}$ and $\mu^{(2)}$ is it possible to estimate $W_p(\mu^{(1)}, \mu^{(2)})$ at a rate faster than $n^{-1/d}$?

• Answer is yes, but we need smoothness and decay assumptions on the measures.

Concentration Assumptions:

A prob. meas. $\mu$ on $\mathbb{R}^d$ satisfy the $T_p(\sigma^2)$ transport inequality if

Wasserstein Projection Pursuit:

• Consider the set $\mathcal{V}_k(\mathbb{R}^d)$ of $k\times d$ matrices with orthonormal rows.
• For $\mu \in \text{Prob}(\mathbb{R}^d)$ and $U\in \mathcal{V}_k(\mathbb{R}^d)$, define $\mu_U$ as the law of $U.Y$ for $Y\sim\mu$.

Wasserstein Projection Pursuit:

Estimation:

Theorem 1: Let $(\mu^{(1)}, \mu^{(2)})$  sats. the Spiked Transport Model (STM). For any $p\in[1,2]$, if $\mu^{(1)}$ and $\mu^{(2)}$ sats. the $T_p(\sigma^2)$ ineq., then

Before proving:

Prop. 3: Under the STM

Proof for estimation:

First notice that

$$\mathbb{E} \left| \widehat W_{p,k} - W_p(\mu^{(1)}, \mu^{(2)}) \right| \leq \mathbb{E} \tilde W_{p,k} (\mu^{(1)}, \mu^{(1)}_n) + \mathbb{E} \tilde W_{p,k} (\mu^{(2)}, \mu^{(2)}_n)$$ 

Then we can focus on bounding $\mathbb{E} \tilde W_{p,k} (\mu, \mu_n)$.

• We assume w.l.g. that $\mu$ has mean 0, and $\sigma=1$.
• Consider $Z_U := W_p(\mu_U, (\mu_n)_U)$

Proof for estimation:

Lemma:  $\exists$ r.v. $L$ s.t. for all $U,V\in \mathcal{V}_k(\mathbb{R}^d)$

$$|Z_U - Z_V| \leq L\|U-V\|_{op}$$

and $\mathbb{E} L \lesssim \sqrt{dp}$

proof: Let $X\sim \mu$, then

Proof for estimation:

• The process $Z_U$ is Lipschitz,
• By Thm. 6 $Z_U$ is $n^{-1}$-subgaussian.

Proof for estimation:

There exist a univ. const. $c$ such that $\mathcal{N}(\mathcal{V}_k, \varepsilon, \|\cdot\|_{op}) \leq dk \log{\frac{c\sqrt{k}}{\varepsilon}}$ for $\varepsilon\in (0,1]$. Choosing $\varepsilon = \sqrt{k/n}$ yelds

Proof for estimation:

Finally we get,

Spike estimation:

Theorem 10: Let $p\in [1,2]$. Assume $(\mu^{(1)}, \mu^{(2)})$  sats. STM and the $T_p(\sigma^2)$ ineq. Let \(\hat \mathcal{U} := \text{span}(\hat U)\), where

Then

Lower Bounds:

Consider a compact metric space $\mathcal{X}$ s.t.

$$c \varepsilon^{-d} \leq \mathcal{N}(\mathcal{X}, \varepsilon) \leq C \varepsilon^{-d}$$

for all $\varepsilon \leq \text{diam}(\mathcal{X})$.

With $\mathcal{P} = \text{Prob}(\mathcal{\mathcal{X}})$ define

$$R(n, \mathcal{P}) := \underset{\hat W}{\inf}\underset{\mu, \nu\in\mathcal{P}}{\sup} \mathbb{E}_{\mu, \nu}|\hat W - W_p(\mu,\nu)|.$$

Lower Bounds:

Theorem 11: Let $d>2p>2$ and $\mathcal{X}$ with the cov. number as before. Then

$$R(n,\mathcal{P}) \geq C_{d,p} (d \log{n})^{-1/d}$$

Lower Bounds:

Prop. 9: Assume $d>2p>2$ and let $m$ be a pos. integer. Let $u=\text{Unif}([m])$. $\exists$ a random function $F:[m]\rightarrow \mathcal{X}$ s.t. for any dist. $q$ on $[m]$,

with prob. at least $.9$

Lower Bounds:

proof: For the lower bound:

Use the cov. num. to get a set of points $G_m = \{x_i\}_{i\in[m]}$ s.t. $d(x_i, x_j) \gtrsim m^{-1/d}$. Then select $F$ unif. at random from the set of bijections from $[m]$ to $G_m$. Now, since

$$d(x,y)^p \gtrsim m^{-p/d} \boldsymbol{1}\{x\neq y\}$$

we can com get for any coupling $\pi$ of $F_\#q$ and $F_\#u$

$$\int d(x,y)^p d\pi(x,y) \gtrsim m^{-p/d}\mathbb{P}[X\neq Y]\geq m^{-p/d} d_{TV}(q, u)$$

Lower Bounds:

For the upper bound: ...

Lower Bounds:

Prop. 10: Fix $n\in\mathbb{N}$ and a cnt. $\delta\in [0, .1]$. Given $m\in\mathbb{N}$, let $D_m$ be the set of prob. $q$ in $[m]$ sats. $\chi^2(q, u)\leq 9$. Denote by $D_{m,\delta}^-$ the subset of $D_m$ where $d_{TV}(q,u)\leq \delta$ and by $D_m^+$ the subset sats. $d_{TV}(q,u)\geq 1/4$. If $m = \lceil C\delta^{-1}n\log{n} \rceil$ for a sufficiently large univ. $C$ and $n$ is sufficiently large, then

$$\underset{\psi}{\inf} \left\{ \underset{q\in D_m^+}{\sup} \mathbb{P}_q [\psi=1] + \underset{q\in D_{m,\delta}^-}{\sup} \mathbb{P}_q [\psi=0] \right\} \geq .9$$

where the $\inf$ is taken over all test based on $n$ samples.

Lower Bounds:

Proof of Thm. 11:

Let $A = \{|\hat W - W_p(F_\#q, F_\#u)|\geq \Delta_d\}$, with $\Delta_d = \frac{1}{16} c^* m^{-1/d}$. Then

We def. the rand. test

$$\psi(X_1,\dots,X_n) := \boldsymbol{1}\{\hat W(F(X_1),\dots,F(X_n);  F(Y_1),\dots,F(Y_n)) \leq 2\Delta_d\}$$

Lower Bounds:

Choosing $m = \lceil C\delta^{-1}n\log{n} \rceil$ for suff. large C and apply prop. 10 gives $\underset{\mu, \nu\in\mathcal{P}}{\sup} \mathbb{P}_{\mu, \nu}\left[|\hat W - W_p(\mu,\nu)|\geq \Delta_d\right] \geq .8$, and Markov's ineq. yields the claim.

Computational-Statistical Gap:

Def. 5: Given a dist. $\mathcal{D}$ on $\mathbb{R}^d$, for any sample size $t>0$ and $f:\mathbb{R}^d\rightarrow [0,1]$, the oracle $\text{VSTAT}(t)$ returns a value  $v\in [p-\tau, p+\tau]$, where $p = \mathbb{E}f(X)$ and $\tau = \frac{1}{t}\vee \sqrt{\frac{p(1-p)}{t}}$

Computational-Statistical Gap:

Theorem 12: There exists a pos. univ. constant $c$ s.t., for any $d$, estimating $W_1(\mu^{(1)}, \mu^{(2)})$ for $\mu^{(1)}$, $\mu^{(2)}$ sats. the STM with $k=1$ to accuracy $\Theta(1/\sqrt{d})$ with prob. at least $2/3$ requires at least $2^{cd}$ queries to $\text{VSTAT}(2^{cd})$.