Scalable Simulation-free Entropic Unbalanced Optimal Transport

Authors: Jaemoo Choi, Jaewoong Choi
Primary area: generative models
Keywords: Optimal Transport, Unbalanced Optimal Transport, Entropic Optimal Transport, Generative models, Image-to-image Translation

Abstract:
The Optimal Transport (OT) problem investigates a transport map that connects two distributions while minimizing a given cost function. Finding such a transport map has diverse applications in machine learning, such as generative modeling and image-to-image translation. In this paper, we introduce a scalable and simulation-free approach for solving the Entropic Unbalanced Optimal Transport (EUOT) problem. We derive the dynamical form of this EUOT problem, which is a generalization of the Schrödinger bridges (SB) problem. Based on this, we derive dual formulation and optimality conditions of the EUOT problem from the stochastic optimal control interpretation. By leveraging these properties, we propose a simulation-free algorithm to solve EUOT, called Simulation-free EUOT (SF-EUOT). While existing SB models require expensive simulation costs during training and evaluation, our model achieves simulation-free training and one-step generation by utilizing the reciprocal property. Our model demonstrates significantly improved scalability in generative modeling and image-to-image translation tasks compared to previous SB methods.

Decision: reject
Reviewer scores: 6, 5, 6, 5, 6

Reviews and discussion:
<thread>
creation date:
1730637205262

thread_type:
official_review

<thread_object>
creator: reviewer_jnr2

<content>
summary:
This paper proposes a scalable and simulation-free approach for solving the EUOT problem. The authors derive the dynamical form and the duall formulation of the EUOT problem and their method is based on these interpretations. The paper also verifies the effectiveness of the proposed method on both toy datasets and image datasets.

strengths:
The studied problem is interesting and important. The writing is mostly clear. The proposed method has good empirical performance.

weaknesses:
Here are my comments:

1. How to guarantee the convergence of the proposed method?
2. What does NFE mean? The paper does not seem to clarify this term or explain why the proposed method results in a smaller NFE.
3. Since the main advantage of the proposed methods is its scability, I think the paper should provide a detailed comparison on the computational cost rather than only showing the accuracy.

questions:
see weakness
</content>
</thread_object>

</thread>
<thread>
creation date:
1730700042483

thread_type:
official_review

<thread_object>
creator: reviewer_alvw

<content>
summary:
The authors propose a new method for solving the dynamic formulation of the Entropic Unbalanced Optimal Transport (EUOT) problem. The method is based on the Stochastic Optimal Control (SOC) formulation of the dynamic EUOT. The authors divide the optimized variables into two parts: value function of the corresponding SOC problem and the underlying distribution dynamics. Authors train them on functionals, derived from the original SOC formulation and the corresponding Hamilton-Jacobi-Bellman equation. Representation of the problem in terms of this coordinate descent results in the ability to parameterize distribution dynamics of EUOT as a stochastic one-step map instead of the stochastic process, thus allowing to achieve simulation-free (SF) procedure, which authors call SF-EUOT. Authors validate the method on toy 2d examples, unconditional generation and unpaired I2I translation, where they beat OT-based counterparts.

strengths:
Originality: construction of the algorithm based on the coordinate descent with respect to the value function and the distribution dynamics, is novel. One of the optimized functionals is based on the HJB equation, which I have not seen successfully optimized in PINN-like [1] fashion in the OT context.

Quality: the paper has sound theory, well marked connections with other methods and more general OT and SOC theory.

Clarity: despite the amount of derivations, the mathematical writing is accurate and easy to follow, the experimental section and the overall presentation is clear.

Significance: the paper introduces a way to solve dynamic (essentially, multi-step) OT problems, while parameterizing the underlying distribution dynamics with a single-step map. This is especially important in case of Entropic OT problems, where one-step parameterization does not directly allow to estimate entropy, which is why many methods resort to simulation-based training.

[1] Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations

weaknesses:
My main concerns are related to the applicability of the method and limited comparison with the other existing models.

* First, in Table 1 the authors compare their method with models based on EOT, Diffusion and Flow Matching. While I agree with the choice of the EOT solvers, most of the chosen diffusion models correspond to results from 2021-2022. They are now outperformed by a large margin by even one-step diffusion distillation methods (CTM [1] and GDD [2] with FID around 1.6-1.7), which should be indicated for a more fair comparison. Overall, while being more efficient than other EOT solvers, the proposed SF-EUOT is still inefficient compared to more practical methods for unconditional generation. Given this, I would consider this experiment more as a benchmark for EOT solvers rather than a broad comparison with competitive methods.

* In Image-to-Image experiments and in Table 2, in particular, authors choose a limited set of methods for comparison. Two of them, DiscoGAN and CycleGAN, perform much worse than newer (but still 2-4 years old) one-step GAN-based methods such as CUT [3], StarGAN-v2 [4], and ITTR [5], which complicates fair comparison. In the Wild->Cat experiment authors choose to compare only with multi-step methods based on Schrödinger Bridge, while the results from DiscoGAN, CycleGAN and NOT, present in the Male->Female experiment, are missing. Besides, there are other SB-based methods proven competitive in practice, e.g. NSB [6], which is, being a few-step method, relevant for comparison. Overall, this section would greatly benefit from more consistent comparisons (e.g. with the same set of models for both experiments) and from adding higher-quality one-step (or few-step) baselines.

* In the Image-to-Image experiments, only FID is measured as a metric of generation quality. However, as it was done by the authors in Figure 2, it is essential to also measure some similarity metric between the input and the output. It will allow comparison of faithfulness-realism tradeoff between the proposed method and the baselines and will exclude situations e.g. in which the method achieves better FID at the cost of poor transport cost.

* Samples obtained by the SF-EUOT in the Wild->Cat translation problem (Figure 8), seem to significantly suffer from mode collapse.

[1] Consistency Trajectory Models: Learning Probability Flow ODE Trajectory of Diffusion

[2] Diffusion Models Are Innate One-Step Generators

[3] Contrastive Learning for Unpaired Image-to-Image Translation

[4] StarGAN v2: Diverse Image Synthesis for Multiple Domains

[5] ITTR: Unpaired Image-to-Image Translation with Transformers

[6] Unpaired Image-to-Image Translation via Neural Schrödinger Bridge

questions:
* Could the authors clarify more on the technical implementation of the gradient norm and the Laplace operator of the value function and their computation time (compared to the standard forward and backward passes)? While it is mentioned that the Laplace operator is approximated by the Hutchinson trace estimator, here it would require calculating a Hessian-vector product, which is much more computationally expensive than the Jacobian-vector product, used in likelihood calculations.

* Did the authors test the method in the ODE scenario (with sigma = 0.0)? While it may be less interesting due to the higher feasibility of the non-entropic OT problem (e.g. solved by NOT), it should substantially decrease training time due to the lack of the Laplace operator in computations.

* Could the authors check whether there is a bug in the line 366? Here, the partial derivative with respect to time is approximated by the discrete-time approximation of the full time-derivative, which seems confusing.

* Despite training a stochastic map, in Figure 3 only one-to-one mappings are visualized. Is it a design choice, or here the generator is deterministic?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731081892419

thread_type:
official_review

<thread_object>
creator: reviewer_ixa3

<content>
summary:
This work introduces a method to solve the (semi-relaxed) entropic unbalanced optimal transport problem. First, different equivalent formulation of the problem (namely a dynamic formulation and a dynamic dual) are derived. Then, it is shown that the problem can be stated as a stochastic optimal control problem with a HJB constraint on the value function and a Fokker-Planck constraint for the distribution. The authors then propose to solve this problem by parameterizing the value function with neural networks. Finally, they experiment on generative modeling tasks, and compare on synthetic datasets how well the method allows to recover the optimal coupling and value.

strengths:
The paper is overall well written. It introduces a well motivated method by deriving a suitable dual dynamic formulation of the EUOT problem, and rewriting it as a stochastic optimal control problem.

The experiments on generative modeling seem good.

weaknesses:
On the presentation of the paper, I believe some things could be improved. For instance, in Section 4.1, a lot of equations are presented, which are obtained by performing only some rewriting (e.g. equation (21) to (22) and then (22) to (24)). While this is very clear, I would suggest to just report the last equation (24) with some description of what has been done, and report the detailed derivations in Appendix. Also, the legend of Figure 3 are too small. A lot of abreviations are not given or are given after their introduction (NFE, IPF, IMF, HJB...). The notations often change between $u_t(x)$ and $u(t,x)$.

The full procedure is not very clear. It would be better to add an Algorithm.

I believe some relevant baselines are missing from the Generative modeling experiment. I am thinking e.g. of [1] who solve the UOT problem and reported a FID of 2.97 on CIFAR10.

While this method computes an unbalanced OT problem, the robustness to outliers is not discussed nor tested.

[1] Choi, Jaemoo, Jaewoong Choi, and Myungjoo Kang. "Generative modeling through the semi-dual formulation of unbalanced optimal transport." Advances in Neural Information Processing Systems 36 (2024).

questions:
Line 31, the paper (Jordan et al, 1998) is cited for generative modeling. I don't think they do generative modeling at all.

Something more subjective, Line 107, I don't like the terminology of "stochastic map". It is not really a map. I would be more accurate to describe it as a (conditional) probability or probability kernel.

Since the first marginal of the problem is fixed, it would be clearer to state the problem as semi-unbalanced. Is it more difficult to derive Theorem 3.2 and Proposition 3.4 with both marginals relaxed?

The method relaxes a lot of constraints into the loss. I wonder how hard it is to make sure every term go to 0. Also, is this method costly compared e.g. to the semi-dual approach of [1]? And is it stable to train?

Line 512, it is stated that "there is no closed-form for the EUOT problem". But I believe there are in [2], and [2] is also the right reference for closed forms of the EOT problem between Gaussian. Therefore, I would suggest to also perform the experiments in Section 5.2 with unbalanced Gaussian.

Some references on generative modeling with unbalanced OT problem could be added, e.g. [3, 4].

Typos:
- Line 175: "Dynamcial"
- Line 513: "Conducted same benchmark"

[1] Choi, Jaemoo, Jaewoong Choi, and Myungjoo Kang. "Generative modeling through the semi-dual formulation of unbalanced optimal transport." Advances in Neural Information Processing Systems 36 (2024).

[2] Janati, H., Muzellec, B., Peyré, G., & Cuturi, M. (2020). Entropic optimal transport between unbalanced gaussian measures has a closed form. Advances in neural information processing systems, 33, 10468-10479.

[3] Dao, Q., Ta, B., Pham, T., & Tran, A. (2023). Robust Diffusion GAN using Semi-Unbalanced Optimal Transport. arXiv preprint arXiv:2311.17101.

[4] Lübeck, F., Bunne, C., Gut, G., del Castillo, J. S., Pelkmans, L., & Alvarez-Melis, D. (2022). Neural unbalanced optimal transport via cycle-consistent semi-couplings. arXiv preprint arXiv:2209.15621.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731116150467

thread_type:
official_review

<thread_object>
creator: reviewer_ygvj

<content>
summary:
This paper proposes a generalization of the Schrödinger Bridge Problem called Entropic Unbalanced Optimal Transport (EUOT). Leveraging a stochastic optimal control interpretation, the authors efficiently solve the dual formulation of EUOT. Through convex optimization, they optimize the value $V$ and path measure $\rho_t$ separately, resulting in two loss functions to construct a fully simulation-free algorithm. The paper claims that EUOT achieves one-step generation and improves generative model scalability, although the experiments don’t clearly support this level of performance.

strengths:
1. The work demonstrates novelty as the first EOT model that does not require pretraining. Based on the known equivalence between Schrödinger Bridge (SB) and EOT, the authors introduced EUOT by fixing the source distribution and proved that it generalizes EOT (and thus SB).

2. Thanks to its reciprocal property and static generator, EUOT is a simulation-free model that does not require SDE or ODE simulations for sampling.

weaknesses:
1. Theoretically, compared to EOT, EUOT only adds an $f$-divergence penalty term in the objective to relax the marginal constraint, which doesn’t seem like a groundbreaking innovation.

2. While the results surpass other OT methods, EUOT does not achieve the best FID scores compared to other image generation models (Tab. 1). Additionally, Fig. 3 indicates that EUOT does not fully approximate a precise optimal transport plan, and Tab. 3 shows that EUOT performs slightly worse than some benchmarks.

questions:
1. In Eq. (7), it appears that a $1 \over \sigma^2$ term may be missing in the objective. Could the authors provide a detailed derivation of Eq. (7) or reference a source that confirms this formulation? This would help clarify whether the equation is correct as presented or if any term is indeed missing.
2. The paper lacks a visualization of wild-to-cat transformations, which would help assess visual quality and compare with DSBM results at 512x512 resolution. Could the authors include wild-to-cat visualizations and provide a more detailed comparison with DSBM? Additionally, an explanation for the high FID in Table 2 compared to DSBM’s performance would clarify any discrepancies in resolution or model effectiveness.
3. Line 458 claims that Table 2 demonstrates the model's superior scalability; however, I don’t observe any changes in model size, and the dataset resolution remains the same. Could the authors clarify this claim?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731362092586

thread_type:
official_review

<thread_object>
creator: reviewer_rwwa

<content>
summary:
The authors establish several new connections between entropic unbalanced optimal transport (EUOT), the Schrödinger Bridge (SB) problem, stochastic control (SOC), and action matching. Furthermore, based on these new connections, the authors derive the dual formulation of the EUOT problem and its optimality conditions. In turn, this novel dual formulation allowed authors to propose a novel maximin simulation-free solver for the EUOT problem. The authors evaluate the proposed solver on 2D toy experiments, experiments with Gaussians, image generation (CIFAR-10), and unpaired image translation (Celeba Male to Female and Wild to Cat datasets).

strengths:
1) The novel dual formulation of Entropic Unbalanced Optimal Transport based on the connections between Entropic Optimal Transport, Schrödinger Bridge, and Action Matching seems promising in unifying the mentioned problems.

2) The obtained experimental results support the author's claim that the new simulation-free solver outperforms the previous approach for SB/Entropic Optimal Transport.

3) The results for image generations presented for CIFAR-10 are comparable with the diffusion/flow models while requiring only 1 NFE for inference.

weaknesses:
1) While the general theory is clear and strict, some questions arise about the claims in sections 4.1 and 4.2 regarding the description and justification of the proposed method (see questions).

questions:
1) In lines 308-309, the authors propose considering the dual formulation of Eq. 23 given by Eq. 25. However, Eq. 23 depends only on $V_1$, i.e., only the values of $V$ in the terminal time $t=1$, while Eq. 25 depends on the all-time moments of $V_t$. Hence, the optimization result given by Eq. 25 should depend only on the values of $V$ at the terminal times, which seems strange. Could you comment on this?

2) Does Eq. 25 require $V$ to be optimal, i.e., satisfy the HJB equation, or can it be used for any time-dependent function $V$? It is interesting because you propose to use the regularizer for the HJB condition in Eq. 26. In practice, at the arbitrary step of the algorithm, the HJB condition will unlikely hold so that it would influence the optimization of $\rho$.

3) In the algorithm (lines 229-230), you propose simultaneously performing optimization updates for both generator $T_{\theta}$ and value network $V$. However, the motivation for this choice is not clear to me. In Section 4.1, you say that you need to solve an additional optimization problem (Eq. 25) for $\rho$, which makes me think that this problem (Eq. 25) should be solved before updating the value function $V$.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper introduces Simulation-free Entropic Unbalanced Optimal Transport (SF-EUOT), a scalable approach to solving the EUOT problem, a generalization of the Schrödinger bridges problem. By deriving its dynamical form, dual formulation, and optimality conditions, the method avoids costly simulations and enables one-step generation via the reciprocal property. SF-EUOT demonstrates scalability and good performances in generative modeling and image-to-image translation tasks compared to traditional Schrödinger bridge-based methods.

Reviewers expressed some concerns related to the general performances of the method wrt. other diffusion/distillation based methods, and the scientific novelty and added value of the proposed approach. Authors provided novel experiments, that are satisfying, and made more explicit the connections with the semi-relaxed unbalanced entropic OT problem. Overall, I believe that the paper might not be ready for publication at ICLR, given the amount of changes that the paper undergone during the rebuttal, and that would deserve a new review round. In this light, I will recommend a reject option, and I encourage the authors to resubmit a new version of their work, focusing eventually on the points raised by the reviewers, i.e. justifying and focusing on the impact of unbalancing the OT (outliers / generalization perspective), and comparisons with diffusion based approaches.