SlowGuess/NeurIPS_2023_OCR huggingface dataset:
There are two key properties we use for the convergence of the base WGD algorithm: 1) a certain monotonicity of the update steps (up to higher order terms, gradient descent improves the objective) and 2) stability of gradients across iterations. If we include the BA step, we find that 1) still holds, but 2) may a-priori be lost. Indeed, 1) holds since BA updates monotonically improve the objective. Using just 1), we can still obtain a Pareto convergence of the gradients for the hybrid algorithm,  $\sum_{t=0}^{\infty} \gamma_t \int \|\nabla V_{\mathcal{L}}(\nu^{(t)})\|^2 d\nu^{(t)} < \infty$  (here  $\nu^{(t)}$  are the outputs from the respective BA
Table 1: Guide to notation and their interpretations in various problem domains. "LVM" stands for latent variable modeling, "NPMLE" stands for non-parametric maximum-likelihood estimation. The R-D problem (3) is equivalent to a projection problem under an entropic optimal transport cost (discussed in Sec. 2.2) and statistical problems involving maximum-likelihood estimation (see discussion in Sec. 2.3 and below).
Context	μ = PX	ρ(x,y)	K = QY|X	ν = QY
OT	source distribution	transport cost	“transport plan”	target distribution
R-D	data distribution	distortion criterion	compression algorithm	codebook distribution
LVM/NPMLE	data distribution	“- log p(x|y)”	variational posterior	prior distribution
deconvolution	noisy measurements	“noise kernel”	—	noiseless model
steps and  $\gamma_{t}$  is the step size of the gradient steps). Without property 2), we cannot conclude  $\int \| \nabla V_{\mathcal{L}}(\nu^{(t)})\| ^2 d\nu^{(t)}\to 0$  for  $t\to \infty$ . We emphasize that in practice, it still appears that 2) holds even after including the BA step. Motivated by this analysis, an adjusted hybrid algorithm where, e.g., the BA update is rejected if it causes a drastic change in the Wasserstein gradient, could guarantee that 2) holds with little practical changes. From a different perspective, we also believe the hybrid algorithm may be tractable to study in relation to gradient flows in the Wasserstein-Fisher-Rao geometry (cf. [Yan et al., 2023]), in which the BA step corresponds to a gradient update in the Fisher-Rao geometry with a unit step size.
In the stochastic setting, the BA (and therefore our hybrid) algorithm does not directly apply, as performing BA updates on mini-batches no longer guarantees monotonic improvement of the overall objective. Extending the BA and hybrid algorithm to this setting would be interesting future work.
10 R-D estimation and variational inference/learning
In this section, we elaborate on the connection between the R-D problem (3) and variational inference and learning in latent variable models, of which maximum likelihood deconvolution (discussed in Sec. 2.3) can be seen as a special case. Also see Section 3 of [Yang et al., 2023] for a related discussion in the context of lossy data compression.
To make the connections clearer to a general machine learning audience, we adopt notation more common in statistics and information theory. Table 1 summarizes the notation and the correspondence to the measure-theoretic notation used in the main text.
In statistics, a common goal is to model an (unknown) data generating distribution  $P_{X}$  by some density model  $\hat{p}(x)$ . In particular, here we will choose  $\hat{p}(x)$  to be a latent variable model, where  $\mathcal{V}$  takes on the role of a latent space, and  $Q_{Y} = \nu$  is the distribution of a latent variable  $Y$  (which may encapsulate the model parameters). As we shall see, the optimization objective defining the rate functional (5) corresponds to an aggregate Evidence LOwer Bound (ELBO) [Blei et al., 2017]. Thus, computing the rate functional corresponds to variational inference [Blei et al., 2017] in a given model (see Sec. 10.2), and the parametric R-D estimation problem, i.e.,
$$
\inf  _ {\nu \in \mathcal {H}} \mathcal {L} _ {B A} (\nu),
$$
is equivalent to estimating a latent variable model using the variational EM algorithm [Beal and Ghahramani, 2003] (see Sec. 10.3). The variational EM algorithm can be seen as a restricted version of the BA algorithm (see Sec. 10.3), whereas the EM algorithm [Dempster et al., 1977] shares its E-step with the BA algorithm but can differ in its M-step (see Sec. 10.4).
10.1 Setup
For concreteness, fix a reference measure  $\zeta$  on  $\mathcal{V}$ , and suppose  $Q_{Y}$  has density  $q(y)$  w.r.t.  $\zeta$ . Often the latent space  $\mathcal{V}$  is a Euclidean space, and  $q(y)$  is the usual probability density function w.r.t. the Lebesgue measure  $\zeta$ ; or when the latent space is discrete/countable,  $\zeta$  is the counting measure and  $q(y)$  is a probability mass function. We will consider the typical parametric estimation problem and choose a particular parametric form for  $Q_{Y}$  indexed by a parameter vector  $\theta$ . This defines a parametric family  $\mathcal{H} = \{Q_{Y}^{\theta} : \theta \in \Theta\}$  for some parameter space  $\Theta$ . Finally, suppose the distortion
function  $\rho$  induces a conditional likelihood density via  $p(x|y) \propto e^{-\lambda \rho(x,y)}$ , with a normalization constant that is constant in  $y$ .
These choices then result in a latent variable model specified by the joint density  $q(y)p(x|y)$ , and we model the data distribution with the marginal density:
$$
\hat {p} (x) = \int_ {\mathcal {Y}} p (x | y) d Q _ {Y} (y) = \int_ {\mathcal {Y}} p (x | y) q (y) \zeta (d y). \tag {22}
$$
As an example, a Gaussian mixture model with isotropic component variances can be specified as follows. Let  $Q_{Y}$  be a mixing distribution on  $\mathcal{X} = \mathcal{Y} = \mathbb{R}^{d}$  parameterized by component weights  $w_{1,\dots,k}$  and locations  $\mu_{1,\dots,k}$ , such that  $Q_{Y} = \sum_{k=1}^{K} w_{k} \delta_{\mu_{k}}$ . Let  $p(x|y) = \mathcal{N}(y, \sigma^{2})$  be a conditional Gaussian density with mean  $y$  and variance  $\sigma^{2}$ . Now formula (22) gives the usual Gaussian mixture density on  $\mathbb{R}^{d}$ .
Maximum-likelihood estimation then ideally maximizes the population log (marginal) likelihood,
$$
\mathbb {E} _ {x \sim P _ {X}} [ \log \hat {p} (x) ] = \int \log \hat {p} (x) P _ {X} (d x) = \int \log \left(\int_ {\mathcal {Y}} p (x | y) d Q _ {Y} (y)\right) P _ {X} (d x). \tag {23}
$$
The maximum-likelihood deconvolution setup can be seen as a special case where the form of the marginal density (22) derives from knowledge of the true data generation process, with  $P_{X} = \alpha * \mathcal{N}(0, \frac{1}{\lambda})$  for some unknown  $\alpha$  and known noise  $\mathcal{N}(0, \frac{1}{\lambda})$  (i.e., the model is well-specified). We note in passing that the idea of estimating an unknown data distribution by adding artificial noise to it also underlies work on spread divergence [Zhang et al., 2020] and denoising score matching [Vincent, 2011].
To deal with the often intractable marginal likelihood in the inner integral, we turn to variational inference and learning [Jordan et al., 1999, Wainwright et al., 2008].
10.2 Connection to variational inference
Given a latent variable model and any data observation  $x$ , a central task in Bayesian statistics is to infer the Bayesian posterior [Jordan, 1999], which we formally view as a conditional distribution  $Q_{Y|X = x}^{*}$ . It is given by
$$
\frac {d Q _ {Y \mid X = x} ^ {*} (y)}{d Q _ {Y} (y)} = \frac {p (x \mid y)}{\hat {p} (x)},
$$
or, in terms of the density  $q(y)$  of  $Q_{Y}$ , given by the following conditional density via the familiar Bayes' rule,
$$
q ^ {*} (y | x) = \frac {p (x | y) q (y)}{\hat {p} (x)} = \frac {p (x | y) q (y)}{\int_ {\mathcal {Y}} p (x | y) q (y) \zeta (d y)}.
$$
Unfortunately, the true Bayesian posterior is typically intractable, as the (marginal) data likelihood in the denominator involves an often high-dimensional integral. Variational inference [Jordan et al., 1999, Wainwright et al., 2008] therefore aims to approximate the true posterior by a variational distribution  $Q_{Y|X = x} \in \mathcal{P}(\mathcal{V})$  by minimizing their relative divergence  $H(Q_{Y|X = x}|Q_{Y|X = x}^{*})$ . The problem is equivalent to maximizing the following lower bound on the marginal log-likelihood, known as the Evidence Lower BOund (ELBO) [Blei et al., 2017]:
$$
\begin{array}{l} \operatorname *{arg  min}_{Q_{Y|X = x}}H(Q_{Y|X = x}|Q_{Y|X = x}^{*}) = \operatorname *{arg  max}_{Q_{Y|X = x}}\operatorname {ELBO}(Q_{Y},x,Q_{Y|X = x}), \\ \operatorname {E L B O} \left(Q _ {Y}, x, Q _ {Y | X = x}\right) = \mathbb {E} _ {y \sim Q _ {Y | X = x}} [ \log p (x | y) ] - H \left(Q _ {Y | X = x} \mid Q _ {Y}\right) \\ = \log \hat {p} (x) - H \left(Q _ {Y | X = x} \mid Q _ {Y | X = x} ^ {*}\right). \tag {24} \\ \end{array}
$$
Recall the rate functional (5) arises through an optimization problem over a transition kernel  $K$ ,
$$
\mathcal {L} _ {B A} (\nu) = \inf  _ {K} \lambda \int \rho d (\mu \otimes K) + H (\mu \otimes K | \mu \otimes \nu).
$$
To be more precise, we always fix a reference measure  $\eta$  on  $\mathcal{X}$ , so that densities such as  $\hat{p}(x)$  and  $p(x|y)$  are with respect to  $\eta$ . In the applications we considered in this work,  $\eta$  is the Lebesgue measure on  $\mathcal{X} = \mathbb{R}^d$ .
Translating the above into the present notation  $(\mu \to P_X, K \to Q_{Y|X}, \nu \to Q_Y$ ; see Table 1), we obtain
$$
\begin{array}{l} \mathcal {L} _ {B A} (Q _ {Y}) = \inf  _ {Q _ {Y | X}} \mathbb {E} _ {x \sim P _ {X}, y \sim Q _ {Y | X = x}} [ - \log p (x | y) ] + \mathbb {E} _ {x \sim P _ {X}} [ H (Q _ {Y | X = x} | Q _ {Y}) ] + \text {c o n s t} \\ = \inf  _ {Q _ {Y \mid X}} \mathbb {E} _ {x \sim P _ {X}} [ - \operatorname {E L B O} \left(Q _ {Y}, x, Q _ {Y \mid X = x}\right) ] + \text {c o n s t .} \tag {25} \\ \end{array}
$$
which allows us to interpret the rate functional as the result of performing variational inference through ELBO optimization. At optimality,  $Q_{Y|X} = Q_{Y|X}^{*}$ , the ELBO (24) is tight and recovers  $\log \hat{p}(x)$ , and the rate functional takes on the form of a (negated) population marginal log likelihood (23), as given earlier by (6) in Sec. 2.1.
10.3 Connection to variational EM
The discussion so far concerns probabilistic inference, where a latent variable model  $(Q_{Y},p(x|y))$  has been given and we saw that computing the rate functional amounts to variational inference. Suppose now we wish to learn a model from data. The R-D problem (4) then corresponds to model estimation using the variational EM algorithm [Beal and Ghahramani, 2003].
To estimate a latent variable model by (approximate) maximum-likelihood, the variational EM algorithm maximizes the population ELBO
$$
\mathbb {E} _ {x \sim P _ {X}} \left[ \operatorname {E L B O} \left(Q _ {Y}, x, Q _ {Y \mid X = x}\right) \right] = \mathbb {E} _ {x \sim P _ {X}, y \sim Q _ {Y \mid X = x}} \left[ \log p (x \mid y) \right] - \mathbb {E} _ {x \sim P _ {X}} \left[ H \left(Q _ {Y \mid X = x} \mid Q _ {Y}\right) \right], \tag {26}
$$
w.r.t.  $Q_{Y}$  and  $Q_{Y|X}$ . This precisely corresponds to the R-D problem  $\inf_{Q_{Y} \in \mathcal{H}} \mathcal{L}_{BA}(Q_{Y})$ , using the form of  $\mathcal{L}_{BA}(Q_{Y})$  from (25).
In popular implementations of variational EM such as the VAE [Kingma and Welling, 2013],  $Q_{Y}$  and  $Q_{Y|X}$  are restricted to parametric families. When they are allowed to range over all of  $\mathcal{P}(\mathcal{V})$  and all conditional distributions, variational EM then becomes equivalent to the BA algorithm.
10.4 The Blahut-Arimoto and (exact) EM algorithms
Here we focus on the distinction between the BA algorithm and the (exact) EM algorithm [Dempster et al., 1977], rather than the variational EM algorithm. Both BA and (exact) EM perform coordinate descent on the same objective function (namely the negative of the population ELBO from (26)), but they define the coordinates slightly differently — the BA algorithm uses  $(Q_{Y|X}, Q_Y)$  with  $Q_Y \in \mathcal{P}(\mathcal{V})$ , whereas the EM algorithm uses  $(Q_{Y|X}, \theta)$  with  $\theta$  indexing a parametric family  $\mathcal{H} = \{Q_Y^\theta : \theta \in \Theta\}$ . Thus the coordinate update w.r.t.  $Q_{Y|X}$  (the "E-step") is the same in both algorithms, but the subsequent "M-step" potentially differs depending on the role of  $\theta$ .
Given the optimization objective, which is simply the negative of (26),
$$
\mathbb {E} _ {x \sim P _ {X}, y \sim Q _ {Y | X = x}} [ - \log p (x | y) ] + H \left(P _ {X} Q _ {Y | X} \mid P _ {X} \otimes Q _ {Y}\right), \tag {27}
$$
both BA and EM optimize the transition kernel  $Q_{Y|X}$  the same way in the E-step, as
$$
\frac {d Q _ {Y \mid X = x} ^ {*}}{d Q _ {Y}} (y) = \frac {p (x \mid y)}{\hat {p} (x)}. \tag {28}
$$
For the M-step, the BA algorithm only minimizes the relative entropy term of the objective (27),
$$
\min  _ {Q _ {Y} \in \mathcal {P} (\mathcal {Y})} H (P _ {X} Q _ {Y | X} ^ {*}; P _ {X} \otimes Q _ {Y}),
$$
(with the optimal  $Q_{Y}$  given by the second marginal of  $P_{X}Q_{Y|X}^{*}$ ) whereas the EM algorithm minimizes the full objective w.r.t. the parameters  $\theta$  of  $Q_{Y}$ ,
$$
\min  _ {\theta \in \Theta} \mathbb {E} _ {(x, y) \sim P _ {X} Q _ {Y | X} ^ {*}} [ - \log p (x | y) ] + H \left(P _ {X} Q _ {Y | X} ^ {*}; P _ {X} \otimes Q _ {Y}\right). \tag {29}
$$
The difference comes from the fact that when we parameterize  $Q_{Y}$  by  $\theta$  in the parameter estimation problem,  $Q_{Y|X}^{*}$  — and consequently both terms in the objective of (29) — will have functional dependence on  $\theta$  through the E-step optimality condition (28).
In the Gaussian mixture example,  $Q_{Y} = \sum_{k=1}^{K} w_{k} \delta_{\mu_{k}}$ , and its parameters  $\theta$  consist of the components weights  $(w_{1}, \ldots, w_{K}) \in \Delta^{d-1}$  and location vectors  $\{\mu_{1}, \ldots, \mu_{K}\} \subset \mathbb{R}^{d}$ . The E-step computes  $Q_{Y|X=x}^{*} = \sum_{k} w_{k} \frac{p(x|\mu_{k})}{p(x)} \delta_{\mu_{k}}$ . For the M-step, if we regard the locations as known so that  $\theta = (w_{1}, \ldots, w_{K})$  only consists of the weights, then the two algorithms perform the same update; however if  $\theta$  also includes the locations, then the M-step of the EM algorithm will not only update the weights as in the BA algorithm, but also the locations, due to the distortion term  $\mathbb{E}_{(x,y)} \sim P_{X} Q_{Y|X}^{*}[-\log p(x|y)] = -\int \sum_{k} w_{k} \frac{p(x|\mu_{k})}{p(x)} \log p(x|\mu_{k}) P_{X}(dx)$ .
11 Further experimental results
Our deconvolution experiments were run on Intel(R) Xeon(R) CPUs, and the rest of the experiments were run on Titan RTX GPUs.
In most experiments, we use the Adam [Kingma and Ba, 2015] optimizer for updating the  $\nu$  particle locations in WGD and for updating the variational parameters in other methods. For our hybrid WGD algorithm, which adjusts the particle weights in addition to their locations, we found that applying momentum to the particle locations can in fact slow down convergence, and therefore use plain gradient descent with a decaying step size.
To trace an R-D upper bound for a given source, we solve the unconstrained R-D problem 3 for a heuristically-chosen grid of  $\lambda$  values similarly to those in [Yang and Mandt, 2022], e.g., on a log-linear grid  $\{1e4, 3e3, 1e3, 3e2, 1e2, \ldots\}$ . Alternatively, a constrained optimization method like the Modified Differential Multiplier Method (MDMM) [Platt and Barr, 1987] can be adopted to directly target  $R(D)$  for various values of  $D\mathrm{s}$ . The latter approach will also help us identify when we run into the  $\log(n)$  rate limit of particle-based methods (Sec. 4.4): suppose we perform constrained optimization with a distortion threshold of  $D$ ; if the chosen  $n$  is not large enough, i.e.,  $\log(n) < R(D)$ , then no  $\nu$  supported on (at most)  $n$  points is feasible for the given constraint (otherwise we have a contradiction). When this happens, the Lagrange multiplier associated with the infeasible constraint ( $\lambda$  in our case) will be observed to increase indefinitely rather than stabilize to a finite value under a method like MDMM.
11.1 Deconvolution
Architectures for the neural baselines For the RD-VAE, we used a similar architecture as the one used on the banana-shaped source in [Yang and Mandt, 2022], consisting of two-layer MLPs for the encoder and decoder networks, and Masked Autoregressive Flow [Papamakarios et al., 2017] for the variational prior. For NERD, we follow similar architecture settings as [Lei et al., 2023a], using a two-layer MLP for the decoder network. Following Yang and Mandt [2022], we initially used softmax activation for the MLP, but found it made the optimization difficult; therefore we switched to ReLU activation which gave much better results. We conjecture that the softmax activation led to overly smooth mappings, which had difficulty matching the optimal  $\nu^{*} = \alpha$  measure which is concentrated on the unit circle and does not have a Lebesgue density.
Experiment setup As discussed in Sec. 4.2, BA and our hybrid algorithms do not apply to the stochastic setting; to be able to include them in our comparison, the input to all the algorithms is an empirical measure  $\mu^m$  (training distribution) with  $m = 100000$  samples, given the same fixed seed. We found the number of training samples is sufficiently large that the losses do not differ significantly on the training distribution v.s. random/held-out samples.
Recall from Sec. 2.3, given data observations following  $\mu = \alpha * \mathcal{N}(0, \sigma^2)$ , if we perform deconvolution with an assumed noise variance  $\frac{1}{\lambda}$  for some  $\frac{1}{\lambda} \leq \sigma^2$ , then the optimal solution to the problem is given by  $\nu^{*} = \alpha_{\lambda} = \alpha * \mathcal{N}(0, \frac{1}{\lambda})$ . We compute the optimal loss  $OPT = \mathcal{L}_{BA}(\nu^{*})$  by a Monte-Carlo estimate, using formula (19) with  $m = 10^4$  samples from  $\mu$  and  $n = 10^6$  samples from  $\nu^{*}$ . The resulting  $\hat{\mathcal{L}}_{BA}$  is positively biased (it overestimates the  $OPT$  in expectation) due to the bias of the plug-in estimator for  $\varphi(x)$  (18), so we use a large  $n$  to reduce this bias. Note the same issue occurs in the NERD algorithm (13).
Another, more sophisticated solution would be annealed importance sampling [Grosse et al., 2015] or a related method developed for estimating marginal likelihoods and partition functions.
Figure 6: Visualizing the converged  $n = 20$  particles of WGD (top row) and hybrid WGD (bottom row) estimated on  $m = 10000$  source samples in the deconvolution problem (Sec. 5.1), for decreasing distortion penalty  $\lambda$ . The case  $\lambda = 10.0 = \sigma^{-2}$  corresponds to "complete noise removal" of the source and recovers the ground truth  $\alpha$  (unit circle).
Optimized reproduction distributions To visualize the (continuous)  $\nu$  learned by RD-VAE and NERD, we fit a kernel density estimator on 10000  $\nu$  samples drawn from each, using seaborn.kdeplot.
In Figure 6, we provide additional visualization for the particles estimated from the training samples by WGD and its hybrid variant, for a decreasing distortion penalty  $\lambda$  (equivalently, increasing entropic regularization strength  $\epsilon$ ).
11.2 Higher-dimensional datasets
For NERD, we set  $n$  to the default 40000 in the code provided by [Lei et al., 2023a], on all three datasets.
For WGD, we used  $n = 4000$  on physics, 200000 on speech, and 40000 on MNIST (see also smaller  $n$  for comparison in Fig. 5).
On speech, both NERD and WGD encountered the issue of a  $\log (n)$  upper bound on the rate estimate as described in Sec. 4.4. Therefore, we increased  $n$  to 200000 for WGD and obtain a tighter R-D upper bound than NERD, particularly in the low-distortion regime. Such a large  $n$  incurred out-of-memory error for NERD.
We borrow the R-D upper and lower bound results of [Yang and Mandt, 2022] from their paper on physics and speech. We obtain sandwich bounds on MNIST using a similar neural network architecture and other hyperparameter choices as in the GAN-generated image experiments of [Yang and Mandt, 2022] (using a ResNet-VAE for the upper bound and a CNN for the lower bound), except we set a larger  $k = 10000$  in the lower bound experiment; we suspect the resulting lower bound is still far from being tight.
12 Example implementation of WGD
We provide a self-contained minimal implementation of Wasserstein gradient descent on  $\mathcal{L}_{BA}$ , using the Jax library [Bradbury et al., 2018]. To compute the Wasserstein gradient, we evaluate the first variation of the rate functional in compute_psi_sum according to (20), yielding  $\sum_{i=1}^{n} \psi^{\nu}(x_i)$ , then simply take its gradient w.r.t. the particle locations  $x_1, \ldots, n$  using Jax's autodiff tool on line 51.
The implementation of WGD on  $\mathcal{L}_{EOT}$  is similar, except the first variation is computed using Sinkhorn's algorithm. Both versions can be easily just-in-time compiled and enjoy GPU acceleration.
Wasserstein GD on the rate functional L.{BA}. import jax numpy as jnp import jax from jax.scipy_special import logsumexp   
# Define the distortion function \rho. squared_diff = lambda x, y: jnp.sum((x - y) ** 2) pairwise_distortion_fn = jax.vmap(jax.vmap(squared_diff, (None, 0)), (0, None))   
def wgrad(mu_x, mu_w, nu_x, nu_w, rd_lambda): " Compute the Wasserstein gradient of the rate functional, which we will use to move the \nu particles. :param mu_x: locations of \mu atoms. :param mu_w: weights of \mu atoms. :param nu_x: locations of \nu atoms. :param nu_w: weights of \nu atoms. :param rd_lambda: R-D tradeoff hyperparameter. :return: "   
def compute_psi_sum(nu_x): "Here we compute a surrogate loss based on the first variation \psi i, which allows jax autodiff to compute the desired Wasserstein gradient. :param nu_x: :return: psi_sum = \sum_i | psi(nu_x[i]) " C = pairwise_distortion_fn(mu_x, nu_x) scaled_C = rd_lambda * C # [m, n] log_nu_w = jnp.log(nu_w) # [1, n] # Solve BA inner problem with a fixed nu. phi = - logsumexp(-scaled_C + log_nu_w, axis=1, keepdims=True) # [m, 1] loss = jnp-sum(mu_w * phi) # Evaluate the rate functional. Eq (6) in paper # Let's also report rate and distortion estimates (discussed in Sec. 4.4 of # Find  $\backslash p i ^ { \text{一} }$  via  $\backslash$  phi pi = jnp.exp(phi - scaled_C) * jnp外出(mu_w, nu_w) # [m, n] distortion = jnp-sum(pi * C) rate = loss - rd_lambda * distortion # Now evaluate  $\backslash$  psi on the atoms of  $\backslash$  nu. phi = jax.lax.stop_gradients(phi) psi = - jnp-sum(jnp.exp(jax.lax.stop_gradients(phi) - scaled_C) * mu_w, axis psi_sum = jnp-sum(psi) # For computing gradient w.r.t. each nu_x atom. return psi_sum, (loss, rate, distortion) # Evaluate the Wasserstein gradient, i.e.,  $\backslash$  nabla  $\backslash$  psi, on nu_x. psi_prime, loss = jax.grad(compute_psi_sum, has(aux=True)(nu_x) return psi_prime, loss   
def wgd(X, n, rd_lambda, num_steps, lr,rng): A basic demo of Wasserstein gradient descent on a discrete distribution. :param X: a 2D array [N, d] of data points defining the source  $\backslash$  mu. :param n: the number of particles to use for  $\backslash$  nu.
1   
2   
3   
4   
5   
6   
7   
8   
9   
10   
11   
12   
13   
14   
15   
16   
17   
18   
19   
20   
21   
22   
23   
24   
25   
26   
27   
28   
29   
30   
31   
32   
33   
34   
35   
36   
37   
38   
39   
40   
41   
42   
43   
44   
45   
46   
47   
48   
49   
50   
51   
52   
53   
54   
55   
56   
57   
58   
59
:param rd_lambda: R-D tradeoff hyperparameter.
:param num_steps: total number of gradient updates.
:param lr: step size.
:paramrng:jaxrandomkey.
:return:  $(nu_{-}x,nu_{-}w)$  , the locations and weights of the final  $\backslash nu$
Set up the source measure \mu.
$\mathfrak{m} = \mathfrak{j}\mathfrak{n}\mathfrak{p}$  .size(X，0)
$\mathbf{mu\_x} = \mathbf{X}$
$\mathrm{mu\_w} = 1 / \mathrm{m} * \mathrm{jnp\_ones}((\mathrm{m}, 1))$
Initialize \nuu atoms using random training samples.
randidx =jax.random.permutation(rng, m)[:n]
nu_x = X[randidx] # Locations of \nu atoms.
nu_w = 1 / n * jnp.ones((1, n)) # Uniform weights.
for step in range(num_steps):
psi_prime, (loss, rate, distortion) = wgrad(mu_x, mu_w, nu_x, nu_w, rd_lambda)