Estimating the Rate-Distortion Function by Wasserstein Gradient Descent

For a memoryless data source $X$ with distribution $P_{X}$, its rate-distortion (R-D) function describes the minimum possible number of bits per sample needed to reproduce the source within a prescribed distortion threshold $D$. Let the source and reproduction take values in two sets ${\mathcal{X}}$ and ${\mathcal{Y}}$, known as the source and reproduction alphabets, and let $\rho:({\mathcal{X}},{\mathcal{Y}})\to[0,\infty)$ be a given distortion function. The R-D function is defined by the following optimization problem (Polyanskiy and Wu, 2022),

where $Q_{Y|X}$ is any Markov kernel from ${\mathcal{X}}$ to ${\mathcal{Y}}$ conceptually associated with a (possibly) stochastic compression algorithm, and $I(X;Y)$ is the mutual information of the joint distribution $P_{X}Q_{Y|X}$.

For ease of presentation, we now switch to a more abstract notation without reference to random variables. We provide the precise definitions in the Supplementary Material. Let ${\mathcal{X}}$ and ${\mathcal{Y}}$ be standard Borel spaces; let $\mu\in\mathcal{P}({\mathcal{X}})$ be a fixed probability measure on ${\mathcal{X}}$, which should be thought of as the source distribution $P_{X}$. For a measure $\pi$ on the product space ${\mathcal{X}}\times{\mathcal{Y}}$, the notation $\pi_{1}$ (or $\pi_{2}$) denotes the first (or second) marginal of $\pi$. For any $\nu\in\mathcal{P}({\mathcal{Y}})$, we denote by $\Pi(\mu,\nu)$ the set of couplings between $\mu$ and $\nu$ (i.e., $\pi_{1}=\mu$ and $\pi_{2}=\nu$). Similarly, $\Pi(\mu,\cdot)$ denotes the set of measures $\pi$ with $\pi_{1}=\mu$. Throughout the paper, $K$ denotes a transition kernel (conditional distribution) from ${\mathcal{X}}$ to ${\mathcal{Y}}$, and $\mu\otimes K$ denotes the product measure formed by $\mu$ and $K$. Then $R(D)$ is equivalent to

where $H$ denotes relative entropy, i.e., for two measures $\alpha,\beta$ defined on a common measurable space, $H(\alpha|\beta):=\int\log(\frac{d\alpha}{d\beta})d\alpha$ when $\alpha$ is absolutely continuous w.r.t $\beta$, and infinite otherwise.

To make the problem more tractable, we follow the approach of the classic Blahut–Arimoto algorithm (Blahut, 1972; Arimoto, 1972) (to be discussed in Sec. 3.1) and work with an equivalent unconstrained Lagrangian problem as follows. Instead of parameterizing the R-D function via a distortion threshold $D$, we parameterize it via a Lagrange multiplier $\lambda\geq 0$. For each fixed $\lambda$ (usually selected from a predefined grid), we aim to solve the following optimization problem,

Geometrically, $F_{\lambda}(\mu)\in\mathbb{R}$ is the y-axis intercept of a tangent line to the $R(D)$ with slope $-\lambda$, and $R(D)$ is determined by the convex envelope of all such tangent lines (Gray, 2011). To simplify notation, we often drop the dependence on $\lambda$ (e.g., we write $F(\mu)=F_{\lambda}(\mu)$) whenever it is harmless.

To set the stage for our later developments, we write the unconstrained R-D problem as

where we refer to the optimization objective $\mathcal{L}_{BA}$ as the rate function (Harrison and Kontoyiannis, 2008). We abuse the notation to write $\mathcal{L}_{BA}(\nu):=\mathcal{L}_{BA}(\mu,\nu)$ when it is viewed as a function of $\nu$ only, and refer to it as the rate functional. The rate function characterizes a generalized Asymptotic Equipartition Property, where $\mathcal{L}_{BA}(\mu,\nu)$ is the asymptotically optimal cost of lossy compression of data $X\sim\mu$ using a random codebook constructed from samples of $\nu$ (Dembo and Kontoyiannis, 2002). Notably, the optimization in (5) can be solved analytically (Csiszár, 1974a, Lemma 1.3), and $\mathcal{L}_{BA}$ simplifies to

In practice, the source $\mu$ is only accessible via independent samples, on the basis of which we propose to estimate its $R(D)$, or equivalently $F(\mu)$. Let $\mu^{m}$ denote an $m$-sample empirical measure of $\mu$, i.e., $\mu^{m}=\sum_{i=1}^{m}\delta_{x_{i}}$ with $x_{1,...,n}$ being independent samples from $\mu$, which should be thought of as the “training data”. Following Harrison and Kontoyiannis (2008), we consider two kinds of (plug-in) estimators for $F(\mu)$: (1) the non-parametric estimator $F(\mu^{m})$, and (2) the parametric estimator $F^{\mathcal{H}}(\mu^{m}):=\inf_{\nu\in\mathcal{H}}\mathcal{L}_{BA}(\mu^{m},\nu)$, where $\mathcal{H}$ is a family of probability measures on ${\mathcal{Y}}$. Harrison and Kontoyiannis (2008) showed that under rather broad conditions, both kinds of estimators are strongly consistent, i.e., $F(\mu^{m})$ converges to $F(\mu)$ (and respectively, $F^{\mathcal{H}}(\mu^{m})$ to $F^{\mathcal{H}}(\mu)$) with probability one as $m\to\infty$. Our algorithm will implement the parametric estimator $F^{\mathcal{H}}(\mu^{m})$ with $\mathcal{H}$ chosen to be the set of probability measures with finite support, and we will develop finite-sample convergence results for both kinds of estimators in the continuous setting (Proposition. 4.3).

The R-D problem turns out to have a close connection to entropic optimal transport (EOT) (Peyré and Cuturi, 2019), which we will exploit in Sec. 4.3 to obtain sample complexity results under our approach. For $\epsilon>0$, the entropy-regularized optimal transport problem is given by

We now consider the problem of projecting $\mu$ onto $\mathcal{P}({\mathcal{Y}})$ under the cost $\mathcal{L}_{EOT}$:

In the OT literature this is known as the (regularized) Kantorovich estimator (Bassetti et al., 2006) for $\mu$, and can also be viewed as a Wasserstein barycenter problem (Agueh and Carlier, 2011).

With the identification $\epsilon=\lambda^{-1}$, problem (8) is in fact equivalent to the R-D problem (4): compared to ${\mathcal{L}}_{BA}$ (5), the extra constraint on the second marginal of $\pi$ in ${\mathcal{L}}_{EOT}$ (7) is redundant at the optimal $\nu$. More precisely, Lemma 7.1 shows that (we omit the notational dependence on $\mu$ when it is fixed):

Existence of a minimizer holds under mild conditions, for instance if ${\mathcal{X}}={\mathcal{Y}}=\mathbb{R}^{d}$ and $\rho(x,y)$ is a coercive lower semicontinuous function of $y-x$ (Csiszár, 1974a, p. 66).

The connection between R-D and maximum-likelihood estimation has been observed in the information theory, machine learning and compression literature (Harrison and Kontoyiannis, 2008; Alemi et al., 2018; Ballé et al., 2017; Theis et al., 2017; Yang et al., 2020; Yang and Mandt, 2022). Here, we bring attention to a basic equivalence between the R-D problem and maximum-likelihood deconvolution, where the connection is particularly natural under a quadratic distortion function. Also see (Rigollet and Weed, 2018) for a related discussion that inspired ours and extension to a non-quadratic distortion. We provide further insight from the view of variational learning and inference in Section 10.

Maximum-likelihood deconvolution is a classical problem of non-parametric statistics and mixture models (Carroll and Hall, 1988; Lindsay and Roeder, 1993). The deconvolution problem is concerned with estimating an unknown distribution $\alpha$ from noise-corrupted observations $X_{1},X_{2},...$, where for each $i\in\mathbb{N}$, we have $X_{i}=Y_{i}+N_{i}$, $Y_{i}\stackrel{{\scriptstyle i.i.d.}}{{\sim}}\alpha$, and $N_{i}$ are i.i.d. independent noise variables with a known distribution. For concreteness, suppose all variables are $\mathbb{R}^{d}$ valued and the noise distribution is $\mathcal{N}(0,\sigma^{2}\text{I}_{d})$ with Lebesgue density $\phi_{\sigma^{2}}$. Denote the distribution of the observations $X_{i}$ by $\mu$. Then $\mu$ has a Lebesgue density given by the convolution $\alpha*\phi_{\sigma^{2}}(x):=\int\phi_{\sigma^{2}}(x-y)\alpha(dy)$. Here, we consider the population-level (instead of the usual sample-based) maximum-likelihood estimator (MLE) for $\alpha$:

and observe that $\nu^{*}=\alpha$. Plugging in the density $\phi_{\sigma^{2}}(x)\propto e^{-\frac{1}{2\sigma^{2}}\|x\|^{2}}$, we see that the MLE problem (10) is equivalent to the R-D problem (4) with $\rho(x,y)=\frac{1}{2}\|x-y\|^{2},\lambda=\frac{1}{\sigma^{2}}$, and ${\mathcal{L}}_{BA}$ given by (6) in the form of a marginal log-likelihood. Thus the R-D problem has the interpretation of estimating a distribution from its noisy observations given through $\mu$, assuming a Gaussian noise with variance $\frac{1}{\lambda}$.

This connection suggests analytical solutions to the R-D problem for a variety of sources that arise from convolving an underlying distribution with Gaussian noise. Consider an R-D problem (4) with ${\mathcal{X}}={\mathcal{Y}}=\mathbb{R}^{d},\rho(x,y)=\frac{1}{2}\|x-y\|^{2}$, and let the source $\mu$ be the convolution between an arbitrary measure $\alpha\in{\mathcal{P}}({\mathcal{Y}})$ and Gaussian noise with known variance $\sigma^{2}$. E.g., using a discrete measure for $\alpha$ results in a Gaussian mixture source with equal covariance among its components. When $\lambda=\frac{1}{\sigma^{2}}$, we recover exactly the population-MLE problem (10) discussed earlier, which has the solution $\nu^{*}=\alpha$. While this allows us to obtain one point of $R(D)$, we can in fact extend this idea to any $\lambda\geq\frac{1}{\sigma^{2}}$ and obtain the analytical form for the corresponding segment of the R-D curve. Specifically, for any $\lambda\geq\frac{1}{\sigma^{2}}$, applying the summation rule for independent Gaussians reveals the source distribution $\mu$ as

i.e., as the convolution between another underlying distribution $\alpha_{\lambda}$ and independent noise with variance $\frac{1}{\lambda}$. A solution to the R-D problem (3) is then analogously given by $\nu^{*}=\alpha_{\lambda}$, with the corresponding optimal coupling given by $\nu^{*}\otimes\tilde{K}$, $\tilde{K}(y,dx)=\mathcal{N}(y,\frac{1}{\lambda})$. ¹¹1 Here $\tilde{K}$ maps from the reproduction to the source alphabet, opposite to the kernel $K$ elsewhere in the text. Evaluating the distortion and mutual information of the coupling then yields the $R(D)$ point associated with $\lambda$. Fig. 1 illustrates the $R(D)$ of a toy Gaussian mixture source, along with the $\nu^{*}$ estimated by our proposed WGD algorithm (Sec. 4); note that $\nu^{*}$ transitions from continuous (a Gaussian mixture with smaller component variances) to singular (a mixture of two Diracs) at $\lambda=\sigma^{-2}$. See caption for more details.

The Blahut–Arimoto (BA) algorithm (Blahut, 1972; Arimoto, 1972) is the default method for computing $R(D)$ for a known and discrete case. For a fixed $\lambda$, BA carries out the optimization problem (3) via coordinate ascent. Starting from an initial measure $\nu^{(0)}\in\mathcal{P}({\mathcal{Y}})$, the BA algorithm at step $t$ computes an updated pair $(\nu^{(t+1)},K^{(t+1)})$ as follows

When the alphabets are finite, the above computation can be carried out in matrix and vector operations, and the resulting sequence $\{(\nu^{(t)},K^{(t)})\}_{t=1}^{\infty}$ can be shown to converge to an optimum of (3); cf. (Csiszár, 1974b, 1984). When the alphabets are not finite, e.g., ${\mathcal{X}}={\mathcal{Y}}=\mathbb{R}^{d}$, the BA algorithm no longer applies, as it is unclear how to digitally represent the measure $\nu$ and kernel $K$ and to tractably perform the integrals required by the algorithm. The common workaround is to perform a discretization step and then apply BA on the resulting discrete problem.

One standard discretization method is to tile up the alphabets with small bins (Gray and Neuhoff, 1998). This quickly becomes infeasible as the number of dimensions increases. We therefore consider discretizing the data space ${\mathcal{X}}$ to be the support of training data distribution $\mu^{m}$, i.e., the discretized alphabet is the set of training samples; this can be justified by the consistency of the parametric R-D estimator $F^{\mathcal{H}}(\mu^{m})$ (Harrison and Kontoyiannis, 2008). It is less clear how to discretize the reproduction space ${\mathcal{Y}}$, especially in high dimensions. Since we work with ${\mathcal{X}}={\mathcal{Y}}$, we will disretize ${\mathcal{Y}}$ similarly and use an $n$-element random subset of the training samples, as also considered by Lei et al. (2023a). As we will show, this rather arbitrary placement of the support of $\nu$ results in poor performance, and can be significantly improved from our perspective of evolving particles.

RD-VAE ((Yang and Mandt, 2022)): To overcome the limitations of the BA algorithm, Yang and Mandt (2022) proposed to parameterize the transition kernel $K$ and reproduction distribution $\nu$ of the BA algorithm by neural density estimators (Papamakarios et al., 2021), and optimize the same objective (3) by (stochastic) gradient descent. They estimate (3) by Monte Carlo using joint samples $(X_{i},Y_{i})\sim\mu\otimes K$; in particular, the relative entropy can be written as $H(\mu\otimes K|\mu\otimes\nu)=\int\int\log\left(\frac{dK(x,\cdot)}{d\nu}(y)\right)K(x,dy)\mu(dx)$, where the integrand is computed exactly via a density ratio. In practice, an alternative parameteriation is often used where the neural density estimators are defined on a lower dimensional latent space than the reproduction alphabet, and the resulting approach is closely related to VAEs (Kingma and Welling, 2013). Yang and Mandt (2022) additionally propose a neural estimator for a lower bound on $R(D)$, based on a dual representation due to Csiszár (1974a).    NERD (Lei et al., 2023a): Instead of working with the transition kernel $K$ as in the RD-VAE, Lei et al. (2023a) considered optimizing the form of the rate functional in (6), via gradient descent on the parameters of $\nu$ parameterized by a neural network. Let $\nu^{{\mathcal{Z}}}$ be a base distribution over ${\mathcal{Z}}=\mathbb{R}^{K}$, such as the standard Gaussian, and $\omega:{\mathcal{Z}}\to{\mathcal{Y}}$ be a decoder network. The variational measure $\nu$ is then modeled as the image measure of $\nu^{{\mathcal{Z}}}$ under $\omega$. To evaluate and optimize the objective (6), the intractable inner integral w.r.t. $\nu$ is replaced with a plug-in estimator, so that for a given $x\in{\mathcal{X}}$,

After training, we estimate an R-D upper bound using $n$ samples from $\nu$ (to be discussed in Sec. 4.4).

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.

Abstractly, Wasserstein gradient descent updates the variational measure $\nu$ to its pushforward $\tilde{\nu}$ under the map $(\text{id}-\gamma\Psi)$, for a function $\Psi:\mathbb{R}^{d}\rightarrow\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})}$.

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)).

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 (Peyré 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

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})$,

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:

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.

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.

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}:=\int\rho d(\mu\otimes K)$ and $\mathcal{R}:=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.

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).

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. 3 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.

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.