Estimating the Rate-Distortion Function by Wasserstein Gradient Descent
Authors: Yibo Yang
Stephan Eckstein
Marcel Nutz
Stephan Mandt
Venue: NeurIPS 2023 poster (2023)
Keywords: information theory
rate-distortion function
optimal transport
TL;DR: We develop a new approach for estimating the rate-distortion function from the perspective of optimal transport.

Abstract: In the theory of lossy compression, the rate-distortion (R-D) function $R(D)$ describes how much a data source can be compressed (in bit-rate) at any given level of fidelity (distortion). Obtaining $R(D)$ for a given data source establishes the fundamental performance limit for all compression algorithms.  We propose a new method to estimate $R(D)$ from the perspective of optimal transport. Unlike the classic Blahut--Arimoto algorithm which fixes the support of the reproduction distribution in advance, our Wasserstein gradient descent algorithm learns the support of the optimal reproduction distribution by moving particles. We prove its local convergence and analyze the sample complexity of our R-D estimator based on a connection to entropic optimal transport.  Experimentally, we obtain comparable or tighter bounds than state-of-the-art neural network methods on low-rate sources while requiring considerably less tuning and computation effort.  We also highlight a connection to maximum-likelihood deconvolution and introduce a new class of sources that can be used as test cases with known solutions to the R-D problem.

Estimating the Rate-Distortion Function by Wasserstein Gradient Descent
Yibo Yang'
Stephan Eckstein² Marcel Nutz
Stephan Mandt
University of California, Irvine 2ETH Zurich
SColumbia University
(yibo.yare mandt) edu
stephan. eckstein@math ethz. ch mnutz@columbia. edu
Abstract
the theory of lossy compression, the rate-distortion (R-D) function R(D) de- scribes how much data source can be compressed (in bit-rate) any given level fidelity (distortion). Obtaining RDD) for given data source establishes the fundamental performance limit for all compression algorithms. We propose new method to estimate R from the perspective of optimal transport. Unlike the classic Blahut-Arimoto algorithm which fixes the support of the reproduction dis- tribution in advance, our Wasserstein gradient descent algorithm learns the support of the optimal reproduction distribution by moving particles. We prove its local convergence and analyze the sample complexity of our R-D estimator based on a obtain comparable
connection to entropic optimal transport. Experimentally, tighter bounds than state-of-te of-thc-art neural network methods on low-rate sources while requiring considerably less tuning and computation effort. We also highlight a connection to maximum maximum-likelibon likelihood deconvolution and introduce a new class of sources that can be used test cases with known solutions to the R-D problem.
Introduction
The rate-distortion (R-D) function R(D) occupies a central place in the theory of lossy compression. For given data source and a fidelity (or distortion) criterion, R(D) characterizes the minimum possible communication cost needed to reproduce a source sample within an error threshold of D, by any compression algorithm [Shannon, 1959]. A basic scientific and practical question therefore establishing R(D) for any given data source of interest, which helps assess the (sub)optimality of the compression algorithms and guide their development. The classic algorithm by Blahut [1972] and Arimoto [1972] assumes a known discrete source and computes its R(D) by an exhaustive optimization procedure. This often has limited applicability in practice, and line of research has sought to instead estimate R( D) from data samples [Harrison and Kontoyiannis, 2008, Gibson, 2017 7], with recent methods (Yang and Mandt, 2022, Lei al., 2023a] inspired by deep generative models.
In this work, we propose a new approach to R-D estimation from the perspective of optimal transport. Our starting point the formulation of the R-D problem as the minimization of certain rate func- tional [Harrison and Kontoyiannis, 2008] over the space of probability measures on the reproduction alphabet. Optimization over such an simittediimesiioaa space has long been studied under gradient flows [Ambrosio et al., 2008], and we consider a concrete algorithmic implementation based on moving particles space. This formulation of the R-D problem also suggests connections to entropic optimal transport and non-parametric statistics, each offering us new insight into the solution of the R-D problem under a quadratic distortion. More specifically, our contributions are three-fold:
First, we introduce a neural-networkk-fre R(D) upper bound estimator for continuous alphabets. We implement the estimator by Wasserstein gradient descent (WGD) over the space of reproduction
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

distributions. Experimentally, we found the method to converge much more quickly than state-of-the art neural methods with hand-tuned architectures, while offering comparable or tighter bounds.
Second, we theoretically characterize convergence of our WGD algorithm and the sample complexity of our estimator. The latter draws connection between the R-D problem and that of minimizing an entropic optimal transport (EOT) cost relative to the source measure, allowing us to turn statistical bounds for EOT [Mena and Niles-Weed, 2019] into finite-sample bounds for R-D estimation.
Finally, we introduce a new, rich class of sources with known ground truth, including Gaussian mixtures, benchmark for algorithms. While the literature relies on the Gaussian Bernoulli for this purpose, we use the connection with maximum likelihood deconvolution to show that a Gaussian convolution of any distribution can serve as a souree with a known solution to the R-D problem.
2 Lossy compression, entropic optimal transport, and MLE
This section introduces the R-D problem and its rich connections to entropic optimal transport and statistics, along with new insights into its solution. Sec. 2.1 sets the stage for our method (Sec. 4) by known formulation of the standard R-D problem as an optimization problem over a space of probability measures. Sec. 2.2 discusses the equivalence between the R-D problem and projection of the source distribution under entropic optimal transport; this is a key to our sample complexity results in Sec. 4.3. Lastly, See 2.3 gives a statistical interpretation of R-D as maximum-likelihon deconvolution and uses to analytically derive a segment of the R-D curve for a new class of sources under quadratic distortion; this allows us to assess the optimality of algorithms in experiment Sec. 5.1.
2.1 Setup
For memoryless data source X with distribution Px, 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 X' and Y known as the source and reproduction alphabets, and let p: (X) - [0, oo) be a given distortion function. The R-D function is defined by the following optimization problem [Polyanskiy and Wu, 2022],
inf R(D) Qrix ( (X( I(:
(1)
where Qrix any Markov kernel from X' to y conceptually associated with a (possibly) stochastic compression algorithm, and I [(;; Y) is the mutual information of the joint distribution Px PxQyxx-
For case 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 X' and y be standard Borel spaces; let μ E P(2') be fixed probability measure on X' which should be thought of as the source distribution Px. For measure on the product space X' V, the notation T. (or T2) denotes the first (or second) marginal of т. For any V E P()).. we denote by II v/ the set of couplings between μ and (i.e., T1 f and πΣ = v). Similarly, [I,,,) denotes the set of measures with I = J.. Throughout the paper, K denotes a transition kernel (conditional distribution) from X' to }, and μ & K denotes the product measure formed by μ and K. Then R(D) is equivalent
inf R(D)= K:f p)(
inf H & Klu 8 (μ @ K)2) *CT((g)): pda
(T| 8 m2),
(2)
where H denotes relative entropy, i.e., for two measures a, 3 defined on a common measurable space, H (alB)) = log( do dd when or is absolutely continuous w.r.t 8, and infinite otherwise.
dp
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 Lagrange multiplier À 2 0. For each fixed À (usually selected from predefined grid), we aim to solve the following optimization problem,
PEP()) rell(a.) > f
f pdx + H( & v).
F (x) == inf inf
(3)
Geometrically, (x) E R is the y-axis intercept of a tangent line to the R(D) with slope -X, 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 À (c.g., we write F(R) Fx(e)) whenever it is harmless.

---

To set the stage for our later developments, we write the unconstrained R-D problem as inf (BA(U,D), vEP())
Fx(1)
(4)
A Dd () - nell(u.)
(BA(H,)) =
K -
inf
inf pdlp & K) + H(u & 0))
(5)
where we refer to the optimization objective LBA the rate function [Harrison and Kontoyiannis, 2008]. We abuse the notation to write (B)(v) := CBA(U, v) when it is viewed as a function of only, and refer to it as the rate functional. The rate function characterizes a generalized Asymptotic Equipartition Property, where LBA(4 (1,, v) is the asymptotically optimal cost of lossy compression of data X ~ y using random codebook constructed from samples of v [Dembo and Kontoyiannis, 2002]. Notably, the optimization in (5) can be solved analytically [Csiszár, 1974a, Lemma 1.3), and CA simplifies to
(RA(KU)= - -log (J, e v(dr).
(6)
In practice, the source i is only accessible via independent samples, on the basis of which we propose estimate its R(D), or equivalently F(u). Let μr" denote an m-sample empirical measure of μ,, i.e.
1 n² = - 8z with 2 11..... being independent samples from which should be thought of as the "training data". Following Harrison and Kontoyiannis [2008], we consider two kinds of (plug-in) estimators for F(z): (1) the non-pat non-paraeetric estimator F F(n"),, and (2) the parametric estimator F ("""")) == infrem LBA(μ"" v). where is family of probability measures on V. Harrison and strongly Kontoyiannis [2008] showed that under rather broad conditions, both kinds of estimators
consistent, i.c., F(4")) converges to F(u) (and respectively, (="(""") to F* (2)) with probability one as + co. Our algorithm will implement the parametric estimator FR F*(""") with H chosen to be the set 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).
2.2 Connection to entropic optimal transport
The R-D problem turns out have a close connection 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 0, the entropy- intropy-regularized optimal transport problem is given by
LEOT (H, v) == inf
meIl(..)) _ pdz + eH (T| @ v).
We now consider the problem of projecting μ onto P(Y)) under the cost LBTT:
inf LEOT (p, (p, v). MEP()) In the OT literature this is known as the (regularized) Kantorovich estimator [Bassetti et al., 2006] for
(8)
μr, and can also be viewed as a Wasserstein barycenter problem [Agueh and Carlier, 2011].
With the identification
X-1, problem (8) is in fact equivalent the R-D problem (4): compared to LA (5), the extra constraint on the second marginal of π in LOOT (7) is redundant at the optimal v. More precisely, Lemma 7.1 shows that (we omit the notational dependence on μ when it is fixed): inf LBOT(V)= = inf X-1LBA(V) vEP())
and
arg min LEOT = arg min (V). (9) vEP()'
vEP())
NEP())
Existence of a minimizer holds under mild conditions, for instance if X' = y = Rd and p(z, y) is coercive lower semicontinuous function of y [Csiszár, 1974a, p. 66).
2.3 Connection to maximum-likelihowd deconvolution
The connection between R-D and Maximum--ikelihon estimation has been observed in the in- formation theory, machine learning and compression literature [Harrison and Kontoyiannis, 2008, Alemi al., 2018, Ballé al., 2017, Theis al., 2017, Yang et al., 2020, Yang and Mandt, 2022]. Here, we bring attention to basic equivalence between the R-D problem and deconvolution, where the connection is particularly natural under a quadratic distortion function. Also see [Rigollet and Weed, 2018) 8] for 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.
3

---

Maximum-likelihoud deconvolution is a classical problem of non- non-parametric statistics and mixture models [Carroll and Hall, 1988, Lindsay and Roeder, 1993]. The deconvolution problem is concerned with estimating an unknown distribu- tion from noise-corrupted observations X1, X2...., where for each i E N, we have X; = Y; + N;, V; s.i.d. ~ a, and N; i.i.d. independent noise variables with known distribu- tion. For concreteness, suppose all variables are E valued and the noise distribution is (0, Ia with Lebesgue den- sity 0g2. Denote the distribution of the observations X; by μ. Then μ has a Lebesgue density given by the convo- lution a * 0,² := Φ, (a: y)) ))((d))) Here, we con- sider the population- lelvel (instead of the usual sample-based) maximum likelihood estimator (MLE) for a:
R(D)
decreasing À A
Figure 1: The R(D) of Gaussian mixture source, and the estimated op- timal reproduction distributions v* (in bar plots) at varying R-D trade- offs. For any X E ,0), the cor- responding R(D) (yellow segment) is known analytically
/
V¹ arg max WEP(RR)
log (v (z))
(10)
and observe that V² = a. Plugging in the density 6.2 (z) x z=2 ||=]|? see that the MLE problem (10) is equivalent +. and L. given by (6) in the form of marginal log-likelihood. Thus the R-D problem has the interpretation of estimating distribution from its noisy observations given through ja,
the optimal reproduction distribution V/ (whose density plotted in gray). For À E (0, a 00,0-2), v" becomes singular and concentrated on two points, collaps- ing to the source mean À + O.
the R-D problem (4) with p(z, y) = 2 || -yll,,
assuming a Gaussian noise with variance -.
This connection suggests analytical solutions to the R-D problem for variety of sources that arise from convolving an underlying distribution with Gaussian noise. Consider an R-D problem (4) with X' = Rd p(x,, y) Iæ y/l?, and let the source μ be the convolution between an arbitrary measure E P())) and Gaussian noise with known variance o². E.g., using discrete measure for results in Gaussian mixture source with equal covariance among its components. When À = = we recover exactly the population-MLE problem (10) discussed earlier, which has the solution v" a. While this allows us to obtain one point of RCD), we can in fact extend this idea to any À 2 + and obtain the analytical form for the corresponding segment of the R-D curve. Specifically, for any A 2 2.. applying the summation rule for independent Gaussians reveals the source distribution j as
1 ) { as NN. }) := * Nx NN }
# = a + N = a NN o² -
ie.. as the convolution between another underlying distribution as and independent noise with variance 1. A solution to the R-D problem (3) is then analogously given by V" = ax with the corresponding optimal coupling given by D² @ K. K(y, dx) = =N(y, })) Evaluating the distortion and mutual information of the coupling then yields the R(D) point associated with A. Fig. I illustrates the R(D) of toy Gaussian mixture source, along with the 1² estimated by our proposed WGD algorithm (Sec. 4); note that n* transitions from continuous (a Gaussian mixture with smaller component variances) to singular (a mixture of two Diracs) x o 2. See caption for more details.
3
Related Work
3.1
Blahut-Arimoto
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 X, BA carries out the optimization problem (3) via coordinate ascent. Starting from an initial measure (0)) E P())), the BA algorithm step / computes an updated pair (u(++1), K(++)) as follows
dx(++)(xx,.)) de(t)
c Ap(a,w)
(y) J Te (o-
Vr Y X,
(11)
((+1) = (n 8 (((+1))2 2-
(12)
Herr K maps from the reproduction the source alphabet. opposite the kernel K elsewhere in the text.

---

When the alphabets are finite, the above computation can be carried out in matrix and vector operations, and the resulting sequence { K [0]= can be shown to converge to an optimum of (3); cf. [Csiszár, 1974b, Csiszár, 1984]. When the alphabets are not finite, e.g., A' = ) = IR the BA algorithm no longer applies, as it is unclear how to digitally represent the measure and kernel K' and to tractably perform the integrals required by the algorithm. The common workaround is to perform discretization step and then apply BA 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 X' to be the support of training data distribution μ² i.c., the discretized alphabet is the set of training samples; this can be justified by the consistency of the parametric R-D estimator JzH (uT)) [Harrison and Kontoyiannis, 2008]. is less clear how to discretize the reproduction space V, especially in high dimensions. Since work with X' = y/ we will disretize 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 results in poor performance, and can be significantly improved from our perspective of evolving particles.
3.2 Neural network-based methods for estimating R(D)
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
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,, Y3) ~ 11 @ K; in particular, the relative entropy can be written as Hp & K|u @ v) =
[log (K(x.)) (V) K (xx, dy) ((dr), where the integrand is computed exactly via a density ratio. In du
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 neural estimator for 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 parameterized by neural network. Let غ be base distribution over 2 = R such as the standard Gaussian, and w 2+ y be a decoder network. The variational measure v is then modeled as the image measure of ~ under To evaluate and optimize the objective (6), the intractable inner integral w.r.t. v is replaced with a plug-in estimator, so that for a given & E X'
-log (Ls e (-Aolaw)y(dd)) x loo
Σ
Y, ~ v.j = 1,2, n. (13)
After training, we estimate an R-D upper bound using samples from v (to be discussed in Sec. 4.4).
3.3 Other related work
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) discrete problem and discuss the connection to scalar quantization. We also experimented with estimating R( () 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 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-like likelihoo deconvolution (10); our work complements their perspective by re- 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- Saaee-RRo-Waaseeee Wasserstein (FRW) geometry [Chizat et al., 2018 8), which our hybrid algorithm can be seen as implementing. Yan et al. [2023] prove that, in an
5

---

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 = = IRd and p 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 L P()) - T& over the space of probability measures, algorithm essentially simulates the gradient flow [Ambrosio et al., 2008] of and follows the trajectory of steepest descent in the Wasserstein geometry. In practice, we represent measure v() E P))) by a collection of particles and each time step update y(t) in a direction of steepest descent of L as given by its (negative) Wasserstein gradient. Denote by Pn, (IR") the set of probability measures IR that are supported at most n points. Our algorithm implements the parametric R-D estimator with the choice H = Pn, (Rd) (see discussions at the end of Sec. 2.1).
4.1 Wasserstein gradient descent (WGD)
Abstractly, Wasserstein gradient descent updates the variational measure to its pushforward i under the map (id yV), for function V IRd -> Rdd called the Wasserstein gradient of L at v (see below) and a step size y. To implement this scheme, represent Dirac measures, // _i=1 w.da, with locations {xr c IK' and weights {w; h= The algorithm moves each particle r; the direction of (x:),, more precisely, Ei=1 =;
convex combination of
Algorithm Wasserstein gradient descent
Inputs: Loss function {CBA, EOTT) ); data distribution μ E P(RR)); the number of particles E N; total number of iterations N E N; step sizes 7....., YN; batch size m E N. for t ,,...,N do
Pick an initial measure n0 E Pa (R"), e.g., setting the particles to n random samples from p. if support of contains more than m points then
nT < m Σi 82, for Ti...., I' independent samples from p
(1()) + Wasserstein gradient C(p") at ((t-1) (see Definition 4.1
else
(t()) + Wasserstein gradient of L(j,,) at (t-1) (see Definition 4.1)
end if
elt) + (id -
vid->volo) v(l)))) (l-1) ("#"" denotes pushforward)
end for Return: ((N)
Since the optimization objectives (4) and (8) appear integrals w.r.t. the data distribution μ4, can also apply stochastic optimization and perform stochastic gradient descent on mini-batches with size m. This allows us handle very large or infinite amount of data samples, when the source continuous. We formalize the procedure in Algorithm 1.
The following gives constructive definition of a Wasserstein gradient which forms the computational basis of our algorithm. In the literature, the Wasserstein gradient instead usually defined as a Fréchet differential (cf. [Ambrosio et al., 2008, Definition 10)1.I]), 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 functional L: P())) lR and E P(Y), we say that Vz(v): Rdd IR is first variation of L at if
lim C((1 E)V + E -())) / Ve(l) d(ū vv for all & E P(Y).
We call its (Euclidean) gradiens VVc(v) Rd - RR, if exists, the Wasserstein gradient of L at
For L = LEOT, the first variation is given by the Kantorovich potential, which the solution of the convex dual of LEOT and commonly computed by Sinkhorn's 's algorithm [Peyré and Cuturi, 2019,

---

Nutz, 2021]. Specifically, let (4", 2/4) be potentials for LEOT(H, v). Then Ve(v)= 6 is the first variation w.r.t v (cf. [Carlier et al., 2022, equation (20))), and hence V the Wasserstein gradient. This gradient exists whenever p is differentiable and the marginals are sufficiently light-tailed; we give details in Sec. 9.1 of the Supplementary Material. For f LB the first variation can be computed explicitly. As derived in Sec. 9.1 of the Supplementary Material, the first variation at is
exp( (xp(-Ap(x,y))
e (v) (0)=--1
exp( sexp( (xp(- (ddj) u(d))
and then the Wasserstein gradient is VLBA(v) V.y. We observe that 3 (v) is computationally cheap; it corresponds to rumning single iteration of Sinkhorm's algorithm. algaooith By contract, finding the potential for LEOT requires running Sinkhorn's 's algorithm to convergence.
Like the usual Euclidean gradient, the Wasserstein gradient can be shown to possess linearization property, whereby the loss functional is reduced by taking a small enough step along its Wasserstein gradient. Following [Carlier et al., 2022], state it as follows: for any P(D)) and i E (),
[(D) - C(v)- -
-w- x)) VVVV((()) ((xx, du) + (fll= =112=112 (dx, du)
(14)
/ /IVVA(W)Pas dv - / [VV22 di (W).
The first line of (14) is proved for LEOT [Carlier al., 2022, Proposition 4.2] in the case that the marginals are compactly supported and p is twice continuously differentiable. In this setting, the second line of (14) follows using a² - 82 = (a + (a b) and a combination of boundedness and Lipschitz continuity of VVc. 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 LOT and LA converges to a stationary point under mild conditions:
Proposition 4.2 (Convergence of Wasserstein gradient descent). Let 2 72 2..2 0 satisfy Ex=1 mh = oo and Σx=1 Th < cc. Let L P(R)) + R be Wasserstein differentiable in the sense that (14) holds. Denoting by y(t the steps in Algorithm 1, and suppose that L(v(O)) is finite and S [VVVc(u(0)]1 doll) is bounded. Then
/ |1 = 1
lim
4.2 Hybrid algorithm
A main limitation of the BA algorithm is that the support of (() is restricted to that of the (possibly bad) initialization (0) On the other hand, Wasserstein gradient descent (Algorithm 1) only evolves the particle locations of (() but not the weights, which are fixed to be uniform by default. We therefore consider 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 conver- gence 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 X y Rdd and p(x, v) | y1 Leveraging work the statistical complexity of EOT [Mena and Niles-Weed 2019], we obtain finite-sample bounds for the theoretical estimators implemented by WGD 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 recent result by Rigollet and Stromme [2022] might be useful for deriving alternative bounds under distortion functions other than the quadratic.
7

---

Proposition 4.3. Let be o'-suaaausiann Then every optimizer v² of (4) and (8) is also o².
subgaussian. Consider C := LEOT. For a constant Ca only depending on d, we have
0 [5d/21+6
min VEP(RR) L(,, min (R.)
((),un)
S Ca
/54//4143
Vn
E
min MEP(RR)
((2,V)- min
vEP(RA) C(u²", v2)
S c
[5d/21+4
1
c/dd/14+3
mm'
o [5d/2]+6
1
E min (EP(RR)
L(U,
min PA (RE)
C(um, ,(n
S Cae
vn Vn)
el54/41+3
v/m
for all n, m E N. where Pn (BR is the set of probability measures over IR supported on at most n points, nT is the empirical measure of μ with m independent samples and the expectation B[1] is
over these samples. The same inequalities hold L == X LBAA with the identification e = =
4.4 Estimation of rate and distortion
Here, we describe our estimator for an upper bound (D, 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 and K, we always have that D = f pd(p K) and R. = Hu @9 K (0) lie on an upper bound of R(D) [Berger, 1971]. The two quantities can be estimated by standard Monte sample from μ @ K and evaluate the density duž dugg (x. v) dk(z.) (y).
Carlo provided
When only v is given, e.g., obtained from optimizing (6) with WGD NERD, we estimate an R-D upper bound as follows. As in the BA algorithm, we construct a kernel K2 similarly to (11), i.e., (x,) dv
a > Je (atyy then we estimate (D, R) using the pair (v, K2) as described earlier.
As NERD uses continuous v, we follow [Lei al., 2023a] and approximate it with its n-sample empirical measure to estimate (D, R). A limitation of NERD, BA, and our method is that they tend converge rate estimate of most log(n), where is the support size of v. This is because as the algorithms approach an n-poinn minimizer V of the R-D problem, the rate estimate R approaches the mutual information of μ Keg which is upper-bounded by log(n) [Eckstein and Nutz, 2022]. In practice, this means if target point of R(D) has rate then we need 2 e" 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 X and n points in y, which usually has cost of (mndd) for 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 sigmificant 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 high-rate setting, a large number of 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 use small batch size m with stochastic optimization).
Experiments
We compare the empirical performance of our proposed method (WGD) and its hybrid variant with Blahut-Arimon (BA) [Blahut, 1972, Arimoto, 1972], RD- VAE [Yang and Mandt, 2022], and NERD [Lei et al., 2022] on the tasks of bounds. While we experimented with WGD for both LBA and LEOT. we found the former to be 10 to 100 times faster computationally while giving similar or better results; we therefore focus on WGD for LB in discussions below. For the neural-network baselines, we use the same (or as similar as possible) network architectures as 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
deconvolution and estimation of R-D upper
8

---

3.0
Banut-Ayimoo
WGD (proposed)
2.9
Hobrid sigoritim (proposed) (2022)
2.8
Lei (2025)
OPT
2.7
WGD (proposed)
Hyled (poopsedd)
2.6
2.5
10
10a
10
10
Iteration
Figure 3: Visualizing i samples (top left), as well the returned by various algorithms compared to the ground truth v* (cyan).
Figure Losses over iterations. Shading corresponds to one standard deviation over random initializations.
2
10
300 (d
Optinality gap Clr)
Blahut-Arimoto
10P
10
Lei et al. (2023)
10-1
10P
WGD (proposed)
10
Hybrid WGD (proposed)
10-1
10-2
10
10-1
10
10'
10P
10a 101
10P
10
10
10P
10P
101
10P
107
7 (number of y particles)
Figure 4: Optimality gap v.s. the number of particles n used, on deconvolution problems with different dimension d and distortion multiplier X. The first three panels fix À = o and increase d, and the right-most panel corresponds to the 2-D problem with a higher (denoising with narrower Gaussian kernel). Overall the WGD methods attain higher accuracy for given budget of n.
with 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, apply them to deconvolution problem with known ground truth (see Sec. 2.3). We adopt the Gaussian noise as before, letting a be the uniform measure the unit circle in IR2 and the source μ - a (0 a²) with o² 0.1.
use n = particles for BA, NERD, WGD and 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- DAE 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 L(a). In Fig. 3, we visualize the final (t) at the end of training, compared to the ground truth // supported on the circle (colored in cyan). Note that we initialize (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 different solution where 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 place the support of their exactly the circle.
We additionally compare how the performance of BA, NERD, and the proposed algorithms scale to higher dimensions and higher À (corresponding to lower entropic regularization in Lgor). Fig. 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.

---

phyais
RD-WME, Mardh 11 2003 WGD, proposed
R.D bounds on the MNIST training set
Ru(77), Ying & 2003
10
Ru(2%), (a 40030)
Rate (nats per semple]
Rate (rats per sample)
Re(2)), WGD (e-1000)
Rer(2), WED (e-10000)
Re(P)) WGD (n=40000)
0.000
0O0GG 0.c0n 30001 0.00an 0,0355
0
NGLAA) Yong & Mandt 2002
Iteration
speech
20
Ballé 2021
HD-AAL Yong Marde al.
Ying Manct Ru
WGD, proposed
Lei Ry(
Los
Proposed WGD (2
Yong & Mavit Rg)
0.00
0.31
0.02
0.08
0.01
0.05
0.00
0.07
0
Distartion (mean squared emrr)
Distortion [mmea squared error)
U0*
Wallimm (seconds)
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 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-dimensionall data
We perform R(D) estimation higher nigher-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 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.
Discussions
In this work, leverage tools from optimal transport to develop new approach for estimating the rate-distortion function in the continuous setting. Compared to state-of-the-ar neural approaches [Yang and Mandt, 2022, Lei et al., 2022], our Wasserstein gradient descent algorithm offers comple- mentary strengths: 1) It requires a single main hyperparameter n (the number of particles) and no network architecture tuning; and 2) empirically we found 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. and 3). However, like NERD [Lei et al., 2022], our method has a fundamental limitation - it requires an 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-VAI YYngg and Mandt, 2022] may still be preferable on high-rate sources, while our particle-based method stands as a state-of-thee -otthe---tt--tt-ee 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-basee algorithm also generalizes optimal quantization (corresponding to E =
LOOT and projection of the source under the Wasserstein distance [Graf and Luschgy, 2007, Gray, 2013]) to incorporate a rate constraint
(e > 0), and would be interesting to explore the use of the resulting rate-distortion optimal quantizer for practical data compression and communication.
10

---

References
CE Shannon. Coding theorems for a discrete source with a fidelity criterion. IRE Nat. Conv. Rec., March 1959, 4:142-163, 163, 1959.
R. Blahut. Computation of channel capacity and rate-distortion functions. IEEE Transactions Information Theory, 18(4):460-473, 1972. doi:
Suguru Arimoto. An algorithm for computing the capacity of arbitrary discrete memoryless channels.
IEEE Transactions on Information Theory, 18(1):14-20, 1972.
Matthew T. Harrison and loannis Kontoyiannis. Estimation of the rate rate- distortion function. JEEE Transactions Information Theory, 54(8):3757-3762, 2008. doi: 10.1 10.1109tti..0008.222 92668887
Jerry Gibson. Rate distortion functions and rate distortion function lower bounds for real-world sources. Entropy, 19(11)6004, 2017.
Yibo Yang and Stephan Mandt. Towards empirical sandwich bounds on the rate-distortion function. In International Conference on Learning Representations, 2022.
Eric Lei, Hamed Hassani, and Shirin Sacedi Bidokhti. Neural estimation of the rate-distortion function with applications to operational source coding. IEEE Journal on Selected Areas in Information Theory, 2023a.
Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré. Gradient flows in metric spaces and in the space of probability measures. Lectures in Mathematics ETH Zürich. Birkhäuser Verlag, Basel, second edition, 2008.
Gonzalo Mena and Jonathan Niles-Weed. Statistical bounds for entropic optimal transport: sample complexity and the central limit theorem. Advances in Neural Information Processing Systems, 32, 2019.
Yury Polyanskiy and Yihong Wu. Information theory: From coding to learning. Book draft, 2022.
Robert M Gray. Entropy and information theory. Springer Science & Business Media, 2011.
Amir Dembo and I. Kontoyiannis. Source coding, large deviations, and approximate pattern matching. IEEE Transactions on Information Theory, 48(6):1590 1615, 2002.
Imre Csiszár. On extremum problem of information theory. Studia Scientiarum Mathematicarum Hungarica, 9, 01 1974a
Gabriel Peyré and Marco Cuturi. Computational optimal transport: With applications to data science. Foundations and Trends Machine Learning. 11(5-6)3355-600 607, 2019.
Federico Bassetti, Antonella Bodini, and Eugenio Regazzini. minimum kantorovich distance 2006.
estimators. Statistics & probability letters,
Martial Agueh and Guillaume Carlier. Barycenters in the wasserstein space. SIAM Journal Mathematical Analysis, 43(2):904-922 2011.
Alexander Alemi, Ben Poole, Ian Fischer, Joshua Dillon, Rif A Saurous, and Kevin Murphy. Fixing broken ELBO. In International Conference on Machine Learning, pages 159-168. PMLR, 2018.
Johannes Ballé, Valero Laparra, and Eero P Simoncelli. End-to-end optimized image compression. International Conference Learning Representutions, 2017.
Lucas Theis, Wenzhe Shi, Andrew Cunningham, and Ferenc Huszár. Lossy image compression with compressive autoencoders. International Conference on Learning Representations, 2017.
Yibo Yang, Robert Bamler, and Stephan Mandt. Improving inference for neural image compression. In Neural Information Processing Systems (NeurIPS), 2020, 2020.
Philippe Rigollet and Jonathan Weed. Entropic optimal transport is maximum-likellhowe deconvolu- tion. Comptes Rendus Mathematique, 356(1 1-2-1:22-2-131 2018.

---

Raymond J Carroll and Peter Hall. Optimal rates of convergence for deconvolving density. Journal of the American Statistical Association, 83(404):1184-1155 1988.
Bruce G Lindsay and Kathryn Roeder. Uniqueness of estimation and identifiability in mixture models. Canadian Journal of Statistics, 21(2):139-447 1993.
Imre Csiszár. On the computation of rate-distortion functions (corresp.). IEEE Transactions on Information Theory, 20(1):122-124, 1974b. doi: 10.1109/TT..174.105555
Imre Csiszár. Information geometry and alternating minimization procedures. Statistics and Decisions, Dedewicz, 1:205-233 1984.
Robert M. Gray and David L. Neuhoff. Quantization. IEEE transactions on information theory, 44 (6):2325-2383 1998.
George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan. Normalizing flows for probabilistic modeling and inference. The Journal of Machine Learning Research, 22(1):2617,2680, 2021.
Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv:1312.61 14, 2013.
arXiv preprint
Shitong Wu, Wenhao Ye, Hao Wu, Huihui Wu, Wenyi Zhang, and Bo Bai. A communication optimal transport approach to the computation of rate distortion functions. arXiv preprint arXiv:2212.10098, 2022.
Eric Lei, Hamed Hassani, and Shirin Saeedi Bidokhti. On a relation between the rate-distortion function and optimal transport, 2023b.
Aude Genevay, Lénaic Chizat, Francis Bach, Marco Cuturi, and Gabriel Peyré. Sample complexity of Sinkhorn divergences. In The 22nd International Conference on Artificial Intelligence and Statistics, pages 1574-1583, PMLR, 2019.
Philippe Rigollet and Austin Stromme. On the sample complexity of entropic optimal transport.
arXiv preprint arXiv:2206.13472 2022.
Yuling Yan, Kaizheng Wang, and Philippe Rigollet. Learning gaussian mixtures using the wasserstein- fisher-rao gradient flow. arXiv preprint arXiv:2301.01766 2023.
Lenaic Chizat, Gabriel Peyré, Bernhard Schmitzer, and François- Xavier Vialard. An interpolat- ing distance between optimal transport and fisher-rao metrics. Foundations of Computational Mathematics, 18:1-44, 2018.
Lénaic Chizat. Mean-field langevin dynamics: Exponential convergence and annealing. arXiv preprint arX1v:2202.01009, 2022.
Marcel Nutz. Introduction to entropic optimal transport. Lecture notes, Columbia University, 2021. attpss//wwxx math. cowumbia. wwt.:uaa.....emmmtt///oooaate/aa ctth//mutt///ccs//oLLLcccceee
pdf.
Guillaume Carlier, Lénaïc Chizat, and Maxime Laborde. Lipschitz continuity of the Schrödinger map in entropic optimal transport. arXiv preprint arXiv:2210.0022. 2022.
Toby Berger. Rate distortion theory, a mathematical basis for data compression. Prentice Hall, 1971.
Stephan Eckstein and Mareel Nutz. Convergence rates for regularized optimal transport via quantiza- tion. arXiv preprint arXiv:2008. /4391, 2022.
Eric Lei, Hamed Hassani, and Shirin Sacedi Bidokhti. Neural estimation of the rate-distortion function for massive datasets. In 2022 IEEE International Symposium on Information Theory (ISIT), pages 108-613. IEEE, 2022.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86( 86(11))2278-222) 1998.
12

---

Andre Wibisono. Sampling as optimization in the space of measures: The langevin dynamics as composite optimization problem. In Sébastien Bubeck, Vianney Perchet, and Philippe Rigollet. editors, Proceedings of the 31st Conference On Learning Theory, volume 75 of Proceedings of Muchine Learning Research, pages 2093-3027. PMLR, 06-09 Jul 2018. URL https:// proceedings. .ml..mmr....... wttpss//////wiiiioss/m
html.
Qiang Liu and Dilin Wang. Stein variational gradient descent: general purpose bayesian inference algorithm. Advances in neural information processing systems, 29, 2016.
Abhishek Shetty, Raaz Dwivedi, and Lester Mackey. Distribution compression in near-linear time. arXiv preprint arXiv:2111.07941 2021.
Siegfried Graf and Harald Luschgy.
Foundations of quantization for probability distributions.
Springer, 2007.
Robert M. Gray. Transportation distance, shannon information, and source coding. GRETSI 2013 Symposium on Signal and Image Processing, 2013. URL htw.h:ttp////ett //ee. http:///eeestt.tfttteeeee wtt.nfrrrdd.oooo edu/"gray/ gretsi.pdr
Erhan Çinlar. Probability and stochastics, volume 261. Springer, 2011.
Gerald B Folland. Real analysis: modern techniques and their applications, volume 40. John Wiley & Sons, 1999.
Yury Polyanskiy and Yihong Wu. Lecture notes on information theory. Lecture Notes for ECE563 (UIUC) and, 6(2012-2016);7, 2014.
Yibo Yang, Stephan Mandt, and Lucas Theis. An introduction to neural data compression. Founda- tions and Trends@ in Computer Graphics and Vision, 15 2023, ISSN 1572-2740. doi: 10. 000156//0000000107. URL httt://t/// doi org/10...0..... 1ttp//10.116000000000000000
David M Blei, Alp Kucukelbir, and Jon D McAuliffe. Variational inference: A review for statisticians. Journal of the American statistical Association, 1125118):59--8772200 2017.
MJ Beal and Z Ghahramani. The variational bayesian algorithm for incomplete data: with application to scoring graphical model structures. Bayesian statistics, 7(453-464):210, 2003.
Arthur P Dempster, Nan M Laird, and Donald B Rubin. Maximum likelihood from incomplete data via the em algorithm. Journal of the royal statistical society: series B (methodological), 39(1): 1-22, 1977.
Mingtian Zhang, Peter Hayes, Thomas Bird, Raza Habib, and David Barber. Spread divergence. In International Conference Machine Learning, pages 11106-11116. PMLR, 2020.
Pascal Vincent. A connection between score matching and denoising autoencoders. Neural computa- tion, 23(7):1661-167, 2011.
Michael Jordan, Zoubin Ghahramani, Tommi Jaakkola, and Lawrence Saul. An introduction to variational methods for graphical models. Machine learning, 37:183-233 21-8833-3,-11 1999.
Martin Wainwright, Michael I Jordan, al. Graphical models, exponential families, and variational inference. Foundations and Trends@ in Machine Learning. 1(1 (1-2):1-305, 2008.
Michael Irwin Jordan. Learning in graphical models. MIT press, 1999.
Diederik P Kingma and Jimmy Lei Ba. Adam: A method for stochastic gradient descent. In International Conference Learning Representations, 2015.
John Platt and Alan Barr. Constrained differential optimization. In Neural Information Processing Systems, 1987.
George Papamakarios, Theo Pavlakou, and lain Murray. Masked autoregressive flow for density estimation. In Advances in Neural Information Processing Systems, pages 2338-2347, 2017.
13

---

Roger 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-Milee and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL
http:// http//gtttttbb com/ google/jax.
Broader Impacts
Improved estimates on the fundamental cost data compression aids the development and analysis of compression algorithms. This helps researchers and engineers make better decisions about where 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 only accurate insofar as the training data representative of the population distribution of interest, and practitioners need to ensure this in the data collection process.
Acknowledgements
We thank anonymous reviewers for feedback on the manuscript. Yibo Yang acknowledges sup- port from the Hasso Plattner Foundation. Marcel Nutz acknowledges support from NSF Grants DMS-18 1812661, DMS- 106056. 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, well gifts from Intel, Disney, and Qualcomm.
14

---

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 the connections between the R-D estimation problem and variational inference/leaminng Section 10, provide additional experimental results and details in Section 11, and list an example implementation of WGD Section 12. Our code and can be found at https: : / het//ithuu com/ htt/yyyybyynnn/e /vgd.
7 Notions from probability theory
In this section we collect notions of probability theory used in the main text. See, e.g. IÇinlar, 2011] or [Folland, 1999] for more background.
Marginal and conditional distributions. The source and reproduction spaces X, are equipped with sigma-algebras Ax and Ay, respectively. Let N' x y denote the product space equipped with the product sigma algebra Ax @ Ay. For any probability measure т on X' V, its first marginal
T (A) = xAA)),, A E Ax,
which is probability measure on X. When æ is the distribution of random vector (X, Y), then T the distribution of X. The second marginal of i is defined analogously as
Tz == π (X x B), BE Ay-
Markov kernel or conditional distribution K(x, dy) a map x Ay - [0, 1] such that
((x,.) a probability measure on y for each X::
2. the function x + (xx, B) is measurable for each set B E Ay-
When speaking of the conditional distribution of a random variable Y given another random vari- able X, we occasionally also use the notation Qryx from information theory [Polyanskiy and Wu, 2014]. Then, Qvix=a(B) K(x, B) is the conditional probability of the event {Y E B} given X
Suppose that probability measure on is given, in addition to a kernel K(x, dy). Together they define unique measure j K the product space &' For rectangle set B E Ax Ay.
K(AxB) S (ldd)) B). A BEAy
The measure := a & K has first marginal #1 = p.
The classic product measure is special case of this construction. Namely, when a measure v on y is given, using the constant kernel K(r, dy) v(dy) (which does not depend gives rise the product measure & V,
P 8 V A B) (BA)(BB), A AEAx,
Under mild conditions (for instance when X. are Polish spaces equipped with their Borel sigma algebras, as in the main text), any probability measure on x y is the above form. Namely, the disintegration theorem asserts that can be written as the joint distribution of random vector (X, Y), this says that there a measurable version of the conditional distribution Qyjx.
K for some kernel K. When is
15

---

Optimal transport. Given a measure on A' and a measurable function T : X - V, the pushforward (or image measure) of μ under T a measure on V given by
T B) = (( (B)), B Ay-
If is seen as random variable and lr as the baseline probability measure, then T is simply the distribution T.
Suppose that and v are probability measures on X' = - Rd with finite second moment. As introduced the main text, I(,,,)) denotes the set of couplings, i.e., measures on & with Ti = μl and TT V. The 2- Wasserstein distance W2(,, v) between μ and v is defined as
W2(u,2) = ( inf
[Ell((u.)) / fly =1I (d1)=(dzdd)))
1/2
This indeed defines a metric on the space probability measures with finite second moment.
We finish this section by giving proof of the basic equivalence between optimization of LEOT and CBA. which goes back to [Csiszár, 1974a, Lemma 1.3 and subsequent discussion]. Recall that 11( denotes the set of measures on the product space X' y with μ.
Lemma 7.1. Set = A-1 It holds that inf LEOT(V)= inf X LBA(D) and PEP())
arg mrg LEOT(V) = arg min LBA(V) vEP())
vEP())
vEP())
Proof. Both statements will follow from simple property of relative entropy [Polyanskiy and Wu, 2022, Theorem 4.1, "golden formula"]: For π E I(,,.) with H( v) < oc, the properties of the logarithm reveal
( n 00 v) H( a 60 R2) + H( v);
(15)
and hence (TT ( 2 H μ 1 (&TT)) π2). This implies
I pdTT + ell(T T2) *ET((r))
(EOT(V) inf vEP(Y)
inf
/ pdr + cH (H(T)
inf inf nell(u.) vEP(o))
inf (BA(U).
X VEPO))
Further, any optimizer v of LBA with corresponding optimal "coupling" πt satisfy H( ['' 0 (otherwise, taking 1 मे has better objective) and thus тミ = V" and n' E vI), therefore n" is also optimizer of LEoT by the above equality. Conversely, any optimizer n² of LOTT with coupling x' I[( v²) C II(,,,)) clearly yields a feasible solution for LBA as well, and hence is also optimizer of LA by the above equality.
I(p,)) must
8 Numerical estimation of rate and distortion from reproduction distribution
Given reproduction distribution and kernel X' Ay [0, 1]. the tuple (D, IR² defined by
/ pdllu @ K)
and
R: H(p & Klu 8)
lies above the R(D) curve. This again follows from the variational inequality (15) 2 H(T (μ x K)2) I(u x K), that R 2 R(D) ( (TT).
Given only a reproduction distribution v, we will construct a kernel K from v and use (μ, K) to
compute (D, R), letting
dK(x,, (v) J --- Te-xpta.ovv(dj)
>xpxxy
16

---

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 LA (5) for the given v. Plugging into the formulas for and R gives
p(ar, ()) /
D
v(dx, dy)
(16)
and
R / log
K(x,dy) v(du)
K (x,, dy)
(17)
= [x(x)) - Ap(a, (aplaa))poee))) @ v(dx, dy),
where we introduced the shorthand
(xxx)= = - log
A(dz))))(dy)-
(18)
Note that we have the following relation (which explains (6))
(BA(V)= XD | ((z))((x))
Let y be an nepoint measure, / Ei= Wgyy,, e.g., the output of WGD or in the inner step of NERD. Then (o(x:) can be evaluated exactly as finite sum over the particles, and the expressions above for 'D and R (which are integrals w.r.t. the product distribution @ v) can be estimated as sample averages. That is, given m independent samples {=,] from μs we compute unbiased estimates
w, wyp(x, yj
w [))p(( Ap(z;, (())))( -Aplos))))))
1=
where (pxx) log "i=1 e0(a)) (A((z))A Similarly, sample mean estimate for LBA is given by
(nA(P) Σ
p(x;) = log
(19)
==1
In practice, found simpler and numerically more stable to instead compute R as
R (BA(V)) XD.
Whenever possible, avoid exponentiation and instead use logsumexp to prevent numerical issues.
9 Wasserstein gradient descent
9.1 On Wasserstein gradients of the EOT and rate f'unctionals
First, elaborate the Wasserstein gradient of the EOT functional LEOT(V). That the dual potential from Sinkhorn's 's algorithm 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 in unbounded settings. While Mena and Niles-Weed [2019] only states the result for quadratic cost, the approach of Proposition therein applies more generally.
Below, compute the Wasserstein gradient of the rate functional (BA(P). Recall from (6),
LBA(V) ff log /
(al(x,8)))((e (op(-Ap(xx)))u(bbu((a)))
17

---

Under sufficient integrability on μ and // to exchange the order of limit and integral, we can calculate the first variation as
C((1 E) + =i) L(V)
1
Teep( (xp(-AA(x, u)) (()))(dd
lim e-+
lim c-2
log
p(dx)
I /exp(- (xxp(-Ap(z.)))u(d0
E
1
(exp(->o(sss))e0
lim
log
1 | (xp(-AA(z,
((drr)
c0 E
exp( (o0(--ppp 1--pp(,)))) v))
exp( - Ap(z, ()) (dy) ((dx) (D (D-v)dyy) vdy),
where the last equality uses lime- & log(1 + ex) x and Fubini's theorem. Thus the first variation v of LAA at v is
cxp(-Ap -p(xp.))))) j((x))
v
(20)
exp(-Appxx 1) (dg)
To find the desired Wasserstein gradient of LBA, it remains to take the Euclidean gradient of v)", ie..
V (nv(v) = Vva Doing so gives us the desired Wasserstein gradient:
VVCBA (v)
V V, (v)
0
exp( - App(z, v))
J exp( Ap(ar, (d)j)ddj) ((dr)
oy
exp(-Ap(az; y)) by ((z,)) y)
u(dx),
(21)
J
()) (djj
again assuming suitable regularity conditions on p and ب 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 2 72 2 ... 2 0 and 0, 2 20, C.IEN,O C> 0 satisfy Σe= = oc, Et 77 < 00, 1= ac7 < do and Ja- - ac+1l < for all e N. Then lim al = 0.
Proof. The conclusion remains unchanged when rescaling a by the constant C. and thus without loss of generality C =
Clearly 7t- + 0 as E x2 < co. Moreover, there exists a subsequence of (at hee which converges to zero (otherwise there exists > 0 such that a; 2 8 > 0 for all but finitely many f, contradicting E 7xde oo).
Arguing by contradiction, suppose that the conclusion fails, i.e., that there exists subsequence of (ar heN which is uniformly bounded away from zero, say 2 8 >> O along that subsequence. Using this subsequence and the convergent subsequence mentioned above, we can construct subsequence
a Oiyy Oig.... where as &e 0 for n odd and a, 2 for even. We will show that
ian
am? 2 82/2 for all N,
contradicting the finiteness of E, Mag (The notation (2) indicates (in)equality up to additive terms converging to zero for + oo.)
To ease notation, fix n and set m 12--1 and M = i2n- We show that ΣM ar% 2 8//2. To this
end, using |a au41l I Y we find
M
M-1
ag 2 RM
Σ
aa > 8
j=k
18

---

Since am 2 0, there exists a largest no N, no 2 m, such that Σ =j=no 76 2 8 (and thus
Ey=ng 7/j S "mo Ro ò as well), We conclude
M
M
M
Σ 1m
7:
Σ
7
Σ I=no
7t
Σ
Tj
2 8
Σ
7171022
1=Ro
j=k
tm j=no
1
82
Σ
8
1t
2
where we used that Σi= 72 Rd This completes the proof.
Proof of Proposition 4.2. Using the linear approximation property in (14), we calculate
[(v()) [(v(0)) Σ C(u(tt))-fae
' / /VVcl(00) dull +20o (
As C(vlo)) is finite and ((v(n) is bounded from below, it follows that
7 / [VVVc((())
dx(0)) 00.
The claim now follow by applying Lemma 9.1 with a; ( 1|2 dv(t); note that the assumption in the lemma, ac+1l S Crte, is satisfied due to the second inequality in (14) and a short calculation below
Gullo 1|2 1 1
(Veuuretty 1,2
S CW2(u(l),
sc
(ff 100 - - voll---nweccooaaaa))a0()
Cn
s Cr sup
v
s C'n
where we use sup, / [1 oe in the last step.
9.3 Proof of Proposition 4.3 (sample complexity)
Recall that X' = = IRd and p(a, y) = || y2 in this proposition. For the proof, we will need the following lemma which is of independent interest. We write v Se μ if v dominated by μ convex order, i.e., f du S f du for all convex functions f Rd + R.
Lemma 9.2. Let have finite second moment. Given v E P(RR), there exists D E P(Rd) with Se p and
LEoT D) < (gOT(U, v).
This inequality is strict if v.Z. In particular, any optimizer v" of (8) satisfies v² Se p.
Proof. Because this proof uses disintegration over D. is convenient to reverse the order of the spaces in the notation and write generic point as (x. ) x X. Consider II(r,, ()) and its disintegration æ = v(dx) @ K(x, dy) over x E V. Define T Rd - Rdd by
T(x) = / K dy).
19

---

Define also ãt := (T.id) T and ;= Æ1. From the definition of T, we see that म is martingale, thus Se μ. Moreover, μ (T.id) 3 & μs. The datc-psogssiing inequality now shows that
1 (
On the other hand, f y (x, dÿ) - y/1 K(x, dy) S J || k dy) since the barycenter minimizes the squared distance, and this inequality is strict whenever # 5K (z. dÿ). Thus
/ | y112 ( dy) - ||= vll2
x xdx dy),
and the inequality is strict unless T(x) = for v-a.e. which in turn is equivalent to ı being martingale, The claims follow.
Proof of Proposition 4.3. Subgaussianity of the optimizer follows directly from Lemma 9.2.
Recalling that inf, LEOT(U) and inf, X- have the same values and minimizers (given by (9) in Sec. 2.2), it suffices to show the claim for L = LEOT. Let v² be optimizer of (8) (i.e., an optimal reproduction distribution) and v" its empirical measure from n samples, then clearly
min LEOT(URM) (2.))
min (EOT(H,U) MEP(R")
(B()) min LEOT(U, VEP(2()
min LEOT(UN))
SE | [EooT(,, - LEOT (1,, w*)l]
where the expectation is taken over samples for 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 Vr, the optimizer for the problem (8) with replaced by pT Similarly to the above, we obtain
E
min (R) min LEOT(U", MEP(Ba)
v)
(EP(R))
(LEOT(,,U) - CEor(u" mill
E
max VE(V" wi.
where the expectation is taken over samples from pmth 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 D E {v" "/"h is through their subgaussianity constants. By Lemma 9.2, these constants are bounded by the corresponding constants of μ and y²" 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 n²") and the triangle inequality, where again use the arguments in the proof of [Mena and Niles-Weed, of n
Weed, 2019, Theorem 2] to bound the expectation over the subgaussianity constants
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.
There are two key properties use for the convergence of the base WGD algorithm: I) certain monotonicity of the update steps (up to higher order terms, gradient descent improves the objective) and 2) stability of gradients across iterations. If 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, Ee Nt. I [ (
(here (()) are the outputs from the respective BA
20

---

Table I: Guide to notation and their interpretations in various problem domains. "LVM" stands for latent variable modeling. "NPMLE" stands for non-parametric maximum- likeliooo estimation. The R-D problem (3) is equivalent to projection problem under an entropic optimal transport cost (discussed in Sec. 2.2) and statistical problems involving maximum-likellhood estimation (see
discussion in Sec. 2.3 and below).
Context OT R-D LVM/NPMILE deconvolution
Px
p(x, y)
K = QYIX
V Qy
source distribution data distribution data distribution noisy measurements
transport cost distortion criterion
"transport plan"
target distribution
compression algorithm variational posterior
codebook distribution prior distribution noiseless model
"- log p(æll)) "noise kernel"
steps and "t is the step size of the gradient steps). Without property 2), cannot conclude I [IVVE(v(O)II du(t) + 0 for t + co. 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 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 Basserstein-Fisher Rao geometry (cf. [Yan et al., 2023]), in which the BA step corresponds to gradient update in the Fisher-Rao geometry with 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, 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 special case. Also see Section 3 of [Yang al., 2023] for 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 statistics and information theory. Table summarizes the notation and the correspondence to the measure-theoretic notation used in the main text.
In statistics, common goal is model (unknown) data generating distribution Px by some density model p(z). In particular, here we will choose p(z) to be a latent variable model, where ) takes on the role of latent space, and Qy V is the distribution of 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 given model (see Sec. 10.2), and the parametric R-D estimation problem, i.e..
inf LBA(V),
DCH
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 Ç on V, and suppose Qy has density q(y) w.r.t. Ç. Often the latent space y is Euclidean space, and q(v) is the usual probability density function w.r.t the Lebesgue measure Ç: or when the latent space is discrete/countable, Ç is the counting measure and ((y) is a probability mass function. We will consider the typical parametric estimation problem and choose a particular parametric form for Qy indexed by a parameter vector 0. This defines a parametric family H {Q& :0 E Q) for some parameter space O. Finally, suppose the distortion
21

---

function induces conditional likelihood density via p(xly) & e e-An(as)), with normalization constant that constant in y.
These choices then result in a latent variable model specified by the joint density ((v)p(r))) and we model the data distribution with the marginal density: 2
p(æ) = S, p(zzy)ddQ
-/. , p(aly)g())((d)))
(22)
As an example, a Gaussian mixture model with isotropic component variances can be specified as follows. Let Qy be a mixing distribution on X' y I parameterized by component weights W and locations pi....,k, such that Qy = > k=1 WKKmm. Let p(arlu) = NN(,, o²) ) be conditional Gaussian density with mean y and variance o². Now formula (22) gives the usual Gaussian mixture density on R..
Maximum-likelihodd estimation then ideally maximizes the population log (marginal) likelihood,
f log Px(x)) = / log (Ls plelw)dQx(w) Px (dx).
Ex ((x)]
(23)
The maximum-likellhonn 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 Px * (O, +) for some unknown a and known noise (0, }) (i.e., the model is well-speciffed), 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 al., 1999, Wainwright al., 2008).
10.2 Connection to variational inference
Given a latent variable model and any data observation æ, central task in Bayesian statistics is to infer the Bayesian posterior [Jordan, 1999], which we formally view as a conditional distribution Qypx=x' is given by
p(xly) p(x)
dQy(v)
or, in terms of the density 4(v) of Qr. given by the following conditional density via the familiar Bayes' rule,
q" (vl=)) p(xll))(())
p(alo)a(y) p(x) Jy P(
Unfortunately, the true Bayesian posterior is typically intractable, as the (marginal) data likelihood in the denominator involves often high-dimensional integral. Variational inference [Jordan et al., 1999, Wainwright et al., 2008] therefore aims to approximate the true posterior by variational distribution Qyjx== P())) by minimizing their relative divergence H(Qv (0l2vx===1111 1Qr The problem is equivalent maximizing the following lower bound the marginal log-likelihood, known as the Evidence Lower BOund (ELBO) [Blei et al., 2017]:
arg min (Qxi xx = arg max ELBOOQQ,F I, Qyix=e),
Qrix-z
Qrix=x
ELBO(QY.C Qvix=x)
Ey-Qvixx [ogg(<ly)) H(Qrix==(Qx)
-log((x) 11(0)ix=1Q11x=2)
(24)
Recall the rate functional (5) arises through an optimization problem over a transition kernel K,
/ pdll ss K) + H(u Klpgv)
(B)(V) inf X K
"To be more precise, we always fix a reference measure 7 on X, so that densities such as p(x) and p(æly) with respect In the applications considered in this work, is the Lebesgue measure X' = Rd
22

---

Translating the above into the present notation (μ - Px.K - Qrix. Qy; see Table 1), we obtain
LBA(Qv)= = inf Bx (ogp(())))) Rezax ((Qx)x=)000) + const
QYix
inf Ex-Px|- Bearc--EEO(Qyy20 Q Qy(X==)] + const. Qrix
(25)
which allows interpret the rate functional the result of performing variational inference through ELBO optimization. At optimality, Qrix = Qy the ELBO (24) is tight and recovers
log ((zr), and the rate functional takes on the form of (negated) population marginal log likelihood (23), as given earlier by (6) Sec. 2.1.
10.3 Connection to variational EM
The discussion so far concerns probabilistic inference, where a latent variable model (Qy,p(<))) has been given and we saw that computing rate functional amounts variational inference. Suppose now we wish to learn 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-likelhhodd the variational EM algorithm maximizes the population ELBO
Ex~Px ELBO(Qy, xy = Px Qu
Eg 1Q
(26)
w.r.t. Qy and Qrjx. This precisely corresponds the R-D problem inf Cna (Qv), using the form of CBA (Qr) from (25).
In popular implementations variational EM such as the VAE [Kingma and Welling, 2013], Qy and Qrix are restricted to parametric families. When they are allowed to range over all of P())) and all conditional distributions, variational EM then becomes equivalent the BA algorithm.
10.4 The Blahut-Arimote and (exact) EM algorithms
Here we focus on the distinction between the BA algorithm and the (exact) EM algorithm [Dempster 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 (Qrix Qy) with Qy E P()), whereas the EM algorithm
(Qyix, 0) with 0 indexing parametric family O). Thus the coordinate update w.r.t Qrix (the "E-step") is the same in both algorithms, but the subseuquent "M-step" potentially differs depending the role of O.
H. = 1Q6
Given the optimization objective, which is simply the negative of (26),
- + H(PxQrIxXPXP Qy),
(27)
both BA and EM optimize the transition kernel Qy the same way in the E-step, as
dQy
p(xly)
(v)
(28)
dQy
p(x)
For the M-step, the BA algorithm only minimizes the relative entropy term of the objective (27),
min @Y
XxxxQVxxxxx & Qy),
(with the optimal Qy given by the second marginal of Px 2x xx1x) whereas the EM algorithm mini- mizes the full objective w.r.t. the parameters 0 of Qy.
min E(nw) log p(xly) + Qi Px & Qy).
BEG
The difference comes from the fact that when parameterize Qy by 0 in the parameter estimation problem, Q and consequently both terms in the objective of (29) will have functional dependence on 0 through the E-step optimality condition (28).
23

---

-EE= wgdpg, and its parameters A consist of the components In the Gaussian mixture example, Qy = I==
weights (W1, ..... wx) E A and location vectors {1I,,.., jK) C IRd. The E-step computes 0Yix=== = Ex Wk pla pk) p(z) ju.- For the M-step, if we regard the locations as known that 0 (WI,..., wK) only consists of the weights, then the two algorithms perform the same update; however 0 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 P(z.))) logp(<]ly)
-
p(zluz) log p()pp))))))))))))) Px(x). x Wi p(a)
11 Further experimental results
Our deconvolution experiments were run on Intel(R) Xeon(R) CPUs, and the rest of the experiments were on Titan RTX GPUs.
In most experiments, we use the Adam [Kingma and Ba, 2015] optimizer for updating the v 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 values similarly to those in [Yang and Mandt, 2022], e.g., on a log-linear grid {1e4, 3e3, Ie3, 3e2, 1e2, 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 Ds. 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 distortion threshold of D; if the chosen n is not large enough, i.e., log(n) < R(D), then no supported on (at most) n points feasible for the given constraint (otherwise we have a contradiction). When this happens, the Lagrange multiplier associated with the infeasiable constraint (A in our case) will be observed to increase indefinitely rather than stabalize to a finite value under a method like MDMM.
11.1 Deconvolution
Architectures for the neural baselines For the RD-VAE, we used 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 softplus 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 softplus activation led to overly smooth mappings, which had difficulty matching the optimal va = measure which is concentrated on the unit eircle and does not have a Lebeggue density.
Experiment setup As discussed in Sec. 4.2, BA and our hybrid algorithms do not apply to the stochastic setting; be able to include them in our comparison, the input to all the algorithms an empirical measure n²'h (training distribution) with vn = 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-on samples.
Recall from Sec. 2.3, given data observations following μ tion with an assumed noise variance for some 1 S o², then the optimal solution to the problem is given by D² = a = (0,, }). We compute the optimal loss OPT = LB (v*) by a Monte-Carlo estimate, using formula (19) with m = 10 samples from μ and = 10t samples from //. The resulting LEΛ is positively biased (it overestimates the OPT in expectation) due to the bias of the plug-in estimator for (xx) (18), so we use large n to reduce this bias." Note the same issue occurs in the NERD algorithm (13).
(0, o²), if we perform deconvolu-
3 Another, more sophisticated solution would be annealed importance sampling [Grosse et al., 2015] or related method developed for estimating marginal likelihoods and partition functions.
24

---

À = 100.0
10.0
X = 3.0
À = 1.0
1
100.0
À = 10.0
= 3.0
A 1.0
Figure 6: Visualizing the converged n = 20 particles of WGD (top row) and hybrid WGD (bottom row) estimated m 10000 source samples in the deconvolution problem (Sec. 5.1), for decreasing distortion penalty A. The case À = 10.0 a-2 corresponds to "complete noise removal" of the source and recovers the ground truth (unit circle).
Optimized reproduction distributions To visualize the (continuous) learned by RD-VAE and NERD, we fit a kernel density estimator on 10000 v 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 decreasing distortion penalty À (equivalently, increasing entropic regularization strength c).
11.2 Higher-dimeenionall datasets
For NERD, set to the default 40000 in the code provided by [Lei et al., 2023a], on all three datasets.
For WGD, we used = 4000 on physics, 200000 on speech, and 40000 on MNIST (see also smaller n for comparison Fig. 5).
On speech, both NERD and WGD encountered the issue of log(n) upper bound on the rate estimate as described 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- 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 MNIST using similar neural network architecture and other hyperparameter choices as in the GAN-generated image experiments of [Yang and Mandt, 2022] (using ResNet-VAE for the upper bound and 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 self-contained minimal implementation of Wasserstein gradient descent on LBA, using the Jax library [Bradbury al., 2018]. To compute the Wasserstein gradient, evaluate the first variation of the rate functional in the particle locations æı.... using Jax's autodiff tool line 51.
according to (20), yielding >m 8 (Iz), then
simply take its gradient
The implementation of WGD on LEOT similar, except the first variation is computed using Sinkhorm's 's algorithm. Both versions can be casily just-in-time compiled and enjoy GPU acceleration.
25

---

Nasserstein GD on the rate functional L_ (BA).
2 import jax. numpy as jup import jax
from jax.scipy. specpyl import logsumexp
6 Define the distortion function 1rho.
squared_dift lambda x, y: jup. sum( y) ** 2) sairwise_distortion.
jax. vmap (aax. vmap diff, (None, 0)), (0, None)) )
10
def wgrad (mu_x, mu_,, nu_x, nu_,, rd_ rd_lambda):
12
Compute the Wasserstein gradient of the rate functional, which wi 12
14
to move the Inu particles.
:param mu_m: locations of 1mu atoms. :param mu_w: weights of Imu atoms. :param nu_x: locations of Inu atoms. :param nu_w: weights of Inu atoms.
18
:param rd_ lambda: R-D tradeoff hyperparameter.
:rtuur::
def
snu_xxx))
Here we compute surrogate loss based on the first variation \pst, which allows jaz autodiff to compute the destred Wasserstein gradient. :param nu_x:
27
:return: pst_ sum lsum_d posinn==(i]) mu.
29
C pairwise_distortionn www.tiiiiiinniiine fn (mu_x, x, nu_x)
31
scaled_C rd_lambda C # [m, n]
lug_nu_w jup log (nu_w) # [1, n]
33
Solve BA inner problem with a fixed phi Logsumexp(-scaled.0 log_mu_w, axis=1, keepdims-True) [m, 1] loss jnp. sum (mu_w phi) Evaluate the rate functional. Eq (6) paper.
35
Let's also report rate and distortion estimates (discussed in Sec. 4.4 of the paper). # Find 1pš-* via Iphi
pi jup. exp Scaled_C) inp.outer (mu_w, nu_w) [m, n]
distortion jnp. sum C) rate
rd_lambda distortion
Now evaluate 1pst on the atoms of Inu. phi = jax. lax atop-gradient (phi) psi jup.sum
46
sum wwpp. lax. lax. stop_stopsrree ceropp_grdientteeee (phi) scaled_C) mu_w, axis=0)
psi_sum sum jnp (psi) # For computing gradient w.r.t. t. each nu atom. return psi_sum, (loss, rate, distortion)
48
50
# Evaluate the Wasserstein gradient, e.e., Inabla lpsé, on nu_a. psi_prime, loss jax. grad (compute_psi__eeee las_aux-True) (nu_x) return psi_prime, loss
def (X, n, rd_lambda, num_steps, 1r, rng): вии
A basic demo of Wasserstein gradient descent a discrete distribution. :param X: a 2D array [N, d] of data points defüning the source lmu. :param n: the number of particles to use for Inu.
58
26

---

:param rd_ lambda: R-D tradeoff hyperparameter.
:param num_steps: total number of gradient updates. :param Lr: step size.
:param rng: jaz random key.
:return: (nu_#, 2, nu_w), the locations and weights of the final Inu. m.n
65
# Set up the source measure Ima.
m jnp.jizze size (X, 0)
mu_x = mu_w 1
69
jnp.onest ones (m, 1))
# Initialize Inu atoms using random training camples.
rand_id = jax random pww.oomttttmmttiooocooom (rng, m) [:n]
nu_x X[rand_idx
Locations of Inu atoms.
nu_ = 1 n onpponne((1,, n)) Uniform weights.
for step in range(num_stee) psi_prime, (loss, rate, distortion) = wgrad (mu_x, mu_w, nu_x, nu_w, rd_lambda) nu_x print (tepps loss= :.o::l0ss:........ 4gh, rate=(rate: : (estoriion(distorttio.....
psi_prime
return nu_x, nu_
if m__mamee..
main_
# Run a toy ezample le 2D Gaussian samples. rng ja..aaaannaa. waando..PNCCee PRNCKey (0)
X = jax. random. normal (rng, [10, 2]) nu_x, nu_w wgd(X, (X, n=4, rd rd__ambda=2., nun_steps=100, 1r=0.1, .1, rng=rng)
86
27


Reviews:
[{"review_id": "mWkQeUVMAH", "reviewer": "Reviewer_jdNJ", "summary": "This paper proposes an estimator using WGD and moving particles. This is different from prior methods that leverage neural networks to fit the unknown high-dimensional support of the optimal reconstruction distribution. The authors also note a connection with entropic OT and provide sample complexity bounds on estimating R(D). Numerical results show that it provides tighter estimates of R(D) compared to prior estimators on high-dimensional data.", "questions": "- Do the neural methods (NERD and Yang and Mandt) provide tighter bounds as the network complexity or architecture increases? If complex enough, could they be comparable with the particle method as well?", "limitations": "Yes", "rating": 5, "confidence": 5, "soundness": 4, "presentation": 3, "contribution": 2, "strengths": "- The method proposed is very principled, and well presented. The ideas are clear and the use of particles is natural for this problem. Showing how it can be implemented in practice was quite interesting.\n- The finite sample complexity result is nice, as it is the first such result in this line of work. Additionally, it shows one interesting use of the connection between R(D) and entropic OT. It would also be good to cite other recent work that discuss this connection (https://arxiv.org/pdf/2212.10098.pdf, https://arxiv.org/pdf/2307.00246.pdf).\n- The experimental results show good support of the method, with regards to the tightness of the bounds.", "weaknesses": "- While interesting, the work appears to be fairly incremental in terms of rate-distortion estimation. It provides slightly tighter bounds but fails to estimate at rates larger than log(n) which is a fundamental problem in rate-distortion (or mutual information) estimation (see https://arxiv.org/abs/1811.04251). Hence, it is difficult to say that this paper provides a wholly new solution to the rate-distortion estimation problem, as it seems to provide a third alternative to the same rate-distortion objective. In practice the prior work by Yang and Mandt is the only one so far applicable to most real-world data like images, where the rate per sample is far larger than log(n) for any feasible n. \n- The computational (specifically, memory) cost of the particle method was not discussed, and this should be mentioned, especially if each particle consists of high-dimension vectors. The benefit of the neural methods is that one does not need to maintain many high dimensional vectors. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your insightful feedback and bringing our attention to recent related work. We have followed your suggestion to cite (Wu et al. 2022, Lei et al. 2023), and added a new experiment on MNIST as well as more discussion around the computational cost of our method; also see points 2, 3, 5 of our main rebuttal.  Below we will go over all concerns raised and detail how they have been addressed.\n\n> While interesting, the work appears to be fairly incremental in terms of rate-distortion estimation. It provides slightly tighter bounds but fails to estimate at rates larger than log(n) which is a fundamental problem in rate-distortion (or mutual information) estimation…\n\nWe agree that the log(n) issue is fundamental to all sample-based estimators of mutual information (and R-D), including ours and NERD, and neural methods are more suitable in certain regimes, e.g., for smooth, high-dimensional distributions. On the other hand, our results demonstrate that directly optimizing over particles can be more effective, for instance in settings where the optimal reproduction distribution is concentrated on low dimensional manifolds and/or have many widely-separated modes. This gives our method several distinct advantages over neural approaches:\n1. **Ease of use for practitioners.** Our method completely removes the need to specify neural net architectures, and involves only choosing 1). a learning rate schedule (which we automate with adaptive methods like Adam) and 2). the number of particles $n$, which controls the solution quality. In our experience, the neural methods also tend to be harder to train: NERD requires experimenting with $n$ to ensure accurate gradient estimation (see the failure case in the left panel of Fig. 1 of the main text), and numerical issues from deep neural density models (deep normalizing flows) can also cause the training of RD-VAE to diverge. \n\n2. **Improved R-D bounds with lower computation complexity.** Without much tuning, our method consistently produces R-D bounds comparable with or better than NERD, as well as RD-VAE up to the rate limit of log(n). At the same approximation quality, our method also tends to be more memory and computationally efficient in our experiments. We give more insight into this in our next response on computation efficiency.  \n\nFinally, we want to highlight another aspect of our contribution to R-D theory, which is the introduction of a new, rich class of examples with known ground truth as a benchmark for algorithms. We discuss it more prominently in our updated manuscript, as explained in point # 5 of our main rebuttal. \n\n> The computational (specifically, memory) cost of the particle method was not discussed, and this should be mentioned, especially if each particle consists of high-dimension vectors. The benefit of the neural methods is that one does not need to maintain many high dimensional vectors.\n\nWe now added a discussion in Section 4 about the computational complexity of WGD vs. the neural methods. Although WGD maintains potentially high-dimensional vectors, its per-iteration complexity (both in computation and memory) is largely comparable to that of NERD given the same $n$; this is because in both methods, the main bottleneck is in computing a pairwise distortion matrix whose cost scales as $O(m n d)$. Depending on the size of the generator network, NERD can be computationally more expensive than WGD, as it additionally requires backpropagation through the network.   RD-VAE also computes a distortion matrix but essentially with $n=1$; however its computation and memory requirements are dominated by that of training (typically large) neural density estimators.  \nAs different network architectures are used on different data sources (with varying levels of GPU support for operations such as convolution), it can be difficult to compare the computational complexity of WGD with the neural methods; nonetheless, we make a timing comparison in Fig 2 of the main rebuttal PDF, where the per-iteration time on MNIST is roughly 1:1.4:5 for WGD:RD-VAE:NERD.  \nFinally, we remark that maintaining particles is inherently more efficient in the low distortion setting where the source consists of distinct atoms; in this regime, the neural methods also have to memorize the training data, but can only store it indirectly in the network parameters. \n\n\n> Do the neural methods (NERD and Yang and Mandt) provide tighter bounds as the network complexity or architecture increases? If complex enough, could they be comparable with the particle method as well?\n\nIn theory, yes, all three methods (NERD, RD-VAE, and ours) converge to the true R-D in the ideal limit of infinite capacity / number of particles & samples. However, in practice it is a different question how efficiently each method estimates R-D given a computation budget. In our experience, increasing the network complexity for the neural methods can help produce tighter bounds, but this does not seem to happen consistently for various reasons. \n\n- On many experiments, the neural architectures we borrowed from the original papers appear to have been well optimized, and tweaking them yielded little or no gains (e.g., the DCGAN architecture for NERD on MNIST, or the hierarchical VAE for RD-VAE). \n- Increasing the network complexity may also make the training more difficult or even diverge, which we observed with RD-VAE.\n- In settings where the optimal reproduction distribution is concentrated on low dimensional manifolds and/or have many widely-separated modes, the neural approach is inherently less efficient because of the built-in bias of neural nets towards learning smooth mappings. The right-hand panel in Fig 1 shows one example where the optimal $\\nu$ is concentrated on the unit circle and is recovered by WGD, while the neural methods produce $\\nu$ smeared around the circle.\nBy contrast, we observe consistently improving R-D estimates when increasing the number of particles in our approach. \n\n"}, {"review_id": "KUr0LhjQu3", "reviewer": "Reviewer_SDqV", "summary": "This paper proposes estimating the rate distortion function using Wesserstein gradient descent. Different from the classical Blahut-Arimoto algorithm in which the support points are fixed, the proposed method is able to learn the support of the optimal distribution. The authors prove finite-sample complexity bounds, and also conduct experiments demonstrating the superiority of the proposed method.", "questions": "1. As the authors mention in Sec. 4.2, the hybrid algorithm alternates between WGD and BA. While this seems like a heuristics-based method, would this hurt the convergence properties of WGD?\n\n2. Lemma 4.2 only implies that WGD would converge to stationary points. Is it possible that under suitable conditions on the objective function, it will converge to minimizers?", "limitations": "The authors have adequately addressed the limitations and potential negative societal impact.", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. By drawing interesting connections between the R-D problem and entropic optimal transport, this paper proposes a novel algorithm which is conceptually simple but effective.\n\n2. The authors combine the advantage of Wasserstein gradient descent with BA and come up with a hybrid method that can update both the particle weights and the support of the distribution.\n\n3. In the experiments section, the authors make a rigorous comparision of the proposed method with various existing methods, from the classic BA algorithm to modern neural network based approaches. ", "weaknesses": "While this paper is technically solid, some parts of the paper is hard to read. For example $I(X;Y)$ in eq.(1) is not defined, and the boundeness assumption in Lemma 4.2 my need additional explanation. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your helpful feedback. We have incorporated your suggestions to improve the writing.  Below we will go over all points raised and detail how they have been addressed.\n\n> Weaknesses:  \n> While this paper is technically solid, some parts of the paper is hard to read. For example \nI(X;Y) in eq.(1) is not defined, and the boundeness assumption in Lemma 4.2 my need additional explanation.\n\n$I(X;Y)$ stands for the mutual information of the joint distribution $P_X Q_{Y|X}$, which we now define after eq (1) in our updated manuscript. \nWe have also clarified the boundedness assumption in Lemma 4.2, specifically the condition $\\sup_t \\int \\| \\nabla V_{\\mathcal{L}}(\\nu^{(t)})\\|^2 \\,d\\nu^{(t)} \\leq \\infty$.  We replaced “assume that” with “suppose that”, to make it clear that this is the hypothesis of the lemma. This condition is satisfied when, e.g., the cost function $\\rho$ has a bounded derivative on the relevant domain, which is the case in our examples.\n\n> Questions:  \n> As the authors mention in Sec. 4.2, the hybrid algorithm alternates between WGD and BA. While this seems like a heuristics-based method, would this hurt the convergence properties of WGD?\n\nThis is an interesting question, and one we can only partly answer. \nThe hybrid algorithm 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 BA steps to not hurt -- but rather accelerate -- the convergence of WGD (see Section 5.2), additional effort is required to theoretically guarantee convergence of the hybrid algorithm.  \n\nThere are two key properties we use for the convergence of WGD: 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 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)} \\rightarrow 0$ for $t\\rightarrow \\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). We have added a discussion in Section 4.2, and leave further investigation to future work.\n\n\n> Lemma 4.2 only implies that WGD would converge to stationary points. Is it possible that under suitable conditions on the objective function, it will converge to minimizers?\n\nAs our objective is non-convex over its domain (probability measures supported on at most $n$ points), we could only show convergence to stationary points. This is a very common limitation in the literature on gradient descent algorithms. It is in contrast to the discrete R-D problem on a fixed finite alphabet, which is convex, finite-dimensional, and solvable (in principle) by the BA algorithm. \nIt is theoretically possible for our algorithm to get stuck at a non-minimizer, but we have yet to observe this on real-world problems.  \nIt is worth mentioning here that in the infinite particle limit, Yan et al. 2023 recently proved that Wasserstein-Fisher-Rao gradient flow for the BA functional can only converge to a global minimum (if it converges). This is important insofar as our hybrid algorithm could be interpreted as an implementation of gradient descent in this Wasserstein-Fisher-Rao geometry. We discuss this in the revised version at the end of Section 3.\n"}, {"review_id": "BpoV6uzFvb", "reviewer": "Reviewer_GB3T", "summary": "The paper proposes an algorithm for estimating the rate-distortion function for (possibly) continuous sources. This is done by presenting the R-D as an entropy-regularized optimal transport problem, and solving via Wasserstein GD. The source distribution is approximated by empirical distributions. This yields an upper bound on the true R,D. Some latest results from EOT are adapted in order to bound the theoretical deviation from the real R-D for sub-Gaussian sources. The convergence of the method is demonstrated through simulations.", "questions": "The paper contains a decent theoretical contribution, but the writing needs a total makeover. I would recommend an acceptance only if this issue could be fixed for the final version, which I doubt.", "limitations": "The Algorithm outputs only the marginal distribution of the reproduction and provides an upper bound for the R-D function only where the Blahut-Arimoto step is tractable. Neither this, nor other limitations, are discussed.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 1, "contribution": 3, "strengths": "The proposed method seems novel, and theoretical contribution about sample complexity (Prop. 4.3) seems important.", "weaknesses": "My concerns here are twofold:\n\n1.\tFirst, the proposed Algorithm outputs the marginal distribution of the reproduction. This allows to compute an upper bound for the R-D function (as described in Sec. 4.4), but only where the Blahut-Arimoto step (11) is tractable, which is a major drawback. It might also incur in an additional error (beyond the result of Prop. 4.3), which is not discussed.\n\n2.\tSecond, and more crucial, Regarding the clarity of the text itself. Writing is *very* unclear and not well-organized. Just to name a few issues: The definition for the Wasserstein gradient is not clear, and it is unclear how to compute it. $\\psi^t$ is used sometimes for a gradient, and sometimes for a potential (whose calculation is also not explained). Eq. (15) is missing some relation sign (probably $\\leq$). In Lemma 4.2, $k$ is not defined (a typo, maybe?). Many cases of \"it is well known\" or \"we know\" without any reference. Many more typos. I'm afraid this does not meet the quality requirements for NeuRIPS.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your helpful suggestions, which have significantly improved our manuscript. Below we go over your concerns and detail how we have addressed them in our updated manuscript.\n\n> First, the proposed Algorithm outputs the marginal distribution of the reproduction. This allows to compute an upper bound for the R-D function (as described in Sec. 4.4), but only where the Blahut-Arimoto step (11) is tractable, which is a major drawback. It might also incur in an additional error (beyond the result of Prop. 4.3), which is not discussed.\n\nOur original writing appears to have been misleading; we have rephrased it to emphasize that the upper bound computation in our method is in fact tractable, and there is no additional error incurred. While the BA-like step (11) involves an integral w.r.t. the $\\nu$ distribution, in our algorithm $\\nu$ is represented as $n$ particles, so the integral reduces to a finite sum, and this step is tractable and no more expensive than one update in the WGD algorithm.  Our original writing on line 242 (“provided the integral in the numerator can be computed exactly”) referred to the general setting which applied to both our method and NERD. Indeed, NERD uses a continuous $\\nu$ and requires additional approximation for this computation (as explained on lines 244-248). \n\n\n> Second, and more crucial, Regarding the clarity of the text itself. Writing is very unclear and not well-organized...\n\nWe sincerely apologize for the unclear writing, numerous typos and poor notation in our rushed write-up ahead of the submission deadline.  We have since made extensive revisions to improve the writing and correct all errors. Going over the list of issues (also see overall response #1, 2):\n> The definition for the Wasserstein gradient is not clear, and it is unclear how to compute it.\n\nWe revised our introduction of the Wasserstein gradient, separating the formal definition (Definition 4.1; as the Euclidean gradient of the first variation) from its linearization property (eq 15). We explain that the former gives us the computational recipe, while the latter characterization justifies its role as a bona fide gradient (i.e., “taking a small enough step along the Wasserstein gradient indeed decreases the loss functional”). We add references to [Ambrosio et al. 2008, Definition 10.1.1], [Chizat 2022, Lemma A.2] and [Carlier et al. 2022, Proposition 4.2] to precisely relate the definition and the linearization property in our setting.\n\n> $\\psi^t$ is used sometimes for a gradient, and sometimes for a potential (whose calculation is also not explained).\n\nThis was an unfortunate notation clash; we have now changed the notation for the Wasserstein gradient to $\\Psi$. We improved the discussion around its calculation based on Def 4.1, clarifying that it comes from a straightforward differentiation of the Sinkhorn potential (i.e. the output of Sinkhorn’s algorithm) in the case of $\\mathcal{L}\\_{EOT}$, and we added in Section 4.1 an explicit formula in the case of $\\mathcal{L}\\_{BA}$  (originally given by eq. 23 in the Supplementary Material). \n\n> Eq. (15) is missing some relation sign (probably).\n\nWe now added the missing comma. \n> In Lemma 4.2, k is not defined (a typo, maybe?).\n\nTypo indeed, corrected as $t$. \n\n> Many cases of \"it is well known\" or \"we know\" without any reference.\n\nWe have eliminated all instances of \"it is well known\" or \"we know\" and replaced them with detailed references.\n\n**References**\n\nLénaïc Chizat. Mean-field langevin dynamics: Exponential convergence and annealing. arXiv:2202.01009, 2022\n\nLuigi Ambrosio, Nicola Gigli, and Giuseppe Savaré. Gradient flows in metric spaces and in the space of probability measures. \n\nGuillaume Carlier, Lénaïc Chizat, and Maxime Laborde. Lipschitz continuity of the Schrödinger map in entropic optimal transport. arXiv:2210.00225, 2022\n\n"}, {"review_id": "XZ9bkd6yON", "reviewer": "Reviewer_4bjw", "summary": "The authors propose a novel method based on Wasserstein Gradient Descent for numerically computing the Rate-Distortion function. Judging from the experiments their method seems to achieve comparable performance to other state of the art methods while having lower computational cost and being conceptually simpler, which reduces the need for tuning (of neural network architectures).", "questions": "Some questions on the theory side:\n- In (15) in line 202, should there be an inequality between the first and second line?\n- Should the integral in Lemma 4.2 not be w.r.t \\nu^t and not \\nu^k ?\n- I am not quite sure I correctly understand the implications of Proposition 4.3. \nHow can anything be implied about WGD if we do not know that WGD converges to a global optimizer i.e. how does WGD relate to \\min considered in the bound of Proposition 4.3 ?\nTo me right now it looks like a theoretical heuristic.\n\nSmall Typos:\n- Line 172: baised should be biased\n- Line 221: optimie should be optimize\n- Line 302: Should \"articles\" be particles?\n- Line 389: One of the authors is shown as Le?nai?c  (not sure if this might be caused on my end)", "limitations": "It would be interesting to have more theoretical guarantees for this method and/or more experiments on other realistic data.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The method is novel and well motivated\n- Good choice of numerical experiments", "weaknesses": "Lack of theoretical analysis/guarantees  (see also questions)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the insightful review. We have incorporated all the typographical/notational suggestions as well as added a new experiment on MNIST (see PDF attached to the main rebuttal). Below we go over all the concerns and how we addressed them in our updated manuscript. \n\n> In (15) in line 202, should there be an inequality between the first and second line?\n \nThere should have been a comma; fixed.\n\n> Should the integral in Lemma 4.2 not be w.r.t \\nu^t and not \\nu^k ?\n\nYes; fixed.\n\n> I am not quite sure I correctly understand the implications of Proposition 4.3. How can anything be implied about WGD if we do not know that WGD converges to a global optimizer i.e. how does WGD relate to \\min considered in the bound of Proposition 4.3 ? To me right now it looks like a theoretical heuristic.\n\nProposition 4.3 is a statement about the approximation capabilities of the extremum estimator $\\min\\_{\\nu \\in \\mathcal{P}\\_n} \\mathcal{L}(\\nu)$ rather than the actual algorithm for computing it. We have now made this aspect clearer in the updated manuscript. \n\nIn statistics, it is common to consider the theoretical merit (e.g., consistency, bias) of an estimator, such as the MLE, separately from the computational procedure for obtaining it; similarly, machine learning often studies the approximation ability of a hypothesis class (such as neural networks) separately from the training problem (e.g., the numerics of SGD). Our Proposition 4.3 is a statement of the former kind. It assures the consistency of the R-D estimator $\\min_{\\nu \\in \\mathcal{P}_n} \\mathcal{L}(\\nu)$ (defined by \"the best R-D loss achievable over all reproduction measures supported on at most $n$ points\") and further quantifies its rate of convergence in terms of: 1). the number of source samples ($m$), characterizing its statistical efficiency; and 2). the number of particles used ($n$), which implies a statement of universal approximation analogous to that of neural networks. While the *asymptotic* consistency of such an R-D estimator has been shown in information theory literature [Harrison and Kontoyiannis, 2008], we are not aware of any previous *finite-sample* bounds like ours. \n\n----\n**References**\n\nMatthew T. Harrison and Ioannis Kontoyiannis. Estimation of the rate–distortion function. IEEE Transactions on Information Theory, 54(8):3757–3762, 2008.\n"}, {"review_id": "1ImFpj7uFY", "reviewer": "Reviewer_RV2s", "summary": "This paper describes the elegant connection between entropic OT and estimation of rate-distortion functions, as well as proposes a novel algorithm based on Wasserstein Gradient Descent (WGD). It outlines prior methods (e.g. Blahut-Arimoto, NERD) in the same (or similar) language, and performs WGD on the two objective functions of interest, detailing the gradient calculations, and provides numerical experiments comparing the methods to prior works in the literature. In these experiments, we see that WGD is able to provide good bounds on the rate distortion function\n", "questions": "Would it be interesting to analyze the \"debiased\" entropic OT dynamics for this problem as in e.g. [Pooladian2022Debiaser] and [Feydy2019Interpolating]?\n\n@inproceedings{pooladian2022debiaser,\n  title={Debiaser beware: Pitfalls of centering regularized transport maps},\n  author={Pooladian, Aram-Alexandre and Cuturi, Marco and Niles-Weed, Jonathan},\n  booktitle={International Conference on Machine Learning},\n  pages={17830--17847},\n  year={2022},\n  organization={PMLR}\n}\n\n@inproceedings{feydy2019interpolating,\n  title={Interpolating between optimal transport and mmd using sinkhorn divergences},\n  author={Feydy, Jean and S{\\'e}journ{\\'e}, Thibault and Vialard, Fran{\\c{c}}ois-Xavier and Amari, Shun-ichi and Trouv{\\'e}, Alain and Peyr{\\'e}, Gabriel},\n  booktitle={The 22nd International Conference on Artificial Intelligence and Statistics},\n  pages={2681--2690},\n  year={2019},\n  organization={PMLR}\n}\n\n\nMinor comments:\n\nL33: “R-D” function instead of $R(D)$ (L86, L99, L101, L110, …) [just to be consistent]\n\nL221: “optimize”\n\nMissing \\mathbb{E} in the statement of Proposition 4.3? You also could write this as one line, saving a lot of space (since the rates are basically the same for all the statements).\n", "limitations": "Yes", "rating": 7, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This paper wrote down a very elegant connection between rate distortions + entropic optimal transport -- as someone who studies the latter, I think this is great. I think it's surprising this connection has not been made explicit between the two communities. Their proposed method is also quite insightful. Numerical and statistical results are convincing as well.", "weaknesses": "There is minimal technical novelty in this paper, though I think that's completely fine given the newfound connections the authors make bridging two academic communities. I think the references are a bit thin, especially on the level of the Wasserstein gradient for $\\mathcal{L}_{EOT}$. A recent paper, for example, is [Rigollet2022Sample], as well as some of the papers it cites, and the papers that cite it.\n\n@article{rigollet2022sample,\n  title={On the sample complexity of entropic optimal transport},\n  author={Rigollet, Philippe and Stromme, Austin J},\n  journal={arXiv preprint arXiv:2206.13472},\n  year={2022}\n}\n\nIt is also worth pointing out that this constant \"$C$\" in (15) of the text scales exponentially poorly with the regularization parameter $\\epsilon$. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the detailed and insightful comments.  We have corrected all the typographical issues as you suggested, and address all your remaining concerns and questions below.\n\n> I think the references are a bit thin, especially on the level of the Wasserstein gradient for $\\mathcal{L}\\_{EOT}$. \n\nThank you for pointing us to the more recent references, which we have included in our revision. Particularly, we added more recent papers like [Chizat 2022] and [Yan et al. 2023] in relation to gradient flows for $\\mathcal{L}\\_{EOT}$ and $\\mathcal{L}\\_{BA}$ . We further added references to [Genevay et al, 2019] and [Rigollet and Stromme, 2022] in relation to the sample complexity result. The result of [Rigollet and Stromme, 2022] is particularly interesting because it could be used to derive a version of our Proposition 4.3 with distortion functions other than the quadratic.\n\n> It is also worth pointing out that this constant \"C\" in (15) of the text scales exponentially poorly with the regularization parameter \n\nThis is true. However, unlike the case where one uses EOT as an approximation of OT, in the compression problem we are not only interested in the small regularization regime. The setting of low rate and high distortion (corresponding to large entropic regularization) has received increasing interest from the generative modeling community [Mentzer et al., 2020; Yang et al., 2023], where techniques such as GANs and diffusion models allow compression algorithms to achieve realistic image reconstructions with extremely low bit-rates, at the cost of a large squared-error distortion. In that sense, the poor scaling of C is less problematic for R-D estimation than for other applications of EOT.\n\n\n> Questions:  \n> Would it be interesting to analyze the \"debiased\" entropic OT?\n\nWe focus on the “biased” entropic OT distance since it has the exact mathematical form of the lossy compression / R-D problem. The “debiased” entropic OT distances may also admit information-theoretic interpretations and would be interesting to investigate in future work. \n\n\n> Missing \\mathbb{E} in the statement of Proposition 4.3? \n\nThe first line of Proposition 4.3 does not have an expectation sign because there the LHS is already deterministic.  The following lines involve the $m$-sample empirical measure $\\mu^m$, which is random and we take expectation over the samples.\n\n\n----\n**References**\n\nLénaïc Chizat. Mean-field langevin dynamics: Exponential convergence and annealing. arXiv:2202.01009, 2022\n\nYuling Yan, Kaizheng Wang, and Philippe Rigollet. Learning gaussian mixtures using the wasserstein-fisher-rao gradient flow. arXiv:2301.01766, 2023\n\nAude Genevay, Lénaic Chizat, Francis Bach, Marco Cuturi, and Gabriel Peyré. Sample complexity of Sinkhorn divergences. AISTATS, 2019\n\nPhilippe Rigollet and Austin J Stromme. On the sample complexity of entropic optimal transport. arXiv:2206.13472, 2022\n\nMentzer, F., Toderici, G. D., Tschannen, M., and Agustsson, E. High-fidelity generative image compression. Advances in Neural Information Processing Systems, 33:11913–11924, 2020.\n\nYang, Y., Mandt, S., and Theis, L. An introduction to neural data compression. arXiv:2202.06533,2022\n"}]

Author rebuttal:
We thank all reviewers for taking the time to review this manuscript and for providing insightful feedback, which has strengthened our submission.

We appreciate that the reviewers recognized our proposed method as “principled” (**jdNJ**), “novel” (**SDqV, GB3T, 4bjw, RV2s**), and “technically solid” (**jdNJ, SDqV, GB3T, RV2s**), and found our numerical results convincing (**jdNJ, SDqV, 4bjw, RV2s**). Moreover, our sample complexity result based on entropic optimal transport (Prop 4.3) “is the first such result in this line of work” (**jdNJ**) and “important” (**GB3T**). 

The reviewers also pointed to some parts of our manuscript that were unclear, particularly the tractability of our upper bound estimate and the computational aspects of the Wasserstein gradient. We have taken the suggestions seriously and improved our manuscript accordingly. As we are not given the option to share our updated manuscript, we summarize the main improvements below:

1. [in response to **GB3T**] We have addressed the confusion around line 242 and made it clear that the upper bound from our method is in fact tractable. While the BA-like step (11) involves an integral w.r.t. the $\nu$ distribution, in our algorithm $\nu$ is represented as $n$ particles, so the integral reduces to a finite sum and this step is no more expensive than one update in the WGD algorithm.  Our original writing on line 242 referred to the general setting which applied to both our method and NERD. Indeed, NERD uses a continuous $\nu$ and requires additional approximation for this computation (as explained in lines 244-248). 
2. [in response to **GB3T**] We have extensively revised the discussion around WGD to make it clearer and more accessible. 
    - In our updated version, we now preface the definition of the Wasserstein gradient with the intuition of gradient flow in the space of probability measures.
    - We added a new explanation that the formal definition (as the Euclidean gradient of the first variation of the loss functional) is the computational basis of the algorithm, whereas the linearization property of the Wasserstein gradient in eq 15 justifies its role as a bona fide gradient and is what allows us to prove convergence of WGD to a stationary point (Lemma 4.2). We added detailed references to [Ambrosio et al. 2008, Definition 10.1.1], [Chizat 2022, Lemma A.2] and [Carlier et al. 2022, Proposition 4.2] to precisely relate the definition and the linearization property in our setting.
   - On the computation of the Wasserstein gradient (WG): for the EOT functional, we now state clearly that its WG is computed by taking the Euclidean gradient of the output of the Sinkhorn algorithm.  For the BA functional, the WG has an explicit formula (derived in Supplementary Material eq 23) which we now include in Section 4.1 of the main text.

3. [in response to **jdNJ, RV2s**] We expanded our references to include related work on the statistical complexity of EOT (Genevay et al. 2019, Riollet and Stromme 2022), non-parametric Gaussian mixture estimation (Yan et al., 2023), as well as recent papers in rate-distortion theory (Wu et al. 2022, Lei et al. 2023) which also note a connection between R-D and EOT. Wu et al. study the finite (and known) alphabet setting, similarly to the BA algorithm, and Lei et al. discuss the relation to optimal scalar quantization. By contrast, our work solves the computational problem of R-D estimation with Wasserstein gradient descent, with a focus on the continuous, high-dimensional, and stochastic optimization setting. Besides our methodological contribution, we also connect literature on related problems in statistics, information theory, and entropic optimal transport, and actually leverage the latter connection to characterize the sample complexity and approximation quality of our R-D estimator. 

4. [in response to **4bjw**] To demonstrate the effectiveness and scalability of our method on more real-world data, we added a new experiment on the MNIST dataset. Again we obtain tighter R-D bounds with less computation than alternative methods. **Please see the attached PDF**.

5. [in response to **jdNJ**] One contribution to the R-D literature, which was not sufficiently emphasized in our first manuscript, is that we introduce a new, rich class of examples with known ground truth as a benchmark for algorithms. While the R-D literature usually uses Gaussian, Laplacian, or Bernoulli sources for that purpose, we leverage the connection with maximum likelihood deconvolution to find that a Gaussian convolution of *any* distribution can serve as a source with a known optimal reproduction distribution. While we already used this result to assess the optimality of various algorithms on the circular source in Section 5.2, we now discuss it more prominently in the updated manuscript. 

----
**References**

Lénaïc Chizat. Mean-field langevin dynamics: Exponential convergence and annealing. arXiv:2202.01009, 2022

Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré. Gradient flows in metric spaces and in the space of probability measures. 

Guillaume Carlier, Lénaïc Chizat, and Maxime Laborde. Lipschitz continuity of the Schrödinger map in entropic optimal transport. arXiv:2210.00225, 2022

Aude Genevay, Lénaic Chizat, Francis Bach, Marco Cuturi, and Gabriel Peyré. Sample complexity of Sinkhorn divergences. AISTATS, 2019

Philippe Rigollet and Austin J Stromme. On the sample complexity of entropic optimal transport. arXiv:2206.13472, 2022


Yuling Yan, Kaizheng Wang, and Philippe Rigollet. Learning gaussian mixtures using the wasserstein-fisher-rao gradient flow. arXiv:2301.01766, 2023

Shitong Wu, Wenhao Ye, Hao Wu, Huihui Wu, Wenyi Zhang, and Bo Bai. A communication optimal transport approach to the computation of rate distortion functions. arXiv:2212.10098, 2022

Eric Lei, Hamed Hassani, and Shirin Saeedi Bidokhti. On a relation between the rate-distortion function and optimal transport. arXiv:2307.00246, 2023




Meta-review: A novel method is proposed for computing the rate-distortion function with inspirations from Wasserstein gradient flows. This is an interesting bridge between optimal transport and information theory. Although the proposed method involves a non-convex optimization, experiments show the proposed method outperforms some existing methods. The authors addressed initial concerns in the rebuttal.
Decision: Accept (poster)