SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Within information theory: Recent work by Wu et al. [2022] and Lei et al. [2023b] also note the connection between the R-D function and entropic optimal transport. Wu et al. [2022] compute the R-D function in the finite and known alphabet setting by solving a version of the EOT problem (8), whereas Lei et al. [2023b] numerically verify the equivalence (9) on a discrete problem and discuss the connection to scalar quantization. We also experimented with estimating  $R(D)$  by solving the EOT problem (8), but found it computationally much more efficient to work with the rate functional (6), and we see the primary benefit of the EOT connection as bringing in tools from statistical OT [Genevay et al., 2019, Mena and Niles-Weed, 2019, Rigollet and Stromme, 2022] for R-D estimation. Outside of information theory: Rigollet and Weed [2018] note a connection between the EOT projection problem (8) and maximum-likelihood deconvolution (10); our work complements their perspective by re-interpreting both problems through the equivalent R-D problem. Unbeknownst to us at the time, Yan et al. [2023] proposed similar algorithms to ours in the context of Gaussian mixture estimation, which we recognize as R-D estimation under quadratic distortion (see Sec. 2.3). Their work is based on gradient flow in the Fisher-Rao-Wasserstein (FRW) geometry [Chizat et al., 2018], which our hybrid algorithm can be seen as implementing. Yan et al. [2023] prove that, in an
idealized setting with infinite particles, FRW gradient descent does not get stuck at local minima; by contrast, our convergence and sample-complexity results (Prop. 4.2, 4.3) hold for any finite number of particles. We additionally consider larger-scale problems and the stochastic optimization setting.
4 Proposed method
For our algorithm, we require  $\mathcal{X} = \mathcal{Y} = \mathbb{R}^d$  and  $\rho$  be continuously differentiable. We now introduce the gradient descent algorithm in Wasserstein space to solve the problems (4) and (8). We defer all proofs to the Supplementary Material. To minimize a functional  $\mathcal{L}:\mathcal{P}(\mathcal{V})\to \mathbb{R}$  over the space of probability measures, our algorithm essentially simulates the gradient flow [Ambrosio et al., 2008] of  $\mathcal{L}$  and follows the trajectory of steepest descent in the Wasserstein geometry. In practice, we represent a measure  $\nu^{(t)}\in \mathcal{P}(\mathcal{V})$  by a collection of particles and at each time step update  $\nu^{(t)}$  in a direction of steepest descent of  $\mathcal{L}$  as given by its (negative) Wasserstein gradient. Denote by  $\mathcal{P}_n(\mathbb{R}^d)$  the set of probability measures on  $\mathbb{R}^d$  that are supported on at most  $n$  points. Our algorithm implements the parametric R-D estimator with the choice  $\mathcal{H} = \mathcal{P}_n(\mathbb{R}^d)$  (see discussions at the end of Sec. 2.1).
4.1 Wasserstein gradient descent (WGD)
Abstractly, Wasserstein gradient descent updates the variational measure  $\nu$  to its pushforward  $\tilde{\nu}$  under the map  $(\mathrm{id} - \gamma \Psi)$ , for a function  $\Psi : \mathbb{R}^d \to \mathbb{R}^d$  called the Wasserstein gradient of  $\mathcal{L}$  at  $\nu$  (see below) and a step size  $\gamma$ . To implement this scheme, we represent  $\nu$  as a convex combination of Dirac measures,  $\nu = \sum_{i=1}^{n} w_i \delta_{x_i}$  with locations  $\{x_i\}_{i=1}^n \subset \mathbb{R}^d$  and weights  $\{w_i\}_{i=1}^n$ . The algorithm moves each particle  $x_i$  in the direction of  $-\Psi(x_i)$ , more precisely,  $\tilde{\nu} = \sum_{i=1}^{n} w_i \delta_{x_i - \gamma \Psi(x_i)}$ .
Algorithm 1 Wasserstein gradient descent
Inputs: Loss function  $\mathcal{L} \in \{\mathcal{L}_{BA}, \mathcal{L}_{EOT}\}$ ; data distribution  $\mu \in \mathcal{P}(\mathbb{R}^d)$ ; the number of particles  $n \in \mathbb{N}$ ; total number of iterations  $N \in \mathbb{N}$ ; step sizes  $\gamma_1, \ldots, \gamma_N$ ; batch size  $m \in \mathbb{N}$ .  
for  $t = 1, \ldots, N$  do  
Pick an initial measure  $\nu^{(0)} \in \mathcal{P}_n(\mathbb{R}^d)$ , e.g., setting the particles to  $n$  random samples from  $\mu$ .  
if support of  $\mu$  contains more than  $m$  points then  
\[
\mu^m \leftarrow \frac{1}{m} \sum_{i=1}^{m} \delta_{x_i} \text{ for } x_1, \ldots, x_m \text{ independent samples from } \mu
\]  
\[
\Psi^{(t)} \leftarrow \text{Wasserstein gradient of } \mathcal{L}(\mu^m, \cdot) \text{ at } \nu^{(t-1)} \text{ (see Definition 4.1)}
\]  
else  
\[
\Psi^{(t)} \leftarrow \text{Wasserstein gradient of } \mathcal{L}(\mu, \cdot) \text{ at } \nu^{(t-1)} \text{ (see Definition 4.1)}
\]  
end if  
\[
\nu^{(t)} \leftarrow (\mathrm{id} - \gamma_t \Psi^{(t)})_\# \nu^{(t-1)} \{\#\} \text{ denotes pushforward}
\]  
end for  
Return:  $\nu^{(N)}$
Since the optimization objectives (4) and (8) appear as integrals w.r.t. the data distribution  $\mu$ , we can also apply stochastic optimization and perform stochastic gradient descent on mini-batches with size  $m$ . This allows us to handle a very large or infinite amount of data samples, or when the source is continuous. We formalize the procedure in Algorithm 1.
The following gives a constructive definition of a Wasserstein gradient which forms the computational basis of our algorithm. In the literature, the Wasserstein gradient is instead usually defined as a Fréchet differential (cf. [Ambrosio et al., 2008, Definition 10.1.1]), but we emphasize that in smooth settings, the given definition recovers the one from the literature (cf. [Chizat, 2022, Lemma A.2]).
Definition 4.1. For a functional  $\mathcal{L}:\mathcal{P}(\mathcal{Y})\to \mathbb{R}$  and  $\nu \in \mathcal{P}(\mathcal{Y})$  , we say that  $V_{\mathcal{L}}(\nu):\mathbb{R}^{d}\rightarrow \mathbb{R}$  is a first variation of  $\mathcal{L}$  at  $\nu$  if
$$
\lim  _ {\varepsilon \rightarrow 0} \frac {\mathcal {L} ((1 - \varepsilon) \nu + \varepsilon \tilde {\nu}) - \mathcal {L} (\nu)}{\varepsilon} = \int V _ {\mathcal {L}} (\nu) d (\tilde {\nu} - \nu) f o r a l l \tilde {\nu} \in \mathcal {P} (\mathcal {Y}).
$$
We call its (Euclidean) gradient  $\nabla V_{\mathcal{L}}(\nu):\mathbb{R}^d\to \mathbb{R}^d$  , if it exists, the Wasserstein gradient of  $\mathcal{L}$  at  $\nu$
For  $\mathcal{L} = \mathcal{L}_{EOT}$ , the first variation is given by the Kantorovich potential, which is the solution of the convex dual of  $\mathcal{L}_{EOT}$  and commonly computed by Sinkhorn's algorithm [Peyre and Cuturi, 2019,
Nutz, 2021]. Specifically, let  $(\varphi^{\nu},\psi^{\nu})$  be potentials for  $\mathcal{L}_{EOT}(\mu ,\nu)$ . Then  $V_{\mathcal{L}}(\nu) = \psi^{\nu}$  is the first variation w.r.t.  $\nu$  (cf. [Carlier et al., 2022, equation (20)]), and hence  $\nabla \psi^{\nu}$  is the Wasserstein gradient. This gradient exists whenever  $\rho$  is differentiable and the marginals are sufficiently light-tailed; we give details in Sec. 9.1 of the Supplementary Material. For  $\mathcal{L} = \mathcal{L}_{BA}$ , the first variation can be computed explicitly. As derived in Sec. 9.1 of the Supplementary Material, the first variation 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)
$$
and then the Wasserstein gradient is  $\nabla \mathcal{L}_{BA}(\nu) = \nabla \psi^{\nu}$ . We observe that  $\psi^\nu (y)$  is computationally cheap; it corresponds to running a single iteration of Sinkhorn's algorithm. By contract, finding the potential for  $\mathcal{L}_{EOT}$  requires running Sinkhorn's algorithm to convergence.
Like the usual Euclidean gradient, the Wasserstein gradient can be shown to possess a linearization property, whereby the loss functional is reduced by taking a small enough step along its Wasserstein gradient. Following [Carlier et al., 2022], we state it as follows: for any  $\tilde{\nu} \in \mathcal{P}(\mathcal{Y})$  and  $\pi \in \Pi(\nu, \tilde{\nu})$ ,
$$
\begin{array}{l} \mathcal {L} (\tilde {\nu}) - \mathcal {L} (\nu) = \int (y - x) ^ {\top} \nabla V _ {\mathcal {L}} (\nu) (x) \pi (d x, d y) + o \left(\int \| y - x \| ^ {2} \pi (d x, d y)\right), \tag {14} \\ \left| \int \| \nabla V _ {\mathcal {L}} (\nu) \| ^ {2} d \nu - \int \| \nabla V _ {\mathcal {L}} (\tilde {\nu}) \| ^ {2} d \tilde {\nu} \right| \leq C W _ {2} (\nu , \tilde {\nu}). \\ \end{array}
$$
The first line of (14) is proved for  $\mathcal{L}_{EOT}$  in [Carlier et al., 2022, Proposition 4.2] in the case that the marginals are compactly supported and  $\rho$  is twice continuously differentiable. In this setting, the second line of (14) follows using  $a^2 - b^2 = (a + b)(a - b)$  and a combination of boundedness and Lipschitz continuity of  $\nabla V_{\mathcal{L}}$ , see [Carlier et al., 2022, Proposition 2.2 and Corollary 2.4].
The linearization property given by (14) enables us to show that Wasserstein gradient descent for  $\mathcal{L}_{EOT}$  and  $\mathcal{L}_{BA}$  converges to a stationary point under mild conditions:
Proposition 4.2 (Convergence of Wasserstein gradient descent). Let  $\gamma_1 \geq \gamma_2 \geq \dots \geq 0$  satisfy  $\sum_{k=1}^{\infty} \gamma_k = \infty$  and  $\sum_{k=1}^{\infty} \gamma_k^2 < \infty$ . Let  $\mathcal{L}: \mathcal{P}(\mathbb{R}^d) \to \mathbb{R}$  be Wasserstein differentiable in the sense that (14) holds. Denoting by  $\nu^{(t)}$  the steps in Algorithm 1, and suppose that  $\mathcal{L}(\nu^{(0)})$  is finite and  $\int \| \nabla V_{\mathcal{L}}(\nu^{(t)}) \|^2 d\nu^{(t)}$  is bounded. Then
$$
\lim _ {t \to \infty} \int \| \nabla V _ {\mathcal {L}} (\nu^ {(t)}) \| ^ {2} d \nu^ {(t)} = 0.
$$
4.2 Hybrid algorithm
A main limitation of the BA algorithm is that the support of  $\nu^{(t)}$  is restricted to that of the (possibly bad) initialization  $\nu^{(0)}$ . On the other hand, Wasserstein gradient descent (Algorithm 1) only evolves the particle locations of  $\nu^{(t)}$ , but not the weights, which are fixed to be uniform by default. We therefore consider a hybrid algorithm where we alternate between WGD and the BA update steps, allowing us to optimize the particle weights as well. Experimentally, this translates to faster convergence than the base WGD algorithm (Sec. 5.1). Note however, unlike WGD, the hybrid algorithm does not directly lend itself to the stochastic optimization setting, as BA updates on mini-batches no longer guarantee monotonic improvement in the objective and can lead to divergence. We treat the convergence of the hybrid algorithm in the Supplementary Material Sec. 9.4.
4.3 Sample complexity
Let  $\mathcal{X} = \mathcal{Y} = \mathbb{R}^d$  and  $\rho (x,y) = \| x - y\| ^2$ . Leveraging work on the statistical complexity of EOT [Mena and Niles-Weed, 2019], we obtain finite-sample bounds for the theoretical estimators implemented by WGD in terms of the number of particles and source samples. The bounds hold for both the R-D problem (4) and EOT projection problem (8) as they share the same optimizers (see Sec. 2.2), and strengthen existing asymptotic results for empirical R-D estimators [Harrison and Kontoyiannis, 2008]. We note that a recent result by Rigollet and Stromme [2022] might be useful for deriving alternative bounds under distortion functions other than the quadratic.
Proposition 4.3. Let  $\mu$  be  $\sigma^2$ -subgaussian. Then every optimizer  $\nu^{*}$  of (4) and (8) is also  $\sigma^2$ -subgaussian. Consider  $\mathcal{L} \coloneqq \mathcal{L}_{EOT}$ . For a constant  $C_d$  only depending on  $d$ , we have
$$
\begin{array}{l} \left| \min  _ {\nu \in \mathcal {P} (\mathbb {R} ^ {d})} \mathcal {L} (\mu , \nu) - \min  _ {\nu_ {n} \in \mathcal {P} _ {n} (\mathbb {R} ^ {d})} \mathcal {L} (\mu , \nu_ {n}) \right| \leq C _ {d} \epsilon \left(1 + \frac {\sigma^ {\lceil 5 d / 2 \rceil + 6}}{\epsilon^ {\lceil 5 d / 4 \rceil + 3}}\right) \frac {1}{\sqrt {n}}, \\ \mathbb {E} \left[ \left| \min  _ {\nu \in \mathcal {P} (\mathbb {R} ^ {d})} \mathcal {L} (\mu , \nu) - \min  _ {\nu \in \mathcal {P} (\mathbb {R} ^ {d})} \mathcal {L} (\mu^ {m}, \nu) \right| \right] \leq C _ {d} \epsilon \left(1 + \frac {\sigma^ {\lceil 5 d / 2 \rceil + 6}}{\epsilon^ {\lceil 5 d / 4 \rceil + 3}}\right) \frac {1}{\sqrt {m}}, \\ \mathbb{E}\left[\left|\min_{\nu \in \mathcal{P}(\mathbb{R}^{d})}\mathcal{L}(\mu ,\nu) - \min_{\nu_{n}\in \mathcal{P}_{n}(\mathbb{R}^{d})}\mathcal{L}(\mu^{m},\nu_{n})\right|\right]\leq C_{d}\epsilon \left(1 + \frac{\sigma^{\lceil 5d / 2\rceil + 6}}{\epsilon^{\lceil 5d / 4\rceil + 3}}\right)\left(\frac{1}{\sqrt{m}} +\frac{1}{\sqrt{n}}\right), \\ \end{array}
$$
for all  $n, m \in \mathbb{N}$ , where  $\mathcal{P}_n(\mathbb{R}^d)$  is the set of probability measures over  $\mathbb{R}^d$  supported on at most  $n$  points,  $\mu^m$  is the empirical measure of  $\mu$  with  $m$  independent samples and the expectation  $\mathbb{E}[\cdot]$  is over these samples. The same inequalities hold for  $\mathcal{L} \coloneqq \lambda^{-1}\mathcal{L}_{BA}$ , with the identification  $\epsilon = \lambda^{-1}$ .
4.4 Estimation of rate and distortion
Here, we describe our estimator for an upper bound  $(\mathcal{D},\mathcal{R})$  of  $R(D)$  after solving the unconstrained problem (3). We provide more details in Sec. 8 of the Supplementary Material.
For any given pair of  $\nu$  and  $K$ , we always have that  $\mathcal{D} \coloneqq \int \rho d(\mu \otimes K)$  and  $\mathcal{R} \coloneqq H(\mu \otimes K|\mu \otimes \nu)$  lie on an upper bound of  $R(D)$  [Berger, 1971]. The two quantities can be estimated by standard Monte Carlo provided we can sample from  $\mu \otimes K$  and evaluate the density  $\frac{d\mu \otimes K}{d\mu \otimes \nu}(x,y) = \frac{dK(x,\cdot)}{d\nu}(y)$ .
When only  $\nu$  is given, e.g., obtained from optimizing (6) with WGD or NERD, we estimate an R-D upper bound as follows. As in the BA algorithm, we construct a kernel  $K_{\nu}$  similarly to (11), i.e.,  $\frac{dK_{\nu}(x,\cdot)}{d\nu} (y) = \frac{e^{-\lambda\rho(x,y)}}{\int e^{-\lambda\rho(x,\tilde{y})}\nu(d\tilde{y})}$ ; then we estimate  $(\mathcal{D},\mathcal{R})$  using the pair  $(\nu ,K_{\nu})$  as described earlier.
As NERD uses a continuous  $\nu$ , we follow [Lei et al., 2023a] and approximate it with its  $n$ -sample empirical measure to estimate  $(\mathcal{D},\mathcal{R})$ . A limitation of NERD, BA, and our method is that they tend to converge to a rate estimate of at most  $\log (n)$ , where  $n$  is the support size of  $\nu$ . This is because as the algorithms approach an  $n$ -point minimizer  $\nu_{n}^{*}$  of the R-D problem, the rate estimate  $\mathcal{R}$  approaches the mutual information of  $\mu \otimes K_{\nu_n^*}$ , which is upper-bounded by  $\log (n)$  [Eckstein and Nutz, 2022]. In practice, this means if a target point of  $R(D)$  has rate  $r$ , then we need  $n\geq e^r$  to estimate it accurately.
4.5 Computational considerations
Common to all the aforementioned methods is the evaluation of a pairwise distortion matrix between  $m$  points in  $\mathcal{X}$  and  $n$  points in  $\mathcal{Y}$ , which usually has a cost of  $\mathcal{O}(mnd)$  for a  $d$ -dimensional source. While RD-VAE uses  $n = 1$  (in the reparameterization trick), the other methods (BA, WGD, NERD) typically use a much larger  $n$  and thus has the distortion computation as their main computation bottleneck. Compared to BA and WGD, the neural methods (RD-VAE, NERD) incur additional computation from neural network operations, which can be significant for large networks.
For NERD and WGD (and BA), the rate estimate upper bound of  $\log (n)$  nats/sample (see Sec. 4.4) can present computational challenges. To target a high-rate setting, a large number of  $\nu$  particles (high  $n$ ) is required, and care needs to be taken to avoid running out of memory during the distortion matrix computation (one possibility is to use a small batch size  $m$  with stochastic optimization).
We compare the empirical performance of our proposed method (WGD) and its hybrid variant with Blahut-Arimoto (BA) [Blahut, 1972, Arimoto, 1972], RD-VAE [Yang and Mandt, 2022], and NERD [Lei et al., 2022] on the tasks of maximum-likelihood deconvolution and estimation of R-D upper bounds. While we experimented with WGD for both  $\mathcal{L}_{BA}$  and  $\mathcal{L}_{EOT}$ , we found the former to be 10 to 100 times faster computationally while giving similar or better results; we therefore focus on WGD for  $\mathcal{L}_{BA}$  in discussions below. For the neural-network baselines, we use the same (or as similar as possible) network architectures as in the original work [Yang and Mandt, 2022, Lei et al., 2023a]. We use the Adam optimizer for all gradient-based methods, except we use simple gradient descent
Figure 2: Losses over iterations. Shading corresponds to one standard deviation over random initializations.
Figure 3: Visualizing  $\mu$  samples (top left), as well as the  $\nu$  returned by various algorithms compared to the ground truth  $\nu^{*}$  (cyan).
Figure 4: Optimality gap v.s. the number of particles  $n$  used, on deconvolution problems with different dimension  $d$  and distortion multiplier  $\lambda$ . The first three panels fix  $\lambda = \sigma^{-2}$  and increase  $d$ , and the right-most panel corresponds to the 2-D problem with a higher  $\lambda$  (denoising with a narrower Gaussian kernel). Overall the WGD methods attain higher accuracy for a given budget of  $n$ .
with a decaying step size in Sec. 5.1 to better compare the convergence speed of WGD and its hybrid variant. Further experiment details and results are given in the Supplementary Material Sec. 11.
5.1 Deconvolution
To better understand the behavior of the various algorithms, we apply them to a deconvolution problem with known ground truth (see Sec. 2.3). We adopt the Gaussian noise as before, letting  $\alpha$  be the uniform measure on the unit circle in  $\mathbb{R}^2$  and the source  $\mu = \alpha * \mathcal{N}(0, \sigma^2)$  with  $\sigma^2 = 0.1$ .
We use  $n = 20$  particles for BA, NERD, WGD and its hybrid variant. We use a two-layer network for NERD and RD-VAE with some hand-tuning (we replace the softplus activation in the original RD-VAE network by ReLU as it led to difficulty in optimization). Fig. 2 plots the resulting loss curves and shows that the proposed algorithms converge the fastest to the ground truth value  $OPT := \mathcal{L}(\alpha)$ . In Fig. 3, we visualize the final  $\nu^{(t)}$  at the end of training, compared to the ground truth  $\nu^{*} = \alpha$  supported on the circle (colored in cyan). Note that we initialize  $\nu^{(0)}$  for BA, WGD, and its hybrid variant to the same  $n$  random data samples. While BA is stuck with the randomly initialized particles and assigns large weights to those closer to the circle, WGD learns to move the particles to uniformly cover the circle. The hybrid algorithm, being able to reweight particles to reduce their transportation cost, learns a different solution where a cluster of particles covers the top-left portion of the circle with small weights while the remaining particles evenly covers the rest. Unlike our particle-based methods, the neural methods generally struggle to place the support of their  $\nu$  exactly on the circle.
We additionally compare how the performance of BA, NERD, and the proposed algorithms scale to higher dimensions and a higher  $\lambda$  (corresponding to lower entropic regularization in  $\mathcal{L}_{EOT}$ ). Fig. 4 plots the gap between the converged and the optimal losses for the algorithms, and demonstrates the proposed algorithms to be more particle-efficient and scale more favorably than the alternatives which also use  $n$  particles in the reproduction space. We additionally visualize how the converged particles for our methods vary across the R-D trade-off in Fig. 6 of the Supplementary Material.
Figure 5: Left, Middle: R-D bound estimates on the physics, speech datasets [Yang and Mandt, 2022] and MNIST training set. Right: Example speed comparisons of WGD and neural upper bound methods on MNIST, with WGD converging at least an order of magnitude faster. On each of the dataset we also include an R-D lower bound estimated using the method of [Yang and Mandt, 2022].
5.2 Higher-dimensional data
We perform  $R(D)$  estimation on higher-dimensional data, including the physics and speech datasets from [Yang and Mandt, 2022] and MNIST [LeCun et al., 1998]. As the memory cost to operating on the full datasets becomes prohibitive, we focus on NERD, RD-VAE, and WGD using mini-batch stochastic gradient descent. BA and hybrid WGD do not directly apply in the stochastic setting, as BA updates on random mini-batches can lead to divergence (as discussed in Sec. 4.2).
Fig. 5 plots the estimated R-D bounds on the datasets, and compares the convergence speed of WGD and neural methods in both iteration count and compute time. Overall, we find WGD to require minimal tuning and obtains the tightest R-D upper bounds within the  $\log (n)$  rate limit (see Sec. 4.4), and consistently obtains tighter bounds than NERD given the same computation budget.
In this work, we leverage tools from optimal transport to develop a new approach for estimating the rate-distortion function in the continuous setting. Compared to state-of-the-art neural approaches [Yang and Mandt, 2022, Lei et al., 2022], our Wasserstein gradient descent algorithm offers complementary strengths: 1) It requires a single main hyperparameter  $n$  (the number of particles) and no network architecture tuning; and 2) empirically we found it to converge significantly faster and rarely end up in bad local optima; increasing  $n$  almost always yielded an improvement (unless the bound is already close to being tight). From a modeling perspective, a particle representation may be inherently more efficient when the optimal reproduction distribution is singular or has many disconnected modes (shown, e.g., in Figs. 1 and 3). However, like NERD [Lei et al., 2022], our method has a fundamental limitation - it requires an  $n$  that is exponential in the rate of the targeted  $R(D)$  point to estimate it accurately (see Sec. 4.4). Thus, a neural method like RD-VAE [Yang and Mandt, 2022] may still be preferable on high-rate sources, while our particle-based method stands as a state-of-the-art solution on lower-rate sources with substantially less tuning or computation requirements.
Besides R-D estimation, our algorithm also applies to the mathematically equivalent problems of maximum likelihood deconvolution and projection under an entropic optimal transport (EOT) cost, and may find other connections and applications. Indeed, the EOT projection view of our algorithm is further related to optimization-based approaches to sampling [Wibisono, 2018], variational inference [Liu and Wang, 2016], and distribution compression [Shetty et al., 2021]. Our particle-based algorithm also generalizes optimal quantization (corresponding to  $\epsilon = 0$  in  $\mathcal{L}_{\mathrm{EOT}}$  and projection of the source under the Wasserstein distance [Graf and Luschgy, 2007, Gray, 2013]) to incorporate a rate constraint  $(\epsilon > 0)$ , and it would be interesting to explore the use of the resulting rate-distortion optimal quantizer for practical data compression and communication.
Roger B Grosse, Zoubin Ghahramani, and Ryan P Adams. Sandwiching the marginal likelihood using bidirectional monte carlo. arXiv preprint arXiv:1511.02543, 2015.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http://github.com/google/jax.
Improved estimates on the fundamental cost of data compression aids the development and analysis of compression algorithms. This helps researchers and engineers make better decisions about where to allocate their resources to improve certain compression algorithms, and can translate to economic gains for the broader society. However, like most machine learning algorithms trained on data, the output of our estimator is only accurate insofar as the training data is representative of the population distribution of interest, and practitioners need to ensure this in the data collection process.