Entropic Neural Optimal Transport via Diffusion Processes

Summary

The paper proposes to solve dynamic entropic optimal tansport (EOT), also known as Schrödinger bridge problem, with nerual solver. Specifically, the authors propose a saddle-point, maximin, formulation of EOT, yielding a GAN-resemble algorithm that can be trained in an end-to-end fashion. Experiments are conducted on 2D toy datasets and images translation.

Strengths

- The saddle point reformulation of EOT is interesting. The proposed algorithm can be viewed as a stochastic control problem, where the terminal cost function $\beta$ is learned so that the resulting policy approaches terminal distribution. - Writing is generally clear and easy to follow. Sufficient related works and preliminaries are included. - Experiment are extensive and include many related baselines.

Weaknesses

- The proposed algorithm aims to minimize (12), which is not well-defined when $\epsilon$=0, as the KL term will blow up. Even though $\epsilon$=0 is algorithmically applicable, it makes the current algorithm disconnected from the mathematical framework. I'll be more convinced if the authors can provide additional justification (maybe connection to OT). - The proposed method is closely related to recent maximim OT [1], which consists of the same two networks (potential + policy) and the same training losses. From my understanding, the two methods coincide when $\epsilon$=0 and $N=1$. Are the authors aware of [1]? Given that the proposed ENOT seems to work best when $\epsilon$=0, can the author compare with [1]? - Given that DiffSB [13] was compared throughout most Sec 5, I suggest the authors to compare in Sec 5.4 to [2], which applies DiffSB to unpaired super-resolution image datasets. - While Sec 2 has introduced sufficient background and comparison between EOT and SB, which I do appreciate, I think their connection to Sec 4, which is the main content, is rather weak. Given that the proposed algorithm closely relates to dual formulation (e.g. (26) in Appendix), I suggest including those parts in Sec 2. [1] On amortizing convex conjugates for optimal transport [2] Conditional Simulation Using Diffusion Schrödinger Bridges

Questions

- Experiments seem to support that $\epsilon$=0, where the dynamics reduce to flow, works much better than SDE. This seems to contradict recent insights from diffusion models [1], where SDE usually performs better than ODE on the same tasks. Can the authors comment on that? - why is there a "1/|f|" term in Alg 1, when computing the KL? And what's $f_{n,m}$? I understand tht $f_n$ is the drift output at time step $n$. [1] Elucidating the Design Space of Diffusion-Based Generative Models

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

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

2 fair

Contribution

2 fair

Limitations

Limitations are included in Sec6, where the authors mentioned the potential computational burden caused by simulation and back-prop through the SDE dynamics during training.

Summary

Inspired by how Sinkhorn duals are derived the authors adapt said derivation to the path measure and via the disintegration theorem they derive a novel unconstrained min-max objective for solving the Schrodinger bridge problem, the authors then proceed to showcase their method in eOT based tasks, introducing a new gaussian benchmark and displaying competitive results to previous approaches. Additionally, the authors quantify errors in sampling and transport of approximate minimisers to their proposed schemes.

Strengths

1. The extension of the Sinkhorn dual to the dynamic setting is rather elegant and certainly novel in the way it is carried out 2. The paper is excellent on the presentation side in regards to technical ideas, whilst the contributions are novel/creative they are presented in such a way that understanding them was not difficult. 3. The new formulation allows for a novel duality gap analysis, which is one of the few works analysing learned SBP methods in the approximate setting. 4. From a purely conceptual viewpoint the work is great and rather complete, just some clarifications/additions on the experimental side and motivations could be enhanced.

Weaknesses

Outside of a concern in how different methods are compared (detailed in the questions), this paper is overall well written and has mostly sound experiments. From the method standpoint, there are several potential weaknesses and lacking ablations which I will detail in the limitations.

Questions

1. Some of the methods you compare to are closed form / non-iterative in the way they solve the half bridges e.g. [1] using GPs. A thing to note is iterative approaches such as SGD can keep resampling from P_0 and P_1 in the toy tasks nonstop, effectively making the dataset set size somewhat proportional to the iterations, this strongly benefits generalisation capabilities, especially in higher dimensions in contrast to a the GP approach in [1] which is likely fitted on a small dataset. If we assume that the dataset is fixed across epochs K_f * batchsize would give a proxy as to the size of the dataset used in the DL approaches , from inspecting the code provided in the supplementary zip in particular the high dim Gaussians and toy examples IPYNB this quantity seems to be above 5000, whilst I suspect for [1] it might be lower. Even then I suspect you resample from the toy distributions (Gaussians) every time as indicated in your pseudocode in which case the dataset sizes are simply not comparable at all and the non-iterative approaches like MLE-SB are prone to be affected significantly by generalisation error (which scales slowly in number of samples for kernel methods, which also don't overcome the curse of dim as well as neural networks) in particular in high dimensions. 2. Are the same number of timesteps used across the DL and the GP approaches when training ? computationally it feels unlikely the gram matrix would fit in memory for the given timesteps used with the DL approaches. 3. A suggestion here would be to compare all approaches in an additional table (e.g. in the appendix) under budgets in the number of steps and number of samples (and fairly ensure these are the same across methods), this can be useful to quantify the data efficiency of each approach. Furthermore, the settings under which [1] was run (steps and dataset size) should be reported.

Rating

9: Very Strong Accept: Technically flawless paper with groundbreaking impact on at least one area of AI/ML and excellent impact on multiple areas of AI/ML, with flawless evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

Finding saddlepoints / solving min-max problems is typical quite a challenging task and was often a huge challenge in the training stability of GANs (very high variance training results). In contrast to IPF which is much nicer from an objective viewpoint (sequence of regression losses), this approach poses the question of how stable training is. As none of the experiments have error bars (on training runs) it is difficult to see if the proposed approach is robust and “easy” to train. I would suggest to the authors provide such results and in addition maybe comparisons to training loss with Chen 2021 joint or De Bortoli 2021.

Summary

This paper focuses on neural optimal transport and more particularly through a "dynamic schrodinger bridge" approach. I am not an expert in this particular topic, but I must say that the authors manage to make it quite readable and a good introduction to the methodology. As far as I can tell, the particularity of the approach is to start from a known key constrained optimization problem in eq (11) that has been presented before and that is progressively explained. Then, departing from previous methods, the authors propose a Lagrangian approach, for which we now have an unconstrained problem and the need to train not only the drift function that allows sampling, but also some "beta network" (coined in as the beta network), in a min-max optimization approach, that is similar --but different -- from the usual GAN methodology. After the rigorous theoretical treatment, the authors present some very nice experiment that seem to support their claims and the interesting features of their approach. I think the paper is quite challenging, but that it is also very stimulating and inspiring. -- after rebuttal, I am still happy with the paper. Maintaining my score.

Strengths

I feel slightly uncomfortable in assessing whether the proposed method is new or is not exactly or if the authors miss some particular recent and/or relevant method in their references, mostly because I cannot be considered an expert on the topic. Still, it looks to me that they are doing their best at providing a very objective account of the relevant literature and giving many pointers, so I am assuming they can be trusted when they present their contributions. The proposed method seems quite ok to implement (cf Algorithm 1) and the results are very nice. I am not really aware of the appropriate metrics and usual evaluation criteria that should be used, but I felt compelled by the experiments, which made me want to try things out. I guess this is the most important aspect. To summarize what I feel are the strengths of the paper: - good theoretical overview, motivations and derivations - the resulting method seems simple to implement - interesting performances

Weaknesses

The paper is a bit difficult to follow, but I don't think it should really be modified to be simpler. I guess it's mostly a matter of myself not being an expert in the topic. Still, some changes here or clarifications here and there could help. I will mention them below.

Questions

- p3 "Hence, one may optimize (5) over processes T for which T|x,y=W_xy for every x, y and set the last term in (6) to zero". How do you actually do that ? By enforcing gaussian dynamics ? - the paragraph "continuous OT" reads a bit weak to me. All the references are just given in a row, without great care for actually discussing them and their connections with the proposed approach. In algorithm 1, I have some questions. - do you actually need to store the gradients when computing {X_n, f_n} for the computation of L_\beta ? It looks to me that Eul-Mar(X_0, T_{f_\theta}}) can actually run in some `no_grad` environment, or am I mistaken ? This would mean there is some possibility in strongly parallelizing this or even use some external workers to compute that ? - On the contrary, in the inner loop (over k), you really need to store the gradients, right ? Is it necessary to use many inner iterations K_f ? Maybe I'm mistaken but I don't see that discussed, although this looks like a key computational burden when I have a look at algorithm 1, right ? Is it feasible to just record gradient for the last k steps ? for some of them ? Some thing that could help going faster ? - I have no clue what this BW2-UVP metric is. Could you at least give us a hint that would help us avoid checking reference [25] ? - You must check the references. Most of them are badly formated. You have "schr\"odinger" everywhere and you are missing many uppercases for proper nouns

Rating

9: Very Strong Accept: Technically flawless paper with groundbreaking impact on at least one area of AI/ML and excellent impact on multiple areas of AI/ML, with flawless evaluation, resources, and 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

4 excellent

Contribution

3 good

Limitations

I would say that the limitations are clearly mentioned

Summary

The main idea of the paper is to estimate a stochastic map for the entropic optimal transport problem using its connection to the dynamic Schrödinger bridge (SDB) problem. The authors formulate the SDB as a saddle point problem of an associated Lagrangian. Then they recover the transport plan as the joint distribution of the solution to the dynamic Schrödinger bridge problem's initial and final values, while the transport is encoded in the drift term of the learned stochastic process that is the solution of the SDB. Compared to previous methods, the method at hand offers more stability for small entropic regularization coefficients. The errors on the drift term solution and on the transport map are quantified given the corresponding duality gaps of the inner and outer optimization problems of the saddle point objective. Finally, the approach is supported by experimental evaluation.

Strengths

The paper is extremely well-written and reader-friendly. - Several remarks are made to facilitate reading and to provide intuition about technical notions. - A guarantee on the quality of the saddle point solution is provided. - Addressing the small $\epsilon$ case, which is a source of instability of several other methods.

Weaknesses

* In line 223, it is mentioned that the negative entropy is not strongly convex. This is false as the function $p\mapsto x\ln(x)$ has second derivate $x \mapsto \frac{1}{x}$ which is bounded from below by $1$ on the interval $[0,1]$. See for example Section 4.1 of reference [2]. As a result, the comparison to [1] (reference [5] in the paper) needs to be reconsidered. * There is no concluding section. * Some experimental section metrics are not introduced. * The $\rm{BW_2^2-UVP}$ is not introduced even in the appendix, although a reference for it is given. * FID: the previous remark applies. * The parametrization $g(X_t,t) = X_t + f(X_t,t)\Delta_t $ should have been indicated in the main paper rather than in the appendix as although it is mathematically equivalent to the parametrization presented in the main paper, it allowed better results on the CelebA dataset according to the appendix ## Minor remarks * Problem with links: For some reason, the bibliographic references links along with links to equations and sections etc. are not working. * $\pi^{W^\epsilon}$ is introduced for the first time in Equation (8) without being defined. * $W_{|x,y}$ is not explicitly defined, although one can infer its meaning from the definition of $T_{|x,y}$. * Theorem 4.1: I think it should be "every pair $(\beta^*,T_{f^*})$ for (13)" rather than "for (12)" as problem (13) is a saddle point problem. * Reference to Algorithm 2: in line 195, Algorithm 2 is referenced. However, it is not indicated that it is written in the appendix. * Suggestion: index $m$ can be removed in the "$\widehat{KL} \leftarrow$" line of Algorithm 1 since the sum terms are already indicated to be the values of $f_n$, or it is possible to indicate $\sum_{m=1}^{|f_n|}$. * Line 158: I think it should be added "that is bounded from above" to "a continuous function". ## References [1] Asadulaev, A., Korotin, A., Egiazarian, V., & Burnaev, E. (2022). Neural optimal transport with general cost functionals. *arXiv preprint arXiv:2205.15403*. [2] Peyré, G., & Cuturi, M. (2019). Computational optimal transport: With applications to data science. *Foundations and Trends® in Machine Learning*, *11*(5-6), 355-607.

Questions

* Computation of $\mathbb{E}\left[\int_0^1\Vert f(X_t,t)\Vert^2{\rm d}t\right]$ : in line 197, it is indicated that the mean of the $f(x,t)$ is used. Is this justified by the Riemann integral discrete approximation? If so, does a trapezoidal rule for approximating the integral improve the result? * Is it straightforward to generalize the approach to costs other than the squared Euclidean distance? Does it fundamentally change the nature of the associated stochastic process * Did the authors try to apply the method to the domain adaptation (DA) problem as several DA methods rely on optimal transport ? ---------- I have read the authors's rebuttal. They have addressed my concerns.

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

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

4 excellent

Presentation

4 excellent

Contribution

4 excellent

Limitations

The impact of the paper and the limitations of the contribution are clearly discussed in Section 6.

Very satisfied with rebuttal and the uploaded comparisons.

Dear Authors, I have read the rebuttal and can comment that you have clarified and resolved all the issues I raised. In particular, the stability plots and the further fair comparisons brought everything to a very clean conclusion. Additionally, I have also gone over the other reviews and the provided responses, overall it looks pretty satisfactory, the non-trivial prior results were quite a nice addition and comment to the robustness/flexibility of the proposed approach. I will increase my score accordingly once the platform allows such, at the time being it seems OpenReview is not allowing this until the discussion period is over, I do agree this work does seem to be the most thorough comparison yet for SBP solvers within the eOT context making this paper a substantial contribution beyond the novel approach that is proposed (which is also a solid and self-contained contribution on its own).

Thanks for the rebuttal

Dear Authors, I would like to thank you for the detailed and very well written rebuttal. It addressed all of my concerns. I update my score.