Expectile Regularization for Fast and Accurate Training of Neural Optimal Transport

Q1: The ENOT and the Monge gap are similar in that they both serve as regularizers for inducing Neural Optimal Transport. Could you provide an additional quantitative comparison between ENOT and Monge gap? Currently, there are only qualitative examples in Fig 1 and 2.

We are thankful for this suggestion. We performed an additional comparison with the Monge Gap on $W_2$ benchmark (dataset from Table 2), which can be found below. Given the potential impact of these results on the ultimate algorithm selection for NOT, we propose to add the table to the main text.

Q2: Table 3 presents the FID results on the Unpaired Image-to-Image Translation task. On the other hand, Table 10 in the appendix includes additional MSE results. For the Unpaired Image-to-Image Translation task, the goal of Neural Optimal Transport is to accurately approximate the target distribution (FID) while minimizing the cost function (MSE). Hence, the MSE results should be included in the main part of the manuscript.

We put the results with the MSE metric in the Appendix, because the majority of baselines we compare against did not include the outcome for those tasks. For convenience, we performed the experiments with the baseline methods and can include them into the main part of the paper:

Other issues: Equation (17): please kindly refer to our response to Reviewer E19z above.

We are also thankful for spotting the typos, which we have corrected.

We are thankful for raising important questions and for the positive feedback.

First, we would like to address the concerns about novelty:

To the best of our knowledge, ENOT is the first approach that ventures into the approximation of the c-transform by means of expectile regularization. Moreover, ENOT is first to apply this approximation for efficient and accurate solution of Neural OT by completely mitigating additional $c$-conjugate optimization as done in prior works. These features provide at least a 3-fold better and a 10-fold faster training performance against the competitors.

Q1: The quantile regularization term biases the estimation of the conjugate operator, making it non-exact. Despite this, the authors note that the regularization enables the method to outperform all exact methods. Could the authors speculate on the reasons for this? Is it solely due to the numerical stability of exact methods?

We acknowledge the need to discuss this and would like to add the following text to the paper. In actuality, the expectile regularization does not introduce an additional bias, neither in theory, nor in practice. The formal convergence for the c-conjugate transform, when $\tau$ converges to 1, is discussed in Appendix D. At the end of training (or upon convergence), we obtain the exact estimate of the $c$-conjugate transformation. Other methods demand near-exact estimation at each optimization step, requiring additional inner optimization and introducing significant overhead. We assume that introduces an imbalance in the simultaneous optimization by $T$ and $g$ in Equation 6, underestimating the OT distance as a result.

Q2: Section 5.4 typo

Thank you for careful reading. We have made the correction.

Thank you for your response. I especially appreciate your discussion on Q1. It would be great to see a similar discussion included in the paper.

I've increased my score to 7. However, I must note that although I've used optimal transport in previous work, I don't consider myself an expert in the field. My support for the paper is mainly because of its strong empirical performance. For this reason, I recommend that the area chair give more weight to the opinions of OT experts when making the final decision.

Q1: Comparison with other methods is focusing mainly on Wasserstein-2 benchmark and image-to-image translation. Comparison with some popular in OT biology problems (see 1, 2) would improve the contribution.

We thank the reviewer for the links to these interesting articles and the application of OT in biology. We analyzed the cited works and concluded that our method is not explicitly suitable for solving the problems from these papers, because they require finding highly non-linear maps. But we have made a rough linear estimation on the proposed biological dataset (taken from TrajectoryNet paper) with our method, given in the table below. Despite learning linear maps, ENOT outperforms baseline methods based on score-matching (DSBM) and flow matching (SB-CFM).

Q2 The authors consider two training modes depending on the ‘is_bidirectional’ parameter. In what cases the bidirectional training mode works better and what it depends on?

The bidirectional mode works only with strongly convex functions $h(x-y)$. It significantly improves the solution when $T(x)$ is a discontinuous function (as in Figure 7). We propose to mention this in the Discussion Section.

Q3 The algorithm (Algorithm 1 line 2) states that in case 'is_bidirectional' is false . From what it follows?

When the parameter is_bidirectional = False, we can use any parameterization for $T(x)$ by a neural network. The expression $T= \nabla f$ is used to avoid introducing additional variables, but here, we assume any function.

Q4 Other algorithms that use expectile regression tend to use 0.9 expectile, why in your method you take the expectile parameter closer to 1?

The point is that with $\tau$ converging to 1, we get a more accurate estimate of the c-conjugate transformation. But from the experiments in Figure 4, we can conclude that the values in the range 0.9-1.0 also give a good result.

Q5 By what metric can we choose the hyperparameters

Conventionally, one could use L2-UVP metric for the experiments on W2 benchmark (Tables 1,2) and FID metric for the Image-to-Image tasks (Table 3).

Q6 Partial OT is often required in practical tasks. It is possible to generalize this method to solve the problem of partial OT?

Our method can be easily adapted for a special case of partial OT, when the whole source measure $\alpha$ and only a part of the target measure $\beta$ are used. Such problem is considered in Gazdieva et al. 2023 (citation from the paper). For the general case, we are currently unable to suggest a proper method of adaptation. Still, if a new method for a neural partial OT is proposed elsewhere, it will also positively benefit from the use of our expectile regularization.

Q7 L2 cost is probably not the best choice in the image-to-image translation task, can we use any cost, e.g. parameterised by the neural network, in coupling with your method?

Indeed, the L2 distance between the activations from VGG16 layers can be used as a cost function. We also conducted experiments with it, but it turned out that the ordinary L2 cost gives more accurate results. Perhaps, this stems from the fact that the faces are centered in all images, enabling a direct comparison.

Q1: Lack of references on the instability of finding optimal c-conjugate (the motivation of this work). What does it mean to make the optimization problem more "balanced" with the expectile regression regularization on (line 86)? The same applies when the authors claim expectile regression is a popular option for estimating maximum of distributions through neural networks (line 130).

We are thankful for questions and for suggesting extra citations. We would like to stress that the main motivation behind our work was in finding fast, simple and accurate solver. The instability of finding accurate c-conjugate transform is the main bottleneck in the existing neural OT solvers. Those problems were thoroughly discussed in the following works (refer to Refs in our paper, e.g., Amos 2022a and Korotin et al 2021). For example, Korotin et al 2021 (Section 4.3, page 7) state that "The exact conjugate solver is slow since each optimization step solves a hard subproblem for computing the conjugate. Solvers that approximate the conjugate are also hard to optimize: they either diverge from the start or diverge after converging to nearly-optimal saddle point". Moreover this claim is empirically supported by the results in Tables 1 and 2 (i.e., W2OT-Objective, W2OT-Cycle diverged in all W2 tasks). Remarkably, our proposed regularization method stabilizes the overall training procedure.

The term "more balanced" here means that we optimize optimal map $T$ and potential $g$ in the OT problem (Equation 6) synchronously and with the same frequency.

Recently, expectile regression was used in some offline Reinforcement Learning algorithms and representation learning approaches. Among them, Implicit Q-Learning (IQL) methods [R1, R2].

[R1] Ilya Kostrikov, Ashvin Nair and Sergey Levine. Offline Reinforcement Learning with Implicit Q-Learning, 2022.

[R2] Dibya Ghosh, Chethan Anand Bhateja, Sergey Levine. Reinforcement Learning from Passive Data via Latent Intentions, 2023.

Q2 Major: for eq (13) and eq (14) to hold you need both f and g to be convex functions (a la Brenier's theorem) with respect to the input; the only option is to use parameterization with ICNN. However, in table 4-5 and 7 the authors used non-convex MLP as parameterization; this is contradicting with the write up in methodology part. Could the authors elaborate on this?

In Section (3), the variables with hats ($\hat{\pi}$, $\hat{T}$, $\hat{f}$, $\hat{g}$) correspond to the solution (argmins) of the OT problem (6). The equations (13) and (14) hold under the specified conditions (source and target domains $X$, $Y$ are compact and the measures are absolutely continuous). So there is no strict requirement on $c$-concavity (or convexity); the potentials $f$ and $g$ must merely be the solution to the OT problem (ref. Santambrogio Theorem 1.17). During the training, we also may use the equations (13) and (14) with non-$c$-concave functions $f$ and $g$, because there are no restrictions on $T(x)$ in problem (6) and we can use any representation for the transport mapping function.

Under weaker constraints (for example $X=Y=R^n$), the $c$-concavity of the potentials $f$ and $g$ may be required. In this case, we can rely on a local $c$-concavity in the data concentration region. In addition, if the conditions in equations (13) and (14) are not satisfied, we can use the mode (is_bidirectional = False), in which we use an arbitrary function for $T(x)$ and do not express it through the potential.

In practice, the is no benefit in using ICNN, because it makes sense only for the squared Euclidean cost, and even in this case the optimization problem deals with additional constraints, becoming more difficult. Similar empirical observations of under-performance of ICNNs were reported in (Amos 2022a).

Q3: Ambiguous report of the empirical evaluations: most of the image task should use standard metric such as the FID. Many of the FIDs for the baselines are missing.

Respectfully, the missing values in Table 3 mean that the corresponding experiments were not performed in the baseline publications. Given the suggestion, we performed additional experiments using the official implementation of the OT methods reported by the other authors to fill in the gaps, which we include in a table below and propose to add to our revision.

Q4: What are $\hat{\pi}, \hat{T}, \hat{f}, \hat{g}$?

We are sorry for the ambiguity. In Section 3, the hats ($\hat{\pi}, \hat{T}, \hat{f}, \hat{g}$) correspond to the solution of the OT problem in equation 6. As written in the beginning of this section, $\hat{\pi}$ is the optimal transport plan, $\hat{T}$ - corresponds to the optimal transport mapping, $\hat{f}, \hat{g}$ are the optimal Kantorovich potentials. The notation without $\hat{\cdot}$ merely means that the optimization is done over the given variable.

Q5: Derivation of Equation 17:

We approximate the maximum of eq. (16) by $\tau$-expectile of $g^T(x) - c(x, y)$ conditioning on $y$. So the target (term $y$ in eq. 15) of the expectile regression here is $g^T(x) - c(x, y)$. The model is then $-g_{\eta}(y)$ with a negative sign as we approximate c-transform of $g^T(x)$ as $\inf_x c(x, y) - g^T(x)$. The corresponding regression loss is: $L_{\tau} \big( g^T(x) - c(x, y) + g_\eta(y) \big).$ Taking definition of $g^T(x)$ from eq. (7), we derive the expression (17): $L_{\tau} \big( c(x, T_\theta(x)) - g_\eta(T_\theta(x)) - c(x, y) + g_\eta(y) \big).$

Q6: Full Objective: is defined in lines 168-169. The expression is $L_g(\eta) + L_f(\theta) + \lambda R_g(\eta)$.

Q1: Thank you for requesting this clarification. Yes, as indicated in Sections 5.1 and 5.3, and in most of our experiments, we used squared Euclidean cost. Under Brenier's theorem conditions for the squared Euclidean cost, it holds that $\hat{T}(x) = \nabla \hat{f} (x)$, where $\hat{f}$ is some convex function. It follows that the optimal potential $\hat{f}$ is convex, even if one uses non-convex potentials $f$ in the training process.

Also, to reiterate our initial answer, there are no restrictions on $T(x)$ in Eq. (6) during the training and we can use any representation for this function, including $T(x) = \nabla f(x)$. Therefore, there is no need for $c$-concavity through the neural network parameterization whatsoever, and we do not use it, as it makes the solution worse in practice.

The beauty is that our expectile regularization actually forces functions to be more $c$-concave according to the $c$-concavity criterion (refer to Villani et al. [2009] Proposition 5.8).

 

Q2: We apologize for the ambiguity, which could have been avoided had we explicitly mentioned that we used flag is_bidirectional=False (meaning the training mode is one-directional in this example). Thank you for spotting the mismatch. So, in Section 5.2, we do not use (13) and (14) and we do not express $T(x)$ through the potentials, and the convexity of the cost function $h$ is not required. As such, we believe the issue should be resolved and there is no contradiction with Remark 1.18 in Santambrogio (2015). We will definitely add this additional clarification to the text.

Please kindly notice that we do mention it in another part of the main paper (citations below). So, by adding this clarification to Section 5.2 also, we avoid any potential for the same confusion of the future readers of the paper. Thank you.

Lines 155-157 (Method section): “The transport mapping $T_θ(x)$ has the same parameters as $f_θ(x)$ if it can be expressed through $f_θ$ (ref. 13), or otherwise, when $f_θ$ is not used (one-directional training), it is its own parameters.”

Lines 204-205 (Section 5.2): “we parametrize the map Tθ as a MLP”

Q1: The expectile loss is well suited to approximate the c-transform only if tau is large enough (close to $1$). Therefore, it would make sense to see a significant drop in performance when tau is close to $0.5$. It would be have been nice to validate this intuition in figure 4 for instance.

Thank you for raising this question. To characterize ENOT performance as a function of expectile $\tau$, we performed additional evaluation with $\tau$ ranging from $0.5$ to $0.999$ on 1) image-to-image translation dataset (CelebA(female) $\Rightarrow$ Anime with dim=$64$, Table $3$), 2) $W_2$ benchmark with dim=$256$ (Table $2$), and 3) 2D dataset from Figure $7$ (second row). In these experiments, we indeed observe a significant drop in performance when $\tau$ approaches $0.5$ ($0.7, 0.6, 0.5$) on the 2D dataset and CelebA (female) $\Rightarrow$ Anime with dim=$64$ (in terms of FID). On $W_2$ benchmark, the tendency is less evident. At the same time, values of $\tau$ in the range $\left[0.90, 1.0\right]$ always demonstrate convergence of ENOT, giving good results in all experiments. We will reflect on these new experiments in the revised manuscript.

Q2: I would expect this to lead to some form of tradeoff. Letting tau->1 should lead to some form of instability and/or slower convergence ? Is that the case ? How difficult is this tradeoff to settle in practice ?

Setting $\tau \simeq 1$ may indeed cause an instability. This can be the case because, under certain conditions, the overall contribution of proposed regularization term will be zero, which means that the potentials can become unbounded. However, in our experiments, such an instability occurred extremely rarely (mostly due to bad optimizer parameters), resulting only in a slight drop in performance. Yet, we carefully show that the optimal $\tau$ is usually in the range $\left[0.9, 1.0\right]$, with different values in this range yielding approximately the same result; so, selecting the optimal value of $\tau$ becomes simple.

Q3: The fitted maps of Sinkhorn in fig1 and fig2 are odd, I suspect the entropic regularization was very high, leading to a very very smooth transportation plan from which the fitted map was obtained by taking the maximum over values very close to each other.

Thank you for pointing this out. In order to be consistent with the experiments from the Monge gap paper (Uscidda and Cuturi [2023]), we used their exact default parameters of the entropic regularization, namely $\epsilon = 0.01$ $\cdot$ mean(C), with C being the cost matrix (in Figure 1, $\epsilon \sim 0.12)$, replicating the 'odd' look from the original paper. Moreover, we observe that lowering the values of $\epsilon$ produces plots of similar transportation smoothness.