Estimating the Rate-Distortion Function by Wasserstein Gradient Descent

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

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. @article{rigollet2022sample, title={On the sample complexity of entropic optimal transport}, author={Rigollet, Philippe and Stromme, Austin J}, journal={arXiv preprint arXiv:2206.13472}, year={2022} } It is also worth pointing out that this constant "$C$" in (15) of the text scales exponentially poorly with the regularization parameter $\epsilon$.

Questions

Would it be interesting to analyze the "debiased" entropic OT dynamics for this problem as in e.g. [Pooladian2022Debiaser] and [Feydy2019Interpolating]? @inproceedings{pooladian2022debiaser, title={Debiaser beware: Pitfalls of centering regularized transport maps}, author={Pooladian, Aram-Alexandre and Cuturi, Marco and Niles-Weed, Jonathan}, booktitle={International Conference on Machine Learning}, pages={17830--17847}, year={2022}, organization={PMLR} } @inproceedings{feydy2019interpolating, title={Interpolating between optimal transport and mmd using sinkhorn divergences}, 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}, booktitle={The 22nd International Conference on Artificial Intelligence and Statistics}, pages={2681--2690}, year={2019}, organization={PMLR} } Minor comments: L33: “R-D” function instead of $R(D)$ (L86, L99, L101, L110, …) [just to be consistent] L221: “optimize” Missing \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).

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Yes

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)

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

It would be interesting to have more theoretical guarantees for this method and/or more experiments on other realistic data.

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.

Strengths

The proposed method seems novel, and theoretical contribution about sample complexity (Prop. 4.3) seems important.

Weaknesses

My concerns here are twofold: 1. 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. 2. Second, 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.

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.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

1 poor

Contribution

3 good

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.

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?

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

The authors have adequately addressed the limitations and potential negative societal impact.

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.

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

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?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

4 excellent

Presentation

3 good

Contribution

2 fair

Limitations

Yes

Follow-up questions

Dear Authors, Thank you very much for your response and clarifications! As AC, I would like to ask follow-up questions that allow me to better understand your contributions. The proposed algorithm (WGD) is similar to particle gradient descent introduced by [Nitanda, Atsushi, and Taiji Suzuki] which basically is gradient descent on the points in the support. In the specific regime of stepsizes $\gamma_k \to 0$, I think the two algorithms are equivalent. If this holds, then the convergence result in Lemma 4.2 can be recovered by invoking the convergence of gradient descent on smooth functions. Would you mind elaborating on this connection? Best regards, AC

Response

I would like to thank the authors for their detailed response. I will keep my rating and still recommend accept.