SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We thank anonymous reviewers for feedback on the manuscript. Yibo Yang acknowledges support from the Hasso Plattner Foundation. Marcel Nutz acknowledges support from NSF Grants DMS-1812661, DMS-2106056. Stephan Mandt acknowledges support from the National Science Foundation (NSF) under the NSF CAREER Award 2047418; NSF Grants 2003237 and 2007719, the Department of Energy, Office of Science under grant DE-SC0022331, the IARPA WRIVA program, as well as gifts from Intel, Disney, and Qualcomm.
Supplementary Material for Estimating the Rate-Distortion Function by Wasserstein Gradient Descent
We review probability theory background and explain our notation from the main text in Section 7, give the formulas we used for numerically estimating an  $R(D)$  upper bound in Section 8, provide additional discussions and proofs regarding Wasserstein gradient descent in Section 9, elaborate on the connections between the R-D estimation problem and variational inference/learning in Section 10, provide additional experimental results and details in Section 11, and list an example implementation of WGD in Section 12. Our code and can be found at https://github.com/yiboyang/wgd.
7 Notions from probability theory
In this section we collect notions of probability theory used in the main text. See, e.g., [Cinlar, 2011] or [Folland, 1999] for more background.
Marginal and conditional distributions. The source and reproduction spaces  $\mathcal{X},\mathcal{Y}$  are equipped with sigma-algebras  $\mathcal{A}_{\mathcal{X}}$  and  $\mathcal{A}_{\mathcal{Y}}$ , respectively. Let  $\mathcal{X}\times \mathcal{Y}$  denote the product space equipped with the product sigma algebra  $\mathcal{A}_{\mathcal{X}}\otimes \mathcal{A}_{\mathcal{Y}}$ . For any probability measure  $\pi$  on  $\mathcal{X}\times \mathcal{Y}$ , its first marginal is
$$
\pi_ {1} (A) := \pi (A \times \mathcal {Y}), \quad A \in \mathcal {A} _ {\mathcal {X}},
$$
which is a probability measure on  $\mathcal{X}$ . When  $\pi$  is the distribution of a random vector  $(X,Y)$ , then  $\pi_1$  is the distribution of  $X$ . The second marginal of  $\pi$  is defined analogously as
$$
\pi_ {2} (B) := \pi (\mathcal {X} \times B), \quad B \in \mathcal {A} _ {\mathcal {Y}}.
$$
A Markov kernel or conditional distribution  $K(x, dy)$  is a map  $\mathcal{X} \times \mathcal{A}_Y \to [0,1]$  such that
1.  $K(x,\cdot)$  is a probability measure on  $\mathcal{V}$  for each  $x\in \mathcal{X}$
2. the function  $x\mapsto K(x,B)$  is measurable for each set  $B\in \mathcal{A}_{\mathcal{Y}}$
When speaking of the conditional distribution of a random variable  $Y$  given another random variable  $X$ , we occasionally also use the notation  $Q_{Y|X}$  from information theory [Polyanskiy and Wu, 2014]. Then,  $Q_{Y|X = x}(B) = K(x,B)$  is the conditional probability of the event  $\{Y\in B\}$  given  $X = x$ .
Suppose that a probability measure  $\mu$  on  $\mathcal{X}$  is given, in addition to a kernel  $K(x,dy)$ . Together they define a unique measure  $\mu \otimes K$  on the product space  $\mathcal{X} \times \mathcal{Y}$ . For a rectangle set  $A \times B \in \mathcal{A}_{\mathcal{X}} \otimes \mathcal{A}_{\mathcal{Y}}$ ,
$$
\mu \otimes K (A \times B) = \int_ {A} \mu (d x) K (x, B), \quad A \in \mathcal {A} _ {\mathcal {X}}, B \in \mathcal {A} _ {\mathcal {Y}}.
$$
The measure  $\pi \coloneqq \mu \otimes K$  has first marginal  $\pi_1 = \mu$ .
The classic product measure is a special case of this construction. Namely, when a measure  $\nu$  on  $\mathcal{V}$  is given, using the constant kernel  $K(x,dy)\coloneqq \nu (dy)$  (which does not depend on  $x$ ) gives rise to the product measure  $\mu \otimes \nu$
$$
\mu \otimes \nu (A \times B) = \mu (A) \nu (B), \quad A \in \mathcal {A} _ {\mathcal {X}}, B \in \mathcal {A} _ {\mathcal {Y}}.
$$
Under mild conditions (for instance when  $\mathcal{X},\mathcal{Y}$  are Polish spaces equipped with their Borel sigma algebras, as in the main text), any probability measure  $\pi$  on  $\mathcal{X}\times \mathcal{Y}$  is of the above form. Namely, the disintegration theorem asserts that  $\pi$  can be written as  $\pi = \pi_1\otimes K$  for some kernel  $K$ . When  $\pi$  is the joint distribution of a random vector  $(X,Y)$ , this says that there is a measurable version of the conditional distribution  $Q_{Y|X}$ .
Optimal transport. Given a measure  $\mu$  on  $\mathcal{X}$  and a measurable function  $T: \mathcal{X} \to \mathcal{Y}$ , the pushforward (or image measure) of  $\mu$  under  $T$  is a measure on  $\mathcal{Y}$ , given by
$$
T _ {\#} \mu (B) = \mu (T ^ {- 1} (B)), \quad B \in \mathcal {A} _ {y}.
$$
If  $T$  is seen as a random variable and  $\mu$  as the baseline probability measure, then  $T_{\#} \mu$  is simply the distribution of  $T$ .
Suppose that  $\mu$  and  $\nu$  are probability measures on  $\mathcal{X} = \mathcal{Y} = \mathbb{R}^d$  with finite second moment. As introduced in the main text,  $\Pi (\mu ,\nu)$  denotes the set of couplings, i.e., measures  $\pi$  on  $\mathcal{X}\times \mathcal{Y}$  with  $\pi_1 = \mu$  and  $\pi_{2} = \nu$ . The 2-Wasserstein distance  $W_{2}(\mu ,\nu)$  between  $\mu$  and  $\nu$  is defined as
$$
W_{2}(\mu ,\nu) = \left(\inf_{\pi \in \Pi (\mu ,\nu)}\int \| y - x\|^{2}\pi (dx,dy)\right)^{1 / 2}.
$$
This indeed defines a metric on the space of probability measures with finite second moment.
We finish this section by giving a proof of the basic equivalence between optimization of  $\mathcal{L}_{EOT}$  and  $\mathcal{L}_{BA}$ , which goes back to [Csiszár, 1974a, Lemma 1.3 and subsequent discussion]. Recall that  $\Pi (\mu ,\cdot)$  denotes the set of measures on the product space  $\mathcal{X}\times \mathcal{Y}$  with  $\pi_1 = \mu$
Lemma 7.1. Set  $\epsilon = \lambda^{-1}$ . It holds that
$$
\inf_{\nu \in \mathcal{P}(\mathcal{Y})}\mathcal{L}_{EOT}(\nu) = \inf_{\nu \in \mathcal{P}(\mathcal{Y})}\lambda^{-1}\mathcal{L}_{BA}(\nu)\hskip 28.452756ptand\hskip 28.452756pt\underset {\nu \in \mathcal{P}(\mathcal{Y})}{\arg \min}\mathcal{L}_{EOT}(\nu) = \underset {\nu \in \mathcal{P}(\mathcal{Y})}{\arg \min}\mathcal{L}_{BA}(\nu).
$$
Proof. Both statements will follow from a simple property of relative entropy [Polyanskiy and Wu, 2022, Theorem 4.1, "golden formula": For  $\pi \in \Pi(\mu, \cdot)$  with  $H(\pi_2 \mid \nu) < \infty$ , the properties of the logarithm reveal
$$
H (\pi \mid \mu \otimes \nu) = H (\pi \mid \mu \otimes \pi_ {2}) + H (\pi_ {2} \mid \nu), \tag {15}
$$
and hence  $H(\pi \mid \mu \otimes \nu)\geq H(\pi \mid \mu \otimes \pi_2)$  . This implies
$$
\begin{array}{l} \inf  _ {\nu \in \mathcal {P} (\mathcal {Y})} \mathcal {L} _ {E O T} (\nu) = \inf  _ {\pi \in \Pi (\mu , \cdot)} \int \rho d \pi + \epsilon H (\pi | \mu \otimes \pi_ {2}) \\ = \inf  _ {\pi \in \Pi (\mu , \cdot)} \inf  _ {\nu \in \mathcal {P} (\mathcal {Y})} \int \rho d \pi + \epsilon H (\pi | \mu \otimes \nu) \\ = \frac{1}{\lambda}\inf_{\nu \in \mathcal{P}(\mathcal{Y})}\mathcal{L}_{BA}(\nu). \\ \end{array}
$$
Further, any optimizer  $\nu^{*}$  of  $\mathcal{L}_{BA}$  with corresponding optimal "coupling"  $\pi^{*} \in \Pi(\mu, \cdot)$  must satisfy  $H(\pi_2^* \mid \nu^*) = 0$  (otherwise, taking  $\nu^{*} = \pi_{2}^{*}$  has better objective) and thus  $\pi_2^* = \nu^*$  and  $\pi^{*} \in \Pi(\mu, \nu^{*})$ , therefore  $\nu^{*}$  is also an optimizer of  $\mathcal{L}_{EOT}$  by the above equality. Conversely, any optimizer  $\nu^{*}$  of  $\mathcal{L}_{EOT}$  with coupling  $\pi^{*} \in \Pi(\mu, \nu^{*}) \subset \Pi(\mu, \cdot)$  clearly yields a feasible solution for  $\mathcal{L}_{BA}$  as well, and hence is also an optimizer of  $\mathcal{L}_{BA}$  by the above equality.
8 Numerical estimation of rate and distortion from a reproduction distribution
Given a reproduction distribution  $\nu$  and a kernel  $\mathcal{X} \times \mathcal{A}_{\mathcal{Y}} \to [0,1]$ , the tuple  $(\mathcal{D},\mathcal{R}) \in \mathbb{R}^2$  defined by
$$
\mathcal {D} := \int \rho d (\mu \otimes K)
$$
$$
\mathcal {R} := H (\mu \otimes K | \mu \otimes \nu)
$$
lies above the  $R(D)$  curve. This again follows from the variational inequality (15)  $\mathcal{R} \geq H(\pi|\mu \times (\mu \times K)_2) = I(\mu \times K)$ , so that  $\mathcal{R} \geq R(\mathcal{D}) \eqqcolon \inf_{\pi \in \Pi(\pi,\cdot):\pi(\rho) \leq \mathcal{D}} I(\pi)$ .
Given only a reproduction distribution  $\nu$ , we will construct a kernel  $K$  from  $\nu$  and use  $(\mu, K)$  to compute  $(\mathcal{D}, \mathcal{R})$ , letting
$$
\frac {d K (x , \cdot)}{d \nu} (y) = \frac {e ^ {- \lambda \rho (x , y)}}{\int e ^ {- \lambda \rho (x , \tilde {y})} \nu (d \tilde {y})}
$$
as in the first step of the BA algorithm. This choice of  $K$  is in fact optimal as it achieves the minimum in the definition of  $\mathcal{L}_{BA}$  (5) for the given  $\nu$ . Plugging  $K$  into the formulas for  $\mathcal{D}$  and  $\mathcal{R}$  gives
$$
\mathcal {D} = \int \rho (x, y) e ^ {\varphi (x) - \lambda \rho (x, y)} \mu \otimes \nu (d x, d y) \tag {16}
$$
$$
\begin{array}{l} \mathcal {R} = \iint \log \left(\frac {K (x , d y)}{\nu (d y)}\right) K (x, d y) \mu (d x) \tag {17} \\ = \int [ \varphi (x) - \lambda \rho (x, y) ] e ^ {\varphi (x) - \lambda \rho (x, y)} \mu \otimes \nu (d x, d y), \\ \end{array}
$$
where we introduced the shorthand
$$
\varphi (x) := - \log \int_ {\mathcal {Y}} e ^ {- \lambda \rho (x, y)} \nu (d y). \tag {18}
$$
Note that we have the following relation (which explains (6))
$$
\mathcal {L} _ {B A} (\nu) = \mathcal {R} + \lambda \mathcal {D} = \int \varphi (x) \mu (d x).
$$
Let  $\nu$  be an  $n$ -point measure,  $\nu = \sum_{i=1}^{n} w_i \delta_{y_i}$ , e.g., the output of WGD or in the inner step of NERD. Then  $\varphi(x)$  can be evaluated exactly as a finite sum over the  $\nu$  particles, and the expressions above for  $\mathcal{D}$  and  $\mathcal{R}$  (which are integrals w.r.t. the product distribution  $\mu \otimes \nu$ ) can be estimated as sample averages. That is, given  $m$  independent samples  $\{x_i\}_{i=1}^m$  from  $\mu$ , we compute unbiased estimates
$$
\hat {\mathcal {D}} = \sum_ {i = 1} ^ {m} \sum_ {j = 1} ^ {n} \frac {1}{m} w _ {j} \rho (x _ {i}, y _ {j}) e ^ {\varphi (x _ {i}) - \lambda \rho (x _ {i}, y _ {j})},
$$
$$
\hat {\mathcal {R}} = \sum_ {i = 1} ^ {m} \sum_ {j = 1} ^ {n} \frac {1}{m} w _ {j} [ \varphi (x _ {i}) - \lambda \rho (x _ {i}, y _ {j}) ] e ^ {\varphi (x _ {i}) - \lambda \rho (x _ {i}, y _ {j})},
$$
where  $\varphi(x) = -\log \sum_{j=1}^{n} e^{\varphi(x_i) - \lambda \rho(x_i, y_j)}$ . Similarly, a sample mean estimate for  $\mathcal{L}_{BA}$  is given by
$$
\hat {\mathcal {L}} _ {B A} (\nu) = \frac {1}{m} \sum_ {i = 1} ^ {m} \varphi \left(x _ {i}\right) = - \frac {1}{m} \sum_ {i = 1} ^ {m} \log \left(\sum_ {j = 1} ^ {n} e ^ {\varphi \left(x _ {i}\right) - \lambda \rho \left(x _ {i}, y _ {j}\right)}\right). \tag {19}
$$
In practice, we found it simpler and numerically more stable to instead compute  $\hat{\mathcal{R}}$  as
$$
\hat {\mathcal {R}} = \hat {\mathcal {L}} _ {B A} (\nu) - \lambda \hat {\mathcal {D}}.
$$
Whenever possible, we avoid exponentiation and instead use logsumexp to prevent numerical issues.
9 Wasserstein gradient descent
9.1 On Wasserstein gradients of the EOT and rate functionals
First, we elaborate on the Wasserstein gradient of the EOT functional  $\mathcal{L}_{EOT}(\nu)$ . That the dual potential from Sinkhorn's algorithm is differentiable follows from the fact that optimal dual potentials satisfy the Schrödinger equations (cf. [Nutz, 2021, Corollary 2.5]). Differentiability was shown in [Genevay et al., 2019, Theorem 2] in the compact setting, and in [Mena and Niles-Weed, 2019, Proposition 1] in unbounded settings. While Mena and Niles-Weed [2019] only states the result for quadratic cost, the approach of Proposition 1 therein applies more generally.
Below, we compute the Wasserstein gradient of the rate functional  $\mathcal{L}_{BA}(\nu)$ . Recall from (6),
$$
\mathcal {L} _ {B A} (\nu) = \int - \log \int \exp (- \lambda \rho (x, y)) \nu (d y) \mu (d x).
$$
Under sufficient integrability on  $\mu$  and  $\nu$  to exchange the order of limit and integral, we can calculate the first variation as
$$
\begin{array}{l} \lim _ {\varepsilon \to 0} \frac {\mathcal {L} ((1 - \varepsilon) \nu + \varepsilon \tilde {\nu}) - \mathcal {L} (\nu)}{\varepsilon} = - \int \lim _ {\varepsilon \to 0} \frac {1}{\varepsilon} \log \left[ \frac {\int \exp (- \lambda \rho (x , y)) (\nu + \varepsilon (\tilde {\nu} - \nu)) (d y)}{\int \exp (- \lambda \rho (x , y)) \nu (d y)} \right] \mu (d x) \\ = - \int_ {\varepsilon \rightarrow 0} \frac {1}{\varepsilon} \log \left[ 1 + \frac {\int \exp (- \lambda \rho (x , y)) \varepsilon (\tilde {\nu} - \nu) (d y)}{\int \exp (- \lambda \rho (x , y)) \nu (d y)} \right] \mu (d x) \\ = \iint - \frac {\exp (- \lambda \rho (x , y))}{\int \exp (- \lambda \rho (x , \tilde {y})) \nu (d \tilde {y})} \mu (d x) (\tilde {\nu} - \nu) (d y), \\ \end{array}
$$
where the last equality uses  $\lim_{\varepsilon \to 0} \frac{1}{\varepsilon} \log (1 + \varepsilon x) = x$  and Fubini's theorem. Thus the first variation  $\psi^\nu$  of  $\mathcal{L}_{BA}$  at  $\nu$  is
$$
\psi^ {\nu} (y) = \int - \frac {\exp (- \lambda \rho (x , y))}{\int \exp (- \lambda \rho (x , \tilde {y})) \nu (d \tilde {y})} \mu (d x). \tag {20}
$$
To find the desired Wasserstein gradient of  $\mathcal{L}_{BA}$ , it remains to take the Euclidean gradient of  $\psi^{\nu}$ , i.e.,  $\nabla \mathcal{L}_{BA}(\nu) = \nabla \psi^{\nu}$ . Doing so gives us the desired Wasserstein gradient:
$$
\begin{array}{l} \nabla V _ {\mathcal {L} _ {B A}} (\nu) [ y ] = \nabla_ {y} \psi^ {\nu} (y) \\ = \frac {\partial}{\partial y} \int - \frac {\exp (- \lambda \rho (x , y))}{\int \exp (- \lambda \rho (x , \tilde {y})) \nu (d \tilde {y})} \mu (d x) \\ = \int \frac {\exp (- \lambda \rho (x , y)) \lambda \frac {\partial}{\partial y} \rho (x , y)}{\int \exp (- \lambda \rho (x , \tilde {y})) \nu (d \tilde {y})} \mu (d x), \tag {21} \\ \end{array}
$$
again assuming suitable regularity conditions on  $\rho$  and  $\mu$  to exchange the order of integral and differentiation.
9.2 Proof of Proposition 4.2 (convergence of Wasserstein gradient descent)
We first provide an auxiliary result.
Lemma 9.1. Let  $\gamma_1 \geq \gamma_2 \geq \dots \geq 0$  and  $a_t \geq 0$ ,  $t \in \mathbb{N}$ ,  $C > 0$  satisfy  $\sum_{t=1}^{\infty} \gamma_t = \infty$ ,  $\sum_{t=1}^{\infty} \gamma_t^2 < \infty$ ,  $\sum_{t=1}^{\infty} a_t \gamma_t < \infty$  and  $|a_t - a_{t+1}| \leq C \gamma_t$  for all  $t \in \mathbb{N}$ . Then  $\lim_{t \to \infty} a_t = 0$ .
Proof. The conclusion remains unchanged when rescaling  $a_{t}$  by the constant  $C$ , and thus without loss of generality  $C = 1$ .
Clearly  $\gamma_t \to 0$  as  $\sum_{t=1}^{\infty} \gamma_t^2 < \infty$ . Moreover, there exists a subsequence of  $(a_t)_{t \in \mathbb{N}}$  which converges to zero (otherwise there exists  $\delta > 0$  such that  $a_t \geq \delta > 0$  for all but finitely many  $t$ , contradicting  $\sum_{t=1}^{\infty} \gamma_t a_t < \infty$ ).
Arguing by contradiction, suppose that the conclusion fails, i.e., that there exists a subsequence of  $(a_{t})_{t\in \mathbb{N}}$  which is uniformly bounded away from zero, say  $a_{t}\geq \delta >0$  along that subsequence. Using this subsequence and the convergent subsequence mentioned above, we can construct a subsequence  $a_{i_1},a_{i_2},a_{i_3},\ldots$  where  $a_{i_n}\approx 0$  for  $n$  odd and  $a_{i_n}\geq \delta$  for  $n$  even. We will show that
$$
\sum_ {t = i _ {2 n - 1}} ^ {i _ {2 n}} a _ {t} \gamma_ {t} \gtrsim \delta^ {2} / 2 \qquad \text {f o r a l l} n \in \mathbb {N},
$$
contradicting the finiteness of  $\sum_{t}\gamma_{t}a_{t}$ . (The notation  $\approx (\gnsim)$  indicates (in)equality up to additive terms converging to zero for  $n\to \infty$ .)
To ease notation, fix  $n$  and set  $m = i_{2n - 1}$  and  $M = i_{2n}$ . We show that  $\sum_{t = m}^{M}a_{t}\gamma_{t}\gtrsim \delta^{2} / 2$ . To this end, using  $|a_{t} - a_{t + 1}|\leq \gamma_{t}$  we find
$$
a _ {t} \geq a _ {M} - \sum_ {j = k} ^ {M - 1} \gamma_ {j} \geq \delta - \sum_ {j = k} ^ {M - 1} \gamma_ {j}.
$$
Since  $a_{m} \approx 0$ , there exists a largest  $n_0 \in \mathbb{N}$ ,  $n_0 \geq m$ , such that  $\sum_{j=n_0}^{M-1} \gamma_j \gtrsim \delta$  (and thus  $\sum_{j=n_0}^{M-1} \gamma_j \lesssim \delta - \gamma_{n_0} \approx \delta$  as well). We conclude
$$
\begin{array}{l} \sum_ {t = m} ^ {M} \gamma_ {t} a _ {t} \geq \sum_ {t = n _ {0}} ^ {M} \gamma_ {t} a _ {t} \geq \sum_ {t = n _ {0}} ^ {M} \gamma_ {t} \left(\delta - \sum_ {j = k} ^ {M - 1} \gamma_ {j}\right) \gtrsim \delta^ {2} - \sum_ {t = n _ {0}} ^ {M} \sum_ {j = n _ {0}} ^ {M} \gamma_ {t} \gamma_ {j} \mathbf {1} _ {\{j \geq k \}} \\ = \delta^ {2} - \frac {1}{2} \left(\sum_ {t = n _ {0}} ^ {M} \gamma_ {t}\right) ^ {2} - \frac {1}{2} \sum_ {t = n _ {0}} ^ {M} \gamma_ {t} ^ {2} \approx \delta^ {2} / 2, \\ \end{array}
$$
where we used that  $\sum_{t=n_0}^{M} \gamma_t^2 \approx 0$ . This completes the proof.
Proof of Proposition 4.2. Using the linear approximation property in (14), we calculate
$$
\begin{array}{l} \mathcal {L} \left(\nu^ {(n)}\right) - \mathcal {L} \left(\nu^ {(0)}\right) = \sum_ {t = 0} ^ {n - 1} \mathcal {L} \left(\nu^ {(t + 1)}\right) - \mathcal {L} \left(\nu^ {(t)}\right) \\ = \sum_ {t = 0} ^ {n - 1} - \gamma_ {t} \int \| \nabla V _ {\mathcal {L}} (\nu^ {(t)}) \| ^ {2} d \nu^ {(t)} + \gamma_ {t} ^ {2} o \left(\int \| \nabla V _ {\mathcal {L}} (\nu^ {(t)}) \| ^ {2} d \nu^ {(t)}\right). \\ \end{array}
$$
As  $\mathcal{L}(\nu^{(0)})$  is finite and  $\mathcal{L}(\nu^{(n)})$  is bounded from below, it follows that
$$
\sum_ {t = 0} ^ {\infty} \gamma_ {t} \int \| \nabla V _ {\mathcal {L}} (\nu^ {(t)}) \| ^ {2} d \nu^ {(t)} <   \infty .
$$
The claim now follows by applying Lemma 9.1 with  $a_{t} = \int \| \nabla \psi^{\nu^{(t)}}\|^{2}d\nu^{(t)}$ ; note that the assumption in the lemma,  $|a_{t} - a_{t + 1}|\leq C\gamma_{t}$ , is satisfied due to the second inequality in (14) and a short calculation below
$$
\begin{array}{l} \left| \int \| \nabla \psi^ {\nu^ {(t)}} \| ^ {2} - \int \| \nabla \psi^ {\nu^ {(t + 1)}} \| ^ {2} \right| \leq C W _ {2} (\nu^ {(t)}, \nu^ {(t + 1)}) \\ \leq C \left(\iint \| x - y \| ^ {2} \delta_ {x - \gamma_ {t} \nabla V _ {\mathcal {L}} (\nu^ {(t)}) [ x ]} (d y) \nu^ {(t)} (d x)\right) ^ {\frac {1}{2}} \\ \leq C \gamma_ {t} \left(\int \| \nabla V _ {\mathcal {L}} (\nu^ {(t)}) \| ^ {2} \nu^ {(t)} (d x)\right) ^ {\frac {1}{2}} \\ \leq C \gamma_ {t} \sup  _ {t ^ {\prime}} \left(\int \| \nabla V _ {\mathcal {L}} \left(\nu^ {\left(t ^ {\prime}\right)}\right) \| ^ {2} \nu^ {\left(t ^ {\prime}\right)} (d x)\right) ^ {\frac {1}{2}} \\ \leq C ^ {\prime} \gamma_ {t}, \\ \end{array}
$$
where we use  $\sup_t\int \| \nabla V_{\mathcal{L}}(\nu^{(t)})\| ^2\nu^{(t)}(dx) <   \infty$  in the last step.
9.3 Proof of Proposition 4.3 (sample complexity)
Recall that  $\mathcal{X} = \mathcal{Y} = \mathbb{R}^d$  and  $\rho(x, y) = \|x - y\|^2$  in this proposition. For the proof, we will need the following lemma which is of independent interest. We write  $\nu \leq_c \mu$  if  $\nu$  is dominated by  $\mu$  in convex order, i.e.,  $\int f \, d\nu \leq \int f \, d\mu$  for all convex functions  $f: \mathbb{R}^d \to \mathbb{R}$ .
Lemma 9.2. Let  $\mu$  have finite second moment. Given  $\nu \in \mathcal{P}(\mathbb{R}^d)$ , there exists  $\tilde{\nu} \in \mathcal{P}(\mathbb{R}^d)$  with  $\tilde{\nu} \leq_c \mu$  and
$$
\mathcal {L} _ {E O T} (\mu , \tilde {\nu}) \leq \mathcal {L} _ {E O T} (\mu , \nu).
$$
This inequality is strict if  $\nu \not\leq_{c}\mu$ . In particular, any optimizer  $\nu^{*}$  of (8) satisfies  $\nu^{*} \leq_{c} \mu$ .
Proof. Because this proof uses disintegration over  $\mathcal{V}$ , it is convenient to reverse the order of the spaces in the notation and write a generic point as  $(x,y) \in \mathcal{V} \times \mathcal{X}$ . Consider  $\pi \in \Pi(\nu, \mu)$  and its disintegration  $\pi = \nu(dx) \otimes K(x,dy)$  over  $x \in \mathcal{V}$ . Define  $T: \mathbb{R}^d \to \mathbb{R}^d$  by
$$
T (x) := \int y K (x, d y).
$$
Define also  $\tilde{\pi} := (T, \mathrm{id})_{\#} \pi$  and  $\tilde{\nu} := \tilde{\pi}_1$ . From the definition of  $T$ , we see that  $\tilde{\pi}$  is a martingale, thus  $\tilde{\nu} \leq_c \mu$ . Moreover,  $\tilde{\nu} \otimes \mu = (T, \mathrm{id})_{\#} \nu \otimes \mu$ . The data-processing inequality now shows that
$$
H (\tilde {\pi} | \tilde {\nu} \otimes \mu) \leq H (\pi | \nu \otimes \mu).
$$
On the other hand,  $\int \| \int \tilde{y} K(x,d\tilde{y}) - y\| ^2 K(x,dy)\leq \int \| x - y\| ^2 K(x,dy)$  since the barycenter minimizes the squared distance, and this inequality is strict whenever  $x\neq \int \tilde{y} K(x,d\tilde{y})$  . Thus
$$
\int \| x - y \| ^ {2} \tilde {\pi} (d x, d y) \leq \int \| x - y \| ^ {2} \pi (d x, d y),
$$
and the inequality is strict unless  $T(x) = x$  for  $\nu$ -a.e.  $x$ , which in turn is equivalent to  $\pi$  being a martingale. The claims follow.
Proof of Proposition 4.3. Subgaussianity of the optimizer follows directly from Lemma 9.2.  $\square$
Recalling that  $\inf_{\nu} \mathcal{L}_{EOT}(\nu)$  and  $\inf_{\nu} \lambda^{-1} \mathcal{L}_{BA}(\nu)$  have the same values and minimizers (given by (9) in Sec. 2.2), it suffices to show the claim for  $\mathcal{L} = \mathcal{L}_{EOT}$ . Let  $\nu^{*}$  be an optimizer of (8) (i.e., an optimal reproduction distribution) and  $\nu^{n}$  its empirical measure from  $n$  samples, then clearly
$$
\begin{array}{l} \left| \min  _ {\nu_ {n} \in \mathcal {P} _ {n} (\mathbb {R} ^ {d})} \mathcal {L} _ {E O T} (\mu , \nu_ {n}) - \min  _ {\nu \in \mathcal {P} (\mathbb {R} ^ {d})} \mathcal {L} _ {E O T} (\mu , \nu) \right| = \min  _ {\nu_ {n} \in \mathcal {P} _ {n} (\mathbb {R} ^ {d})} \mathcal {L} _ {E O T} (\mu , \nu_ {n}) - \min  _ {\nu \in \mathcal {P} (\mathbb {R} ^ {d})} \mathcal {L} _ {E O T} (\mu , \nu) \\ \leq \mathbb {E} \left[ \left| \mathcal {L} _ {E O T} (\mu , \nu^ {n}) - \mathcal {L} _ {E O T} (\mu , \nu^ {*}) \right| \right] \\ \end{array}
$$
where the expectation is taken over samples for  $\nu^n$ . The first inequality of Proposition 4.3 now follows from the sample complexity result for entropic optimal transport in [Mena and Niles-Weed, 2019, Theorem 2].
Denote by  $\nu_{m}^{*}$  the optimizer for the problem (8) with  $\mu$  replaced by  $\mu^m$ . Similarly to the above, we obtain
$$
\begin{array}{l} \mathbb {E} \left[ \left| \min  _ {\nu \in \mathcal {P} (\mathbb {R} ^ {d})} \mathcal {L} _ {E O T} (\mu , \nu) - \min  _ {\nu \in \mathcal {P} (\mathbb {R} ^ {d})} \mathcal {L} _ {E O T} (\mu^ {m}, \nu) \right| \right] \\ \leq \mathbb {E} \left[ \max  _ {\nu \in \{\nu^ {*}, \nu_ {m} ^ {*} \}} | \mathcal {L} _ {E O T} (\mu , \nu) - \mathcal {L} _ {E O T} (\mu^ {m}, \nu) | \right], \\ \end{array}
$$
where the expectation is taken over samples from  $\mu^m$ . In this situation, we cannot directly apply [Mena and Niles-Weed, 2019, Theorem 2]. However, the bound given by [Mena and Niles-Weed, 2019, Proposition 2] still applies, and the only dependence on  $\nu \in \{\nu^{*}, \nu_{m}^{*}\}$  is through their subgaussianity constants. By Lemma 9.2, these constants are bounded by the corresponding constants of  $\mu$  and  $\mu^m$ . Thus, the arguments in the proof of [Mena and Niles-Weed, 2019, Theorem 2] can be applied, yielding the second inequality of Proposition 4.3.
The final inequality of Proposition 4.3 follows from the first two inequalities (the first one being applied with  $\mu^m$ ) and the triangle inequality, where we again use the arguments in the proof of [Mena and Niles-Weed, 2019, Theorem 2] to bound the expectation over the subgaussianity constants of  $\mu^m$ .
9.4 Convergence of the proposed hybrid algorithm
In our present work, our hybrid algorithm targets the non-stochastic setting and is motivated by the fact that the BA update is in some sense orthogonal to the Wasserstein gradient update, and can only monotonically improve the objective. While empirically we observe the additional BA steps to not hurt - but rather accelerate - the convergence of WGD (see 5.1), additional effort is required to theoretically guarantee convergence of the hybrid algorithm.