Title: Estimating the Rate-Distortion Function by Wasserstein Gradient Descent
Venue: NeurIPS 2023
Accepted: yes
Keywords: information theory

rate-distortion function

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

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.

1 Introduction

The rate-distortion (R-D) function R ( D ) occupies a central place in the theory of lossy compression. For a 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 is 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 a line of research has sought to instead estimate R ( D ) from data samples [Harrison and Kontoyiannis, 2008, Gibson, 2017], with recent methods [Yang and Mandt, 2022, Lei et 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 is the formulation of the R-D problem as the minimization of a certain rate functional [Harrison and Kontoyiannis, 2008] over the space of probability measures on the reproduction alphabet. Optimization over such an infinite-dimensional space has long been studied under gradient flows [Ambrosio et al., 2008], and we consider a concrete algorithmic implementation based on moving particles in 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-network-free R ( D ) upper bound estimator for continuous alphabets. We implement the estimator by Wasserstein gradient descent (WGD) over the space of reproduction

distributions. Experimentally, we found the method to converge much more quickly than state-of-theart 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 on a 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, as a benchmark for algorithms. While the literature relies on the Gaussian or 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 source 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 a 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 a projection of the source distribution under entropic optimal transport; this is a key to our sample complexity results in Sec. 4.3. Lastly, Sec 2.3 gives a statistical interpretation of R-D as maximum-likelihood deconvolution and uses it 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 a memoryless data source X with distribution P X , its rate-distortion (R-D) function describes the minimum possible number of bits per sample needed to reproduce the source within a prescribed distortion threshold D . Let the source and reproduction take values in two sets X and Y , known as the source and reproduction alphabets , and let ρ : ( X , Y ) → [0 , ∞ ) be a given distortion function. The R-D function is defined by the following optimization problem [Polyanskiy and Wu, 2022],

where Q Y | X is any Markov kernel from X to Y conceptually associated with a (possibly) stochastic compression algorithm, and I ( X ; Y ) is the mutual information of the joint distribution P X Q Y | X .

For ease of presentation, we now switch to a more abstract notation without reference to random variables. We provide the precise definitions in the Supplementary Material. Let X and Y be standard Borel spaces; let µ ∈ P ( X ) be a fixed probability measure on X , which should be thought of as the source distribution P X . For a measure π on the product space X ×Y , the notation π 1 (or π 2 ) denotes the first (or second) marginal of π . For any ν ∈ P ( Y ) , we denote by Π( µ, ν ) the set of couplings between µ and ν (i.e., π 1 = µ and π 2 = ν ). Similarly, Π( µ, · ) denotes the set of measures π with π 1 = µ . Throughout the paper, K denotes a transition kernel (conditional distribution) from X to Y , and µ ⊗ K denotes the product measure formed by µ and K . Then R ( D ) is equivalent to

where H denotes relative entropy, i.e., for two measures α, β defined on a common measurable space, H ( α | β ) := ∫ log( dα dβ ) dα when α is absolutely continuous w.r.t β , and infinite otherwise.

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

Geometrically, F λ ( µ ) ∈ R is the y-axis intercept of a tangent line to the R ( D ) with slope -λ , 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 λ (e.g., we write F ( µ ) = F λ ( µ ) ) whenever it is harmless.

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

where we refer to the optimization objective L BA as the rate function [Harrison and Kontoyiannis, 2008]. We abuse the notation to write L BA ( ν ) := L BA ( µ, ν ) 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 L BA ( µ, ν ) is the asymptotically optimal cost of lossy compression of data X ∼ µ using a random codebook constructed from samples of ν [Dembo and Kontoyiannis, 2002]. Notably, the optimization in (5) can be solved analytically [Csiszár, 1974a, Lemma 1.3], and L BA simplifies to

In practice, the source µ is only accessible via independent samples, on the basis of which we propose to estimate its R ( D ) , or equivalently F ( µ ) . Let µ m denote an m -sample empirical measure of µ , i.e., µ m = ∑ m i =1 δ x i with x 1 ,...,n 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 ( µ ) : (1) the non-parametric estimator F ( µ m ) , and (2) the parametric estimator F H ( µ m ) := inf ν ∈H L BA ( µ m , ν ) , where H is a family of probability measures on Y . Harrison and Kontoyiannis [2008] showed that under rather broad conditions, both kinds of estimators are strongly consistent, i.e., F ( µ m ) converges to F ( µ ) (and respectively, F H ( µ m ) to F H ( µ ) ) with probability one as m →∞ . Our algorithm will implement the parametric estimator F H ( µ m ) with H chosen to be the set of probability measures with finite support, and we will develop finite-sample convergence results for both kinds of estimators in the continuous setting (Proposition. 4.3).

2.2 Connection to entropic optimal transport

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

We now consider the problem of projecting µ onto P ( Y ) under the cost L EOT :

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

With the identification glyph[epsilon1] = λ -1 , problem (8) is in fact equivalent to the R-D problem (4): compared to L BA (5), the extra constraint on the second marginal of π in L EOT (7) is redundant at the optimal ν . More precisely, Lemma 7.1 shows that (we omit the notational dependence on µ when it is fixed):

Existence of a minimizer holds under mild conditions, for instance if X = Y = R d and ρ ( x, y ) is a coercive lower semicontinuous function of y -x [Csiszár, 1974a, p. 66].

2.3 Connection to maximum-likelihood deconvolution

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

Maximum-likelihood deconvolution is a classical problem of non-parametric statistics and mixture models [Carroll and Hall, 1988, Lindsay and Roeder, 1993]. The deconvolution problem is concerned with estimating an unknown distribution α from noise-corrupted observations X 1 , X 2 , ... , where for each i ∈ N , we have X i = Y i + N i , Y i i.i.d. ∼ α , and N i are i.i.d. independent noise variables with a known distribution. For concreteness, suppose all variables are R d valued and the noise distribution is N (0 , σ 2 I d ) with Lebesgue density φ σ 2 . Denote the distribution of the observations X i by µ . Then µ has a Lebesgue density given by the convolution α ∗ φ σ 2 ( x ) := ∫ φ σ 2 ( x -y ) α ( dy ) . Here, we consider the population-level (instead of the usual sample-based) maximum-likelihood estimator (MLE) for α :

and observe that ν ∗ = α . Plugging in the density φ σ 2 ( x ) ∝ e -1 2 σ 2 ‖ x ‖ 2 , we see that the MLE problem (10) is equivalent to the R-D problem (4) with ρ ( x, y ) = 1 2 ‖ x -y ‖ 2 , λ = 1 σ 2 , and L BA given by (6) in the form of a marginal log-likelihood. Thus the R-D problem has the interpretation of estimating a distribution from its noisy observations given through µ , assuming a Gaussian noise with variance 1 λ .

R

(

D

)

λ

=

σ

-

2

decreasing

λ

→

This connection suggests analytical solutions to the R-D problem for a variety of sources that arise from convolving an underlying distribution with Gaussian noise. Consider an R-D problem (4) with X = Y = R d , ρ ( x, y ) = 1 2 ‖ x -y ‖ 2 , and let the source µ be the convolution between an arbitrary measure α ∈ P ( Y ) and Gaussian noise with known variance σ 2 . E.g., using a discrete measure for α results in a Gaussian mixture source with equal covariance among its components. When λ = 1 σ 2 , we recover exactly the population-MLE problem (10) discussed earlier, which has the solution ν ∗ = α . While this allows us to obtain one point of R ( D ) , we can in fact extend this idea to any λ ≥ 1 σ 2 and obtain the analytical form for the corresponding segment of the R-D curve. Specifically, for any λ ≥ 1 σ 2 , applying the summation rule for independent Gaussians reveals the source distribution µ as

i.e., as the convolution between another underlying distribution α λ and independent noise with variance 1 λ . A solution to the R-D problem (3) is then analogously given by ν ∗ = α λ , with the corresponding optimal coupling given by ν ∗ ⊗ ˜ K , ˜ K ( y, dx ) = N ( y, 1 λ ) . 1 Evaluating the distortion and mutual information of the coupling then yields the R ( D ) point associated with λ . Fig. 1 illustrates the R ( D ) of a toy Gaussian mixture source, along with the ν ∗ estimated by our proposed WGDalgorithm (Sec. 4); note that ν ∗ transitions from continuous (a Gaussian mixture with smaller component variances) to singular (a mixture of two Diracs) at λ = σ -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 λ , BA carries out the optimization problem (3) via coordinate ascent. Starting from an initial measure ν (0) ∈ P ( Y ) , the BA algorithm at step t computes an updated pair ( ν ( t +1) , K ( t +1) ) as follows

When the alphabets are finite, the above computation can be carried out in matrix and vector operations, and the resulting sequence { ( ν ( t ) , K ( t ) ) } ∞ t =1 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., X = Y = R d , 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 a discretization step and then apply BA on the resulting discrete problem.

One standard discretization method is to tile up the alphabets with small bins [Gray and Neuhoff, 1998]. This quickly becomes infeasible as the number of dimensions increases. We therefore consider discretizing the data space X to be the support of training data distribution µ m , i.e., the discretized alphabet is the set of training samples; this can be justified by the consistency of the parametric R-D estimator F H ( µ m ) [Harrison and Kontoyiannis, 2008]. It is less clear how to discretize the reproduction space Y , especially in high dimensions. Since we 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 i , Y i ) ∼ µ ⊗ K ; in particular, the relative entropy can be written as H ( µ ⊗ K | µ ⊗ ν ) = ∫ ∫ log ( dK ( x, · ) dν ( y ) ) K ( x, dy ) µ ( dx ) , where the integrand is computed exactly via a density ratio. In practice, an alternative parameteriation is often used where the neural density estimators are defined on a lower dimensional latent space than the reproduction alphabet, and the resulting approach is closely related to V AEs [Kingma and Welling, 2013]. Yang and Mandt [2022] additionally propose a neural estimator for a lower bound on R ( D ) , based on a dual representation due to Csiszár [1974a]. NERD [Lei et al., 2023a]: Instead of working with the transition kernel K as in the RD-VAE, Lei et al. [2023a] considered optimizing the form of the rate functional in (6), via gradient descent on the parameters of ν parameterized by a neural network. Let ν Z be a base distribution over Z = R K , such as the standard Gaussian, and ω : Z → Y be a decoder network. The variational measure ν is then modeled as the image measure of ν Z under ω . To evaluate and optimize the objective (6), the intractable inner integral w.r.t. ν is replaced with a plug-in estimator, so that for a given x ∈ X ,

After training, we estimate an R-D upper bound using n samples from ν (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) on a discrete problem and discuss the connection to scalar quantization. We also experimented with estimating R ( D ) by solving the EOT problem (8), but found it computationally much more efficient to work with the rate functional (6), and we see the primary benefit of the EOT connection as bringing in tools from statistical OT [Genevay et al., 2019, Mena and Niles-Weed, 2019, Rigollet and Stromme, 2022] for R-D estimation. Outside of information theory: Rigollet and Weed [2018] note a connection between the EOT projection problem (8) and maximum-likelihood deconvolution (10); our work complements their perspective by re-interpreting both problems through the equivalent R-D problem. Unbeknownst to us at the time, Yan et al. [2023] proposed similar algorithms to ours in the context of Gaussian mixture estimation, which we recognize as R-D estimation under quadratic distortion (see Sec. 2.3). Their work is based on gradient flow in the Fisher-Rao-Wasserstein (FRW) geometry [Chizat et al., 2018], which our hybrid algorithm can be seen as implementing. Yan et al. [2023] prove that, in an

idealized setting with infinite particles, FRW gradient descent does not get stuck at local minima; by contrast, our convergence and sample-complexity results (Prop. 4.2, 4.3) hold for any finite number of particles. We additionally consider larger-scale problems and the stochastic optimization setting.

4 Proposed method

For our algorithm, we require X = Y = R d and ρ 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 ( Y ) → R over the space of probability measures, our algorithm essentially simulates the gradient flow [Ambrosio et al., 2008] of L and follows the trajectory of steepest descent in the Wasserstein geometry. In practice, we represent a measure ν ( t ) ∈ P ( Y ) by a collection of particles and at each time step update ν ( t ) in a direction of steepest descent of L as given by its (negative) Wasserstein gradient . Denote by P n ( R d ) the set of probability measures on R d that are supported on at most n points. Our algorithm implements the parametric R-D estimator with the choice H = P n ( R d ) (see discussions at the end of Sec. 2.1).

4.1 Wasserstein gradient descent (WGD)

Abstractly, Wasserstein gradient descent updates the variational measure ν to its pushforward ˜ ν under the map ( id -γ Ψ) , for a function Ψ : R d → R d called the Wasserstein gradient of L at ν (see below) and a step size γ . To implement this scheme, we represent ν as a convex combination of Dirac measures, ν = ∑ n i =1 w i δ x i with locations { x i } n i =1 ⊂ R d and weights { w i } n i =1 . The algorithm moves each particle x i in the direction of -Ψ( x i ) , more precisely, ˜ ν = ∑ n i =1 w i δ x i -γ Ψ( x i ) .

Algorithm 1 Wasserstein gradient descent

Since the optimization objectives (4) and (8) appear as integrals w.r.t. the data distribution µ , we can also apply stochastic optimization and perform stochastic gradient descent on mini-batches with size m . This allows us to handle a very large or infinite amount of data samples, or when the source is continuous. We formalize the procedure in Algorithm 1.

The following gives a constructive definition of a Wasserstein gradient which forms the computational basis of our algorithm. In the literature, the Wasserstein gradient is instead usually defined as a Fréchet differential (cf. [Ambrosio et al., 2008, Definition 10.1.1]), but we emphasize that in smooth settings, the given definition recovers the one from the literature (cf. [Chizat, 2022, Lemma A.2]).

Definition 4.1. For a functional L : P ( Y ) → R and ν ∈ P ( Y ) , we say that V L ( ν ) : R d → R is a first variation of L at ν if

We call its (Euclidean) gradient ∇ V L ( ν ) : R d → R d , if it exists, the Wasserstein gradient of L at ν .

For L = L EOT , the first variation is given by the Kantorovich potential, which is the solution of the convex dual of L EOT and commonly computed by Sinkhorn's algorithm [Peyré and Cuturi, 2019,

Nutz, 2021]. Specifically, let ( ϕ ν , ψ ν ) be potentials for L EOT ( µ, ν ) . Then V L ( ν ) = ψ ν is the first variation w.r.t. ν (cf. [Carlier et al., 2022, equation (20)]), and hence ∇ ψ ν is the Wasserstein gradient. This gradient exists whenever ρ is differentiable and the marginals are sufficiently light-tailed; we give details in Sec. 9.1 of the Supplementary Material. For L = L BA , the first variation can be computed explicitly. As derived in Sec. 9.1 of the Supplementary Material, the first variation at ν is

and then the Wasserstein gradient is ∇L BA ( ν ) = ∇ ψ ν . We observe that ψ ν ( y ) is computationally cheap; it corresponds to running a single iteration of Sinkhorn's algorithm. By contract, finding the potential for L EOT requires running Sinkhorn's algorithm to convergence.

Like the usual Euclidean gradient, the Wasserstein gradient can be shown to possess a linearization property, whereby the loss functional is reduced by taking a small enough step along its Wasserstein gradient. Following [Carlier et al., 2022], we state it as follows: for any ˜ ν ∈ P ( Y ) and π ∈ Π( ν, ˜ ν ) ,

The first line of (14) is proved for L EOT in [Carlier et al., 2022, Proposition 4.2] in the case that the marginals are compactly supported and ρ is twice continuously differentiable. In this setting, the second line of (14) follows using a 2 -b 2 = ( a + b )( a -b ) and a combination of boundedness and Lipschitz continuity of ∇ V L , see [Carlier et al., 2022, Proposition 2.2 and Corollary 2.4].

The linearization property given by (14) enables us to show that Wasserstein gradient descent for L EOT and L BA converges to a stationary point under mild conditions:

Proposition 4.2 (Convergence of Wasserstein gradient descent) . Let γ 1 ≥ γ 2 ≥ · · · ≥ 0 satisfy ∑ ∞ k =1 γ k = ∞ and ∑ ∞ k =1 γ 2 k < ∞ . Let L : P ( R d ) → R be Wasserstein differentiable in the sense that (14) holds. Denoting by ν ( t ) the steps in Algorithm 1, and suppose that L ( ν (0) ) is finite and ∫ ‖∇ V L ( ν ( t ) ) ‖ 2 dν ( t ) is bounded. Then

4.2 Hybrid algorithm

A main limitation of the BA algorithm is that the support of ν ( t ) 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 ν ( t ) , but not the weights, which are fixed to be uniform by default. We therefore consider a hybrid algorithm where we alternate between WGD and the BA update steps, allowing us to optimize the particle weights as well. Experimentally, this translates to faster convergence than the base WGD algorithm (Sec. 5.1). Note however, unlike WGD, the hybrid algorithm does not directly lend itself to the stochastic optimization setting, as BA updates on mini-batches no longer guarantee monotonic improvement in the objective and can lead to divergence. We treat the convergence of the hybrid algorithm in the Supplementary Material Sec. 9.4.

4.3 Sample complexity

Let X = Y = R d and ρ ( x, y ) = ‖ x -y ‖ 2 . Leveraging work on the statistical complexity of EOT [Mena and Niles-Weed, 2019], we obtain finite-sample bounds for the theoretical estimators implemented by WGD in terms of the number of particles and source samples. The bounds hold for both the R-D problem (4) and EOT projection problem (8) as they share the same optimizers (see Sec. 2.2), and strengthen existing asymptotic results for empirical R-D estimators [Harrison and Kontoyiannis, 2008]. We note that a recent result by Rigollet and Stromme [2022] might be useful for deriving alternative bounds under distortion functions other than the quadratic.

Proposition 4.3. Let µ be σ 2 -subgaussian. Then every optimizer ν ∗ of (4) and (8) is also σ 2 -subgaussian. Consider L := L EOT . For a constant C d only depending on d , we have

for all n, m ∈ N , where P n ( R d ) is the set of probability measures over R d supported on at most n points, µ m is the empirical measure of µ with m independent samples and the expectation E [ · ] is over these samples. The same inequalities hold for L := λ -1 L BA , with the identification glyph[epsilon1] = λ -1 .

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 := ∫ ρd ( µ ⊗ K ) and R := H ( µ ⊗ K | µ ⊗ ν ) lie on an upper bound of R ( D ) [Berger, 1971]. The two quantities can be estimated by standard Monte Carlo provided we can sample from µ ⊗ K and evaluate the density dµ ⊗ K dµ ⊗ ν ( x, y ) = dK ( x, · ) dν ( y ) .

When only ν is given, e.g., obtained from optimizing (6) with WGD or NERD, we estimate an R-D upper bound as follows. As in the BA algorithm, we construct a kernel K ν similarly to (11), i.e., dK ν ( x, · ) dν ( y ) = e -λρ ( x,y ) ∫ e -λρ ( x, ˜ y ) ν ( d ˜ y ) ; then we estimate ( D , R ) using the pair ( ν, K ν ) as described earlier.

As NERD uses a continuous ν , we follow [Lei et 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 to converge to a rate estimate of at most log( n ) , where n is the support size of ν . This is because as the algorithms approach an n -point minimizer ν ∗ n of the R-D problem, the rate estimate R approaches the mutual information of µ ⊗ K ν ∗ n , which is upper-bounded by log( n ) [Eckstein and Nutz, 2022]. In practice, this means if a target point of R ( D ) has rate r , then we need n ≥ e r to estimate it accurately.

4.5 Computational considerations

Common to all the aforementioned methods is the evaluation of a pairwise distortion matrix between m points in X and n points in Y , which usually has a cost of O ( mnd ) for a d -dimensional source. While RD-VAE uses n = 1 (in the reparameterization trick), the other methods (BA, WGD, NERD) typically use a much larger n and thus has the distortion computation as their main computation bottleneck. Compared to BA and WGD, the neural methods (RD-VAE, NERD) incur additional computation from neural network operations, which can be significant for large networks.

For NERD and WGD (and BA), the rate estimate upper bound of log( n ) nats/sample (see Sec. 4.4) can present computational challenges. To target a high-rate setting, a large number of ν particles (high n ) is required, and care needs to be taken to avoid running out of memory during the distortion matrix computation (one possibility is to use a small batch size m with stochastic optimization).

5 Experiments

We compare the empirical performance of our proposed method (WGD) and its hybrid variant with Blahut-Arimoto (BA) [Blahut, 1972, Arimoto, 1972], RD-VAE [Yang and Mandt, 2022], and NERD [Lei et al., 2022] on the tasks of maximum-likelihood deconvolution and estimation of R-D upper bounds. While we experimented with WGD for both L BA and L EOT , we found the former to be 10 to 100 times faster computationally while giving similar or better results; we therefore focus on WGDfor L BA in discussions below. For the neural-network baselines, we use the same (or as similar as possible) network architectures as in the original work [Yang and Mandt, 2022, Lei et al., 2023a]. We use the Adam optimizer for all gradient-based methods, except we use simple gradient descent

3.0

Blahut-Arimoto

WGD (proposed)

2.9

Hybrid algorithm (proposed)

Yang & Mandt (2022)

2.8

Loss

OPT

2.7

2.6

2.5

101

102

103

104

Iteration

Source samples

RD-VAE (Yang & Mandt, 2022)

Blahut-Arimoto

WGD (proposed)

Hybrid WGD (proposed)

10

1

10

2

10

3

10

-

2

10

-

1

Optimality gap to

L

(

ν

∗

)

d

= 2

10

1

10

2

10

3

10

-

2

10

-

1

10

0

d

= 5

10

1

10

2

10

3

10

-

1

10

0

d

= 10

10

1

10

2

10

3

10

-

1

10

0

10

1

λ

= 300 (

d

= 2)

Blahut-Arimoto

WGD

(proposed)

Hybrid WGD

(proposed)

n

(number of

particles)

Y

with a decaying step size in Sec. 5.1 to better compare the convergence speed of WGD and its hybrid variant. Further experiment details and results are given in the Supplementary Material Sec. 11.

5.1 Deconvolution

To better understand the behavior of the various algorithms, we apply them to a deconvolution problem with known ground truth (see Sec. 2.3). We adopt the Gaussian noise as before, letting α be the uniform measure on the unit circle in R 2 and the source µ = α ∗ N (0 , σ 2 ) with σ 2 = 0 . 1 .

We use n = 20 particles for BA, NERD, WGD and its hybrid variant. We use a two-layer network for NERD and RD-VAE with some hand-tuning (we replace the softplus activation in the original RD-VAE network by ReLU as it led to difficulty in optimization). Fig. 2 plots the resulting loss curves and shows that the proposed algorithms converge the fastest to the ground truth value OPT := L ( α ) . 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 a different solution where a cluster of particles covers the top-left portion of the circle with small weights while the remaining particles evenly covers the rest. Unlike our particle-based methods, the neural methods generally struggle to place the support of their ν exactly on the circle.

We additionally compare how the performance of BA, NERD, and the proposed algorithms scale to higher dimensions and a higher λ (corresponding to lower entropic regularization in L EOT ). Fig. 4 plots the gap between the converged and the optimal losses for the algorithms, and demonstrates the proposed algorithms to be more particle-efficient and scale more favorably than the alternatives which also use n particles in the reproduction space. We additionally visualize how the converged particles for our methods vary across the R-D trade-off in Fig. 6 of the Supplementary Material.

0

.

0000

0

.

0005

0

.

0010

0

.

0015

0

.

0020

0

.

0025

0

2

4

6

physics

0

.

5

1

.

0

1

.

5

2

.

0

0

10

20

speech

Ball´

e et al. 2021

Yang & Mandt

ˆ

RU

(

D

)

ˆ

RU

(

D

)

Proposed WGD

ˆ

RU

(

D

)

Yang & Mandt

ˆ

RL

(

D

)

Distortion (mean squared error)

Rate (nats per sample)

0

.

00

0

.

01

0

.

02

0

.

03

0

.

04

0

.

05

0

.

06

0

.

07

Distortion (mean squared error)

0

2

4

6

8

10

12

Rate (nats per sample)

R-D bounds on the MNIST training set

ˆ

R

U

(

D

)

, Yang & Mandt 2022

ˆ

R

U

(

D

)

, NERD (

n

= 40000

)

ˆ

R

U

(

D

)

, WGD (n=1000)

ˆ

R

U

(

D

)

, WGD (n=10000)

ˆ

R

U

(

D

)

,

WGD (n=40000)

ˆ

R

L

(

D

)

, Yang & Mandt 2022

10

2

10

3

10

4

10

5

Iteration

10

20

30

40

50

Loss

RD-VAE, Yang & Mandt 2022

WGD, proposed

10

1

10

2

10

3

10

4

Walltime (seconds)

10

20

30

40

50

Loss

RD-VAE, Yang & Mandt 2022

WGD, proposed

5.2 Higher-dimensional data

We perform R ( D ) estimation on higher-dimensional data, including the physics and speech datasets from [Yang and Mandt, 2022] and MNIST [LeCun et al., 1998]. As the memory cost to operating on the full datasets becomes prohibitive, we focus on NERD, RD-VAE, and WGD using mini-batch stochastic gradient descent. BA and hybrid WGD do not directly apply in the stochastic setting, as BA updates on random mini-batches can lead to divergence (as discussed in Sec. 4.2).

Fig. 5 plots the estimated R-D bounds on the datasets, and compares the convergence speed of WGD and neural methods in both iteration count and compute time. Overall, we find WGD to require minimal tuning and obtains the tightest R-D upper bounds within the log( n ) rate limit (see Sec. 4.4), and consistently obtains tighter bounds than NERD given the same computation budget.

6 Discussions

In this work, we leverage tools from optimal transport to develop a new approach for estimating the rate-distortion function in the continuous setting. Compared to state-of-the-art neural approaches [Yang and Mandt, 2022, Lei et al., 2022], our Wasserstein gradient descent algorithm offers complementary strengths: 1) It requires a single main hyperparameter n (the number of particles) and no network architecture tuning; and 2) empirically we found it to converge significantly faster and rarely end up in bad local optima; increasing n almost always yielded an improvement (unless the bound is already close to being tight). From a modeling perspective, a particle representation may be inherently more efficient when the optimal reproduction distribution is singular or has many disconnected modes (shown, e.g., in Figs. 1 and 3). However, like NERD [Lei et al., 2022], our method has a fundamental limitation - it requires an n that is exponential in the rate of the targeted R ( D ) point to estimate it accurately (see Sec. 4.4). Thus, a neural method like RD-V AE [Yang and Mandt, 2022] may still be preferable on high-rate sources, while our particle-based method stands as a state-of-the-art solution on lower-rate sources with substantially less tuning or computation requirements.

Besides R-D estimation, our algorithm also applies to the mathematically equivalent problems of maximum likelihood deconvolution and projection under an entropic optimal transport (EOT) cost, and may find other connections and applications. Indeed, the EOT projection view of our algorithm is further related to optimization-based approaches to sampling [Wibisono, 2018], variational inference [Liu and Wang, 2016], and distribution compression [Shetty et al., 2021]. Our particle-based algorithm also generalizes optimal quantization (corresponding to glyph[epsilon1] = 0 in L EOT and projection of the source under the Wasserstein distance [Graf and Luschgy, 2007, Gray, 2013]) to incorporate a rate constraint ( glyph[epsilon1] > 0 ), and it would be interesting to explore the use of the resulting rate-distortion optimal quantizer for practical data compression and communication.

References

CE Shannon. Coding theorems for a discrete source with a fidelity criterion. IRE Nat. Conv. Rec., March 1959 , 4:142-163, 1959.

Blahut. Computation of channel capacity and rate-distortion functions. IEEE Transactions on Information Theory , 18(4):460-473, 1972. doi: 10.1109/TIT.1972.1054855.

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 Ioannis Kontoyiannis. Estimation of the rate-distortion function. IEEE Transactions on Information Theory , 54(8):3757-3762, 2008. doi: 10.1109/tit.2008.926387.

Jerry Gibson. Rate distortion functions and rate distortion function lower bounds for real-world sources. Entropy , 19(11):604, 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 Saeedi 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 L Kontoyiannis. Source coding, large deviations, and approximate pattern matching. IEEE Transactions on Information Theory , 48(6):1590-1615, 2002.

Imre Csiszár. On an 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 in Machine Learning , 11(5-6):355-607, 2019.

Federico Bassetti, Antonella Bodini, and Eugenio Regazzini. On minimum kantorovich distance estimators. Statistics & probability letters , 76(12):1298-1302, 2006.

Martial Agueh and Guillaume Carlier. Barycenters in the wasserstein space. SIAM Journal on Mathematical Analysis , 43(2):904-924, 2011.

Alexander Alemi, Ben Poole, Ian Fischer, Joshua Dillon, Rif A Saurous, and Kevin Murphy. Fixing a 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 on Learning Representations , 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-likelihood deconvolution. Comptes Rendus Mathematique , 356(11-12):1228-1235, 2018.

Raymond J Carroll and Peter Hall. Optimal rates of convergence for deconvolving a density. Journal of the American Statistical Association , 83(404):1184-1186, 1988.

Bruce G Lindsay and Kathryn Roeder. Uniqueness of estimation and identifiability in mixture models. Canadian Journal of Statistics , 21(2):139-147, 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/TIT.1974.1055146.

Imre Csiszár. Information geometry and alternating minimization procedures. Statistics and Decisions, Dedewicz , 1:205-237, 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 preprint arXiv:1312.6114 , 2013.

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 J 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 wassersteinfisher-rao gradient flow. arXiv preprint arXiv:2301.01766 , 2023.

Lenaic Chizat, Gabriel Peyré, Bernhard Schmitzer, and François-Xavier Vialard. An interpolating distance between optimal transport and fisher-rao metrics. Foundations of Computational Mathematics , 18:1-44, 2018.

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

Marcel Nutz. Introduction to entropic optimal transport. Lecture notes, Columbia University , 2021. https://www.math.columbia.edu/~mnutz/docs/EOT_lecture_notes.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.00225 , 2022.

Toby Berger. Rate distortion theory, a mathematical basis for data compression . Prentice Hall, 1971.

Stephan Eckstein and Marcel Nutz. Convergence rates for regularized optimal transport via quantization. arXiv preprint arXiv:2208.14391 , 2022.

Eric Lei, Hamed Hassani, and Shirin Saeedi Bidokhti. Neural estimation of the rate-distortion function for massive datasets. In 2022 IEEE International Symposium on Information Theory (ISIT) , pages 608-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(11):2278-2324, 1998.

Andre Wibisono. Sampling as optimization in the space of measures: The langevin dynamics as a 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 Machine Learning Research , pages 2093-3027. PMLR, 06-09 Jul 2018. URL https:// proceedings.mlr.press/v75/wibisono18a.html .

Qiang Liu and Dilin Wang. Stein variational gradient descent: A 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 https://ee.stanford.edu/~gray/ gretsi.pdf .

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. Foundations and Trends® in Computer Graphics and Vision , 15(2):113-200, 2023. ISSN 1572-2740. doi: 10.1561/0600000107. URL http://dx.doi.org/10.1561/0600000107 .

David M Blei, Alp Kucukelbir, and Jon D McAuliffe. Variational inference: A review for statisticians. Journal of the American statistical Association , 112(518):859-877, 2017.

MJ Beal and Z Ghahramani. The variational bayesian em 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 on Machine Learning , pages 11106-11116. PMLR, 2020.

Pascal Vincent. A connection between score matching and denoising autoencoders. Neural computation , 23(7):1661-1674, 2011.

Michael I Jordan, Zoubin Ghahramani, Tommi S Jaakkola, and Lawrence K Saul. An introduction to variational methods for graphical models. Machine learning , 37:183-233, 1999.

Martin J Wainwright, Michael I Jordan, et al. Graphical models, exponential families, and variational inference. Foundations and Trends® in Machine Learning , 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 on Learning Representations , 2015.

John Platt and Alan Barr. Constrained differential optimization. In Neural Information Processing Systems , 1987.

George Papamakarios, Theo Pavlakou, and Iain Murray. Masked autoregressive flow for density estimation. In Advances in Neural Information Processing Systems , pages 2338-2347, 2017.

Roger B Grosse, Zoubin Ghahramani, and Ryan P Adams. Sandwiching the marginal likelihood using bidirectional monte carlo. arXiv preprint arXiv:1511.02543 , 2015.

James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http://github.com/google/jax .

Broader Impacts

Improved estimates on the fundamental cost of data compression aids the development and analysis of compression algorithms. This helps researchers and engineers make better decisions about where to allocate their resources to improve certain compression algorithms, and can translate to economic gains for the broader society. However, like most machine learning algorithms trained on data, the output of our estimator is only accurate insofar as the training data is representative of the population distribution of interest, and practitioners need to ensure this in the data collection process.

Supplementary Material for Estimating the Rate-Distortion Function by Wasserstein Gradient Descent

We review probability theory background and explain our notation from the main text in Section 7, give the formulas we used for numerically estimating an R ( D ) upper bound in Section 8, provide additional discussions and proofs regarding Wasserstein gradient descent in Section 9, elaborate on the connections between the R-D estimation problem and variational inference/learning in Section 10, provide additional experimental results and details in Section 11, and list an example implementation of WGD in Section 12. Our code and can be found at https://github.com/yiboyang/wgd .

7 Notions from probability theory

In this section we collect notions of probability theory used in the main text. See, e.g., [Çinlar, 2011] or [Folland, 1999] for more background.

Marginal and conditional distributions. The source and reproduction spaces X , Y are equipped with sigma-algebras A X and A Y , respectively. Let X × Y denote the product space equipped with the product sigma algebra A X ⊗A Y . For any probability measure π on X × Y , its first marginal is

which is a probability measure on X . When π is the distribution of a random vector ( X,Y ) , then π 1 is the distribution of X . The second marginal of π is defined analogously as

A Markov kernel or conditional distribution K ( x, dy ) is a map X × A Y → [0 , 1] such that

K ( x, · ) is a probability measure on Y for each x ∈ X ;

the function x ↦→ K ( x, B ) is measurable for each set B ∈ A Y .

When speaking of the conditional distribution of a random variable Y given another random variable X , we occasionally also use the notation Q Y | X from information theory [Polyanskiy and Wu, 2014]. Then, Q Y | X = x ( B ) = K ( x, B ) is the conditional probability of the event { Y ∈ B } given X = x .

Suppose that a probability measure µ on X is given, in addition to a kernel K ( x, dy ) . Together they define a unique measure µ ⊗ K on the product space X ×Y . For a rectangle set A × B ∈ A X ⊗A Y ,

The measure π := µ ⊗ K has first marginal π 1 = µ .

The classic product measure is a special case of this construction. Namely, when a measure ν on Y is given, using the constant kernel K ( x, dy ) := ν ( dy ) (which does not depend on x ) gives rise to the product measure µ ⊗ ν ,

Under mild conditions (for instance when X , Y are Polish spaces equipped with their Borel sigma algebras, as in the main text), any probability measure π on X × Y is of the above form. Namely, the disintegration theorem asserts that π can be written as π = π 1 ⊗ K for some kernel K . When π is the joint distribution of a random vector ( X,Y ) , this says that there is a measurable version of the conditional distribution Q Y | X .

Optimal transport. Given a measure µ on X and a measurable function T : X → Y , the pushforward (or image measure) of µ under T is a measure on Y , given by

If T is seen as a random variable and µ as the baseline probability measure, then T # µ is simply the distribution of T .

Suppose that µ and ν are probability measures on X = Y = R d with finite second moment. As introduced in the main text, Π( µ, ν ) denotes the set of couplings, i.e., measures π on X × Y with π 1 = µ and π 2 = ν . The 2-Wasserstein distance W 2 ( µ, ν ) between µ and ν is defined as

This indeed defines a metric on the space of probability measures with finite second moment.

We finish this section by giving a proof of the basic equivalence between optimization of L EOT and L BA , which goes back to [Csiszár, 1974a, Lemma 1.3 and subsequent discussion]. Recall that Π( µ, · ) denotes the set of measures on the product space X × Y with π 1 = µ .

Lemma 7.1. Set glyph[epsilon1] = λ -1 . It holds that

Proof. Both statements will follow from a simple property of relative entropy [Polyanskiy and Wu, 2022, Theorem 4.1, 'golden formula']: For π ∈ Π( µ, · ) with H ( π 2 | ν ) < ∞ , the properties of the logarithm reveal

and hence H ( π | µ ⊗ ν ) ≥ H ( π | µ ⊗ π 2 ) . This implies

Further, any optimizer ν ∗ of L BA with corresponding optimal 'coupling' π ∗ ∈ Π( µ, · ) must satisfy H ( π ∗ 2 | ν ∗ ) = 0 (otherwise, taking ν ∗ = π ∗ 2 has better objective) and thus π ∗ 2 = ν ∗ and π ∗ ∈ Π( µ, ν ∗ ) , therefore ν ∗ is also an optimizer of L EOT by the above equality. Conversely, any optimizer ν ∗ of L EOT with coupling π ∗ ∈ Π( µ, ν ∗ ) ⊂ Π( µ, · ) clearly yields a feasible solution for L BA as well, and hence is also an optimizer of L BA by the above equality.

8 Numerical estimation of rate and distortion from a reproduction distribution

Given a reproduction distribution ν and a kernel X ×A Y → [0 , 1] , the tuple ( D , R ) ∈ R 2 defined by

lies above the R ( D ) curve. This again follows from the variational inequality (15) R ≥ H ( π | µ × ( µ × K ) 2 ) = I ( µ × K ) , so that R≥ R ( D ) =: inf π ∈ Π( π, · ): π ( ρ ) ≤D I ( π ) .

Given only a reproduction distribution ν , we will construct a kernel K from ν and use ( µ, K ) to compute ( D , R ) , letting

and

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 L BA (5) for the given ν . Plugging K into the formulas for D and R gives

and

where we introduced the shorthand

Note that we have the following relation (which explains (6))

Let ν be an n -point measure, ν = ∑ n i =1 w i δ y i , e.g., the output of WGD or in the inner step of NERD. Then ϕ ( x ) can be evaluated exactly as a finite sum over the ν particles, and the expressions above for D and R (which are integrals w.r.t. the product distribution µ ⊗ ν ) can be estimated as sample averages. That is, given m independent samples { x i } i m =1 from µ , we compute unbiased estimates

where ϕ ( x ) = -log ∑ n j =1 e ϕ ( x i ) -λρ ( x i ,y j ) . Similarly, a sample mean estimate for L BA is given by

In practice, we found it simpler and numerically more stable to instead compute ˆ R as

Whenever possible, we avoid exponentiation and instead use logsumexp to prevent numerical issues.

9 Wasserstein gradient descent

9.1 On Wasserstein gradients of the EOT and rate functionals

First, we elaborate on the Wasserstein gradient of the EOT functional L EOT ( ν ) . That the dual potential from Sinkhorn's algorithm is differentiable follows from the fact that optimal dual potentials satisfy the Schrödinger equations (cf. [Nutz, 2021, Corollary 2.5]). Differentiability was shown in [Genevay et al., 2019, Theorem 2] in the compact setting, and in [Mena and Niles-Weed, 2019, Proposition 1] in unbounded settings. While Mena and Niles-Weed [2019] only states the result for quadratic cost, the approach of Proposition 1 therein applies more generally.

Below, we compute the Wasserstein gradient of the rate functional L BA ( ν ) . Recall from (6),

Under sufficient integrability on µ and ν to exchange the order of limit and integral, we can calculate the first variation as

where the last equality uses lim ε → 0 1 ε log(1 + εx ) = x and Fubini's theorem. Thus the first variation ψ ν of L BA at ν is

To find the desired Wasserstein gradient of L BA , it remains to take the Euclidean gradient of ψ ν , i.e., ∇L BA ( ν ) = ∇ ψ ν . Doing so gives us the desired Wasserstein gradient:

again assuming suitable regularity conditions on ρ and µ to exchange the order of integral and differentiation.

9.2 Proof of Proposition 4.2 (convergence of Wasserstein gradient descent)

We first provide an auxiliary result.

Proof. The conclusion remains unchanged when rescaling a t by the constant C , and thus without loss of generality C = 1 .

Clearly γ t → 0 as ∑ ∞ t =1 γ 2 t < ∞ . Moreover, there exists a subsequence of ( a t ) t ∈ N which converges to zero (otherwise there exists δ > 0 such that a t ≥ δ > 0 for all but finitely many t , contradicting ∑ ∞ t =1 γ t a t < ∞ ).

Arguing by contradiction, suppose that the conclusion fails, i.e., that there exists a subsequence of ( a t ) t ∈ N which is uniformly bounded away from zero, say a t ≥ δ > 0 along that subsequence. Using this subsequence and the convergent subsequence mentioned above, we can construct a subsequence a i 1 , a i 2 , a i 3 , . . . where a i n ≈ 0 for n odd and a i n ≥ δ for n even. We will show that

contradicting the finiteness of ∑ t γ t a t . (The notation ≈ ( glyph[greaterorsimilar] ) indicates (in)equality up to additive terms converging to zero for n →∞ .)

To ease notation, fix n and set m = i 2 n -1 and M = i 2 n . We show that ∑ M t = m a t γ t glyph[greaterorsimilar] δ 2 / 2 . To this end, using | a t -a t +1 | ≤ γ t we find

Since a m ≈ 0 , there exists a largest n 0 ∈ N , n 0 ≥ m , such that ∑ M -1 j = n 0 γ j glyph[greaterorsimilar] δ (and thus ∑ M -1 j = n 0 γ j glyph[lessorsimilar] δ -γ n 0 ≈ δ as well). We conclude

where we used that ∑ M t = n 0 γ 2 t ≈ 0 . This completes the proof.

Proof of Proposition 4.2. Using the linear approximation property in (14), we calculate

As L ( ν (0) ) is finite and L ( ν ( n ) ) is bounded from below, it follows that

The claim now follow by applying Lemma 9.1 with a t = ∫ ‖∇ ψ ν ( t ) ‖ 2 dν ( t ) ; note that the assumption in the lemma, | a t -a t +1 | ≤ Cγ t , is satisfied due to the second inequality in (14) and a short calculation below

where we use sup t ∫ ‖∇ V L ( ν ( t ) ) ‖ 2 ν ( t ) ( dx ) < ∞ in the last step.

9.3 Proof of Proposition 4.3 (sample complexity)

Recall that X = Y = R d and ρ ( x, y ) = ‖ x -y ‖ 2 in this proposition. For the proof, we will need the following lemma which is of independent interest. We write ν ≤ c µ if ν is dominated by µ in convex order, i.e., ∫ f dν ≤ ∫ f dµ for all convex functions f : R d → R .

Lemma 9.2. Let µ have finite second moment. Given ν ∈ P ( R d ) , there exists ˜ ν ∈ P ( R d ) with ˜ ν ≤ c µ and

This inequality is strict if ν glyph[negationslash]≤ c µ . In particular, any optimizer ν ∗ of (8) satisfies ν ∗ ≤ c µ .

Proof. Because this proof uses disintegration over Y , it is convenient to reverse the order of the spaces in the notation and write a generic point as ( x, y ) ∈ Y × X . Consider π ∈ Π( ν, µ ) and its disintegration π = ν ( dx ) ⊗ K ( x, dy ) over x ∈ Y . Define T : R d → R d by

Define also ˜ π := ( T, id ) # π and ˜ ν := ˜ π 1 . From the definition of T , we see that ˜ π is a martingale, thus ˜ ν ≤ c µ . Moreover, ˜ ν ⊗ µ = ( T, id ) # ν ⊗ µ . The data-processing inequality now shows that

On the other hand, ∫ ‖ ∫ ˜ y K ( x, d ˜ y ) -y ‖ 2 K ( x, dy ) ≤ ∫ ‖ x -y ‖ 2 K ( x, dy ) since the barycenter minimizes the squared distance, and this inequality is strict whenever x = ∫ ˜ yK ( x, d ˜ y ) . Thus glyph[negationslash]

and the inequality is strict unless T ( x ) = x for ν -a.e. x , which in turn is equivalent to π being a martingale. The claims follow.

Proof of Proposition 4.3. Subgaussianity of the optimizer follows directly from Lemma 9.2.

Recalling that inf ν L EOT ( ν ) and inf ν λ -1 L BA ( ν ) have the same values and minimizers (given by (9) in Sec. 2.2), it suffices to show the claim for L = L EOT . Let ν ∗ be an optimizer of (8) (i.e., an optimal reproduction distribution) and ν n its empirical measure from n samples, then clearly

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 ν ∗ m the optimizer for the problem (8) with µ replaced by µ m . Similarly to the above, we obtain

where the expectation is taken over samples from µ m . In this situation, we cannot directly apply [Mena and Niles-Weed, 2019, Theorem 2]. However, the bound given by [Mena and Niles-Weed, 2019, Proposition 2] still applies, and the only dependence on ν ∈ { ν ∗ , ν ∗ m } is through their subgaussianity constants. By Lemma 9.2, these constants are bounded by the corresponding constants of µ and µ m . Thus, the arguments in the proof of [Mena and Niles-Weed, 2019, Theorem 2] can be applied, yielding the second inequality of Proposition 4.3.

The final inequality of Proposition 4.3 follows from the first two inequalities (the first one being applied with µ m ) and the triangle inequality, where we again use the arguments in the proof of [Mena and Niles-Weed, 2019, Theorem 2] to bound the expectation over the subgaussianity constants of µ m .

9.4 Convergence of the proposed hybrid algorithm

In our present work, our hybrid algorithm targets the non-stochastic setting and is motivated by the fact that the BA update is in some sense orthogonal to the Wasserstein gradient update, and can only monotonically improve the objective. While empirically we observe the additional BA steps to not hurt - but rather accelerate - the convergence of WGD (see 5.1), additional effort is required to theoretically guarantee convergence of the hybrid algorithm.

There are two key properties we use for the convergence of the base WGD algorithm: 1) a certain monotonicity of the update steps (up to higher order terms, gradient descent improves the objective) and 2) stability of gradients across iterations. If we include the BA step, we find that 1) still holds, but 2) may a-priori be lost. Indeed, 1) holds since BA updates monotonically improve the objective. Using just 1), we can still obtain a Pareto convergence of the gradients for the hybrid algorithm, ∑ ∞ t =0 γ t ∫ ‖∇ V L ( ν ( t ) ) ‖ 2 dν ( t ) < ∞ (here ν ( t ) are the outputs from the respective BA

steps and γ t is the step size of the gradient steps). Without property 2), we cannot conclude ∫ ‖∇ V L ( ν ( t ) ) ‖ 2 dν ( t ) → 0 for t →∞ . We emphasize that in practice, it still appears that 2) holds even after including the BA step. Motivated by this analysis, an adjusted hybrid algorithm where, e.g., the BA update is rejected if it causes a drastic change in the Wasserstein gradient, could guarantee that 2) holds with little practical changes. From a different perspective, we also believe the hybrid algorithm may be tractable to study in relation to gradient flows in the Wasserstein-Fisher-Rao geometry (cf. [Yan et al., 2023]), in which the BA step corresponds to a gradient update in the Fisher-Rao geometry with a unit step size.

In the stochastic setting, the BA (and therefore our hybrid) algorithm does not directly apply, as performing BA updates on mini-batches no longer guarantees monotonic improvement of the overall objective. Extending the BA and hybrid algorithm to this setting would be interesting future work.

10 R-D estimation and variational inference/learning

In this section, we elaborate on the connection between the R-D problem (3) and variational inference and learning in latent variable models, of which maximum likelihood deconvolution (discussed in Sec. 2.3) can be seen as a special case. Also see Section 3 of [Yang et al., 2023] for a related discussion in the context of lossy data compression.

To make the connections clearer to a general machine learning audience, we adopt notation more common in statistics and information theory. Table 1 summarizes the notation and the correspondence to the measure-theoretic notation used in the main text.

In statistics, a common goal is to model an (unknown) data generating distribution P X by some density model ˆ p ( x ) . In particular, here we will choose ˆ p ( x ) to be a latent variable model, where Y takes on the role of a latent space, and Q Y = ν is the distribution of a latent variable Y (which may encapsulate the model parameters). As we shall see, the optimization objective defining the rate functional (5) corresponds to an aggregate Evidence LOwer Bound (ELBO) [Blei et al., 2017]. Thus, computing the rate functional corresponds to variational inference [Blei et al., 2017] in a given model (see Sec. 10.2), and the parametric R-D estimation problem, i.e.,

is equivalent to estimating a latent variable model using the variational EM algorithm [Beal and Ghahramani, 2003] (see Sec. 10.3). The variational EM algorithm can be seen as a restricted version of the BA algorithm (see Sec. 10.3), whereas the EM algorithm [Dempster et al., 1977] shares its E-step with the BA algorithm but can differ in its M-step (see Sec. 10.4).

10.1 Setup

For concreteness, fix a reference measure ζ on Y , and suppose Q Y has density q ( y ) w.r.t. ζ . Often the latent space Y is a Euclidean space, and q ( y ) 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 q ( y ) is a probability mass function. We will consider the typical parametric estimation problem and choose a particular parametric form for Q Y indexed by a parameter vector θ . This defines a parametric family H = { Q θ Y : θ ∈ Θ } for some parameter space Θ . Finally, suppose the distortion

function ρ induces a conditional likelihood density via p ( x | y ) ∝ e -λρ ( x,y ) , with a normalization constant that is constant in y .

These choices then result in a latent variable model specified by the joint density q ( y ) p ( x | y ) , and we model the data distribution with the marginal density: 2

As an example, a Gaussian mixture model with isotropic component variances can be specified as follows. Let Q Y be a mixing distribution on X = Y = R d parameterized by component weights w 1 ,...,k and locations µ 1 ,...,k , such that Q Y = ∑ K k =1 w k δ µ k . Let p ( x | y ) = N ( y, σ 2 ) be a conditional Gaussian density with mean y and variance σ 2 . Now formula (22) gives the usual Gaussian mixture density on R d .

Maximum-likelihood estimation then ideally maximizes the population log (marginal) likelihood,

The maximum-likelihood deconvolution setup can be seen as a special case where the form of the marginal density (22) derives from knowledge of the true data generation process, with P X = α ∗ N (0 , 1 λ ) for some unknown α and known noise N (0 , 1 λ ) (i.e., the model is well-specified). We note in passing that the idea of estimating an unknown data distribution by adding artificial noise to it also underlies work on spread divergence [Zhang et al., 2020] and denoising score matching [Vincent, 2011].

To deal with the often intractable marginal likelihood in the inner integral, we turn to variational inference and learning [Jordan et al., 1999, Wainwright et al., 2008].

10.2 Connection to variational inference

Given a latent variable model and any data observation x , a central task in Bayesian statistics is to infer the Bayesian posterior [Jordan, 1999], which we formally view as a conditional distribution Q ∗ Y | X = x . It is given by

or, in terms of the density q ( y ) of Q Y , given by the following conditional density via the familiar Bayes' rule,

Unfortunately, the true Bayesian posterior is typically intractable, as the (marginal) data likelihood in the denominator involves an often high-dimensional integral. Variational inference [Jordan et al., 1999, Wainwright et al., 2008] therefore aims to approximate the true posterior by a variational distribution Q Y | X = x ∈ P ( Y ) by minimizing their relative divergence H ( Q Y | X = x | Q ∗ Y | X = x ) . The problem is equivalent to maximizing the following lower bound on the marginal log-likelihood, known as the Evidence Lower BOund (ELBO) [Blei et al., 2017]:

Recall the rate functional (5) arises through an optimization problem over a transition kernel K ,

Translating the above into the present notation ( µ → P X , K → Q Y | X , ν → Q Y ; see Table 1), we obtain

which allows us to interpret the rate functional as the result of performing variational inference through ELBO optimization. At optimality, Q Y | X = Q ∗ Y | X , the ELBO (24) is tight and recovers log ˆ p ( x ) , and the rate functional takes on the form of a (negated) population marginal log likelihood (23), as given earlier by (6) in Sec. 2.1.

10.3 Connection to variational EM

The discussion so far concerns probabilistic inference , where a latent variable model ( Q Y , p ( x | y )) has been given and we saw that computing the rate functional amounts to variational inference. Suppose now we wish to learn a model from data. The R-D problem (4) then corresponds to model estimation using the variational EM algorithm [Beal and Ghahramani, 2003].

To estimate a latent variable model by (approximate) maximum-likelihood, the variational EM algorithm maximizes the population ELBO

w.r.t. Q Y and Q Y | X . This precisely corresponds to the R-D problem inf Q Y ∈H L BA ( Q Y ) , using the form of L BA ( Q Y ) from (25).

In popular implementations of variational EM such as the V AE [Kingma and Welling, 2013], Q Y and Q Y | X are restricted to parametric families. When they are allowed to range over all of P ( Y ) and all conditional distributions, variational EM then becomes equivalent to the BA algorithm.

10.4 The Blahut-Arimoto and (exact) EM algorithms

Here we focus on the distinction between the BA algorithm and the (exact) EM algorithm [Dempster et al., 1977], rather than the variational EM algorithm. Both BA and (exact) EM perform coordinate descent on the same objective function (namely the negative of the population ELBO from (26)), but they define the coordinates slightly differently - the BA algorithm uses ( Q Y | X , Q Y ) with Q Y ∈ P ( Y ) , whereas the EM algorithm uses ( Q Y | X , θ ) with θ indexing a parametric family H = { Q θ Y : θ ∈ Θ } . Thus the coordinate update w.r.t. Q Y | X (the 'E-step') is the same in both algorithms, but the subseuquent 'M-step' potentially differs depending on the role of θ .

Given the optimization objective, which is simply the negative of (26),

both BA and EM optimize the transition kernel Q Y | X the same way in the E-step, as

For the M-step, the BA algorithm only minimizes the relative entropy term of the objective (27),

(with the optimal Q Y given by the second marginal of P X Q ∗ Y | X ) whereas the EM algorithm minimizes the full objective w.r.t. the parameters θ of Q Y ,

The difference comes from the fact that when we parameterize Q Y by θ in the parameter estimation problem, Q ∗ Y | X -and consequently both terms in the objective of (29) - will have functional dependence on θ through the E-step optimality condition (28).

In the Gaussian mixture example, Q Y = ∑ K k =1 w k δ µ k , and its parameters θ consist of the components weights ( w 1 , ..., w K ) ∈ ∆ d -1 and location vectors { µ 1 , ..., µ K } ⊂ R d . The E-step computes Q ∗ Y | X = x = ∑ k w k p ( x | µ k ) p ( x ) δ µ k . For the M-step, if we regard the locations as known so that θ = ( w 1 , ..., w K ) only consists of the weights, then the two algorithms perform the same update; however if θ 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 E ( x,y ) ∼ P X Q ∗ Y | X [ -log p ( x | y )] =

11 Further experimental results

Our deconvolution experiments were run on Intel(R) Xeon(R) CPUs, and the rest of the experiments were run on Titan RTX GPUs.

In most experiments, we use the Adam [Kingma and Ba, 2015] optimizer for updating the ν 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 { 1 e 4 , 3 e 3 , 1 e 3 , 3 e 2 , 1 e 2 , ... } . Alternatively, a constrained optimization method like the Modified Differential Multiplier Method (MDMM) [Platt and Barr, 1987] can be adopted to directly target R ( D ) for various values of D s. The latter approach will also help us identify when we run into the log( n ) rate limit of particle-based methods (Sec. 4.4): suppose we perform constrained optimization with a distortion threshold of D ; if the chosen n is not large enough, i.e., log( n ) < R ( D ) , then no ν supported on (at most) n points is feasible for the given constraint (otherwise we have a contradiction). When this happens, the Lagrange multiplier associated with the infeasiable constraint ( λ 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 a similar architecture as the one used on the banana-shaped source in [Yang and Mandt, 2022], consisting of two-layer MLPs for the encoder and decoder networks, and Masked Autoregressive Flow [Papamakarios et al., 2017] for the variational prior. For NERD, we follow similar architecture settings as [Lei et al., 2023a], using a two-layer MLP for the decoder network. Following Yang and Mandt [2022], we initially used 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 ν ∗ = α measure which is concentrated on the unit circle and does not have a Lebesgue density.

Experiment setup As discussed in Sec. 4.2, BA and our hybrid algorithms do not apply to the stochastic setting; to be able to include them in our comparison, the input to all the algorithms is an empirical measure µ m (training distribution) with m = 100000 samples, given the same fixed seed. We found the number of training samples is sufficiently large that the losses do not differ significantly on the training distribution v.s. random/held-out samples.

Recall from Sec. 2.3, given data observations following µ = α ∗ N (0 , σ 2 ) , if we perform deconvolution with an assumed noise variance 1 λ for some 1 λ ≤ σ 2 , then the optimal solution to the problem is given by ν ∗ = α λ = α ∗N (0 , 1 λ ) . We compute the optimal loss OPT = L BA ( ν ∗ ) by a Monte-Carlo estimate, using formula (19) with m = 10 4 samples from µ and n = 10 6 samples from ν ∗ . The resulting ˆ L BA is positively biased (it overestimates the OPT in expectation) due to the bias of the plug-in estimator for ϕ ( x ) (18), so we use a large n to reduce this bias. 3 Note the same issue occurs in the NERD algorithm (13).

-

1

0

1

-

1

0

1

λ

= 100

.

0

-

1

0

1

-

1

0

1

λ

= 10

.

0

-

1

0

1

-

1

0

1

λ

= 3

.

0

-

1

0

1

-

1

0

1

λ

= 1

.

0

1

0

1

-

1

0

1

λ

= 100

.

0

1

0

1

-

1

0

1

λ

= 10

.

0

1

0

1

-

1

0

1

λ

= 3

.

0

1

0

1

-

1

0

1

λ

= 1

.

0

-

-

-

-

Optimized reproduction distributions To visualize the (continuous) ν learned by RD-VAE and NERD, we fit a kernel density estimator on 10000 ν samples drawn from each, using seaborn.kdeplot .

In Figure 6, we provide additional visualization for the particles estimated from the training samples by WGD and its hybrid variant, for a decreasing distortion penalty λ (equivalently, increasing entropic regularization strength glyph[epsilon1] ).

11.2 Higher-dimensional datasets

For NERD, we set n to the default 40000 in the code provided by [Lei et al., 2023a], on all three datasets.

For WGD, we used n = 4000 on physics , 200000 on speech , and 40000 on MNIST (see also smaller n for comparison in Fig. 5).

On speech , both NERD and WGD encountered the issue of a log( n ) upper bound on the rate estimate as described in Sec. 4.4. Therefore, we increased n to 200000 for WGD and obtain a tighter R-D upper bound than NERD, particularly in the low-distortion regime. Such a large n incurred out-of-memory error for NERD.

We borrow the R-D upper and lower bound results of [Yang and Mandt, 2022] from their paper on physics and speech . We obtain sandwich bounds on MNIST using a similar neural network architecture and other hyperparameter choices as in the GAN-generated image experiments of [Yang and Mandt, 2022] (using a ResNet-VAE for the upper bound and a CNN for the lower bound), except we set a larger k = 10000 in the lower bound experiment; we suspect the resulting lower bound is still far from being tight.

12 Example implementation of WGD

We provide a self-contained minimal implementation of Wasserstein gradient descent on L BA , using the Jax library [Bradbury et al., 2018]. To compute the Wasserstein gradient, we evaluate the first variation of the rate functional in compute_psi_sum according to (20), yielding ∑ n i =1 ψ ν ( x i ) , then simply take its gradient w.r.t. the particle locations x 1 ,...n using Jax's autodiff tool on line 51.

The implementation of WGD on L EOT is similar, except the first variation is computed using Sinkhorn's algorithm. Both versions can be easily just-in-time compiled and enjoy GPU acceleration.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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\nSTRENGTHS:\nThis 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.\n\nWEAKNESSES:\nThere 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$.\n\nQUESTIONS:\nWould 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\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-25"}

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

STRENGTHS:
- The method is novel and well motivated
- Good choice of numerical experiments

WEAKNESSES:
Lack of theoretical analysis/guarantees  (see also questions)

QUESTIONS:
Some questions on the theory side:
- In (15) in line 202, should there be an inequality between the first and second line?
- Should the integral in Lemma 4.2 not be w.r.t \nu^t and not \nu^k ?
- 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.

Small Typos:
- Line 172: baised should be biased
- Line 221: optimie should be optimize
- Line 302: Should "articles" be particles?
- 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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-06-26

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe proposed method seems novel, and theoretical contribution about sample complexity (Prop. 4.3) seems important.\n\nWEAKNESSES:\nMy 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\nQUESTIONS:\nThe 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.\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-01"}

comments: 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.

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.

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

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

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?

2. 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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-07

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\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 (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.\n\nQUESTIONS:\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\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-08"}