Inverse Entropic Optimal Transport Solves Semi-supervised Learning via Data Likelihood Maximization

Dear reviewer, we thank you for already considering the general response to all the reviewers. Still we want to highlight that this answer does not take into account all your questions, so we provide a detailed answers to them below.

1. What is the intuitive interpretation of the new objective? This seems a bit ad-hoc/heuristic to introduce $f(y)$ the term for the sole purpose of using marginal samples from $p(y)$.

Our objective is the direct log-likelihood maximization of data. In particular, introducing $f_{\theta}$ is theoretically justified because instead of just parameterizing $E^{\theta}(y|x)$ with a function of two variables (as it should be), we parameterize it as a function of two variables $c^{\theta}(x,y)$ plus a function of one variable $f^{\theta}(y)$. While this parameterization is more expressive (one may just set $f^{\theta}(y)=0$ and $c^{\theta}(x,y)=E^{\theta}(y|x)$), it allows to estimate the part of the loss using unpaired samples. Our ablation studies, e.g., Table 2 (Real data) or Appendix B.2.4 (Swiss Roll), show that this technique of variable separation indeed practically works and allows unpaired samples to provide notable assistance to paired ones.

Also note that in the scenarios when the cost function is fixed (not learnable), our objective function (Equation 12) simplifies to the objective introduced [1, Equation 17], i.e., solving forward OT problem between unpaired samples for a given cost function. This connection provides further justification for the inclusion of the $f^\theta(y)$ term and highlights the generality of our approach.

2. Intuition and theoretic understanding of objective <...> What is the purpose of the objectives corresponding to the Y and X marginals? Why does this make sense? What if these are misestimated?

Our loss is simply the KL divergence with a carefully designed energy-based reparameterization (Equation 9 and 11). The three components of the loss are not separate but rather integral parts of a single log-likelihood loss. Consequently, removing or misestimating any of these components would result in an incorrect estimation of the overall KL divergence.

3. Why use the parametrization in 16 and 17? What is the intuition of this parametrization? <...> Can you give examples of cost functions that do not have this form? <...> How do you ensure $c^\theta(x, y)$ is a valid cost function?

We selected this parametrization because it allows us to derive a closed-form solution for the normalization constant $Z_{\theta}$, which is essential for efficient computation. While it may be hard establish an intuition for such a parameterization, it is possible to explain how we come up with this parameterization. In fact, we analyzed the parametrization used in [2, 3] for the inner product OT cost $\langle x, y \rangle$, and found that the normalizing constant remains computable if one considers cost $\langle b_m(x), y \rangle$. Then we found that multiple such costs may be also smoothed with the log-sum-exp function without losing the necessary integration properties.

Elaborating your question, we derived a new Theorem 3.1 which shows that our proposed parameterization allows to universally approximate a broad class of conditional distributions (and, therefore, transport costs). Furthermore, in response to questions of the reviewers, we added an additional Appendix A showing that our framework supports arbitrary neural parametrizations.

4. Also, why is it required to have an energy-based model? <...> score-based/diffusion-based generative models are far more common now...

The main motivation for the energy-based model parametrization is because it allows us enables a clear separation between paired and the marginal distribution (Section 3.1). At the same time, we emphasize that according to work [4], Energy-based models can be viewed as a general framework encapsulating score and diffusion models. Therefore, it is highly expected that our ideas open the door to future extensions, such as incorporating score-based or diffusion models, but this is left for further research.

5. How does this compare to the most natural likelihood objective that is $\mathcal{L}(\theta)=\mathbb{E}_{x, y \sim \pi^\star}[-\log\pi^\theta(y\vert x)] + \mathbb{E}_{y \sim \pi^\star_y}[-\log\mathbb{E}_{x \sim \pi^\star_x}[\pi^\star(y \vert x)]]$?

Thank you for highlighting this interesting loss. In fact, we already mentioned such approaches [5, 6] in Appendix B.2.2 of the original submission. As you correctly observe, this objective, while conceptually straightforward, presents computational challenges. Following you request, we added the detailed discussion of this approach to Appendix B.2.2 and performed the comparison with it in the revised paper. As Figure 2i indicates, it overfits with limited data, whereas our method (Figure 2j) remains robust.

Hi authors, I have read and reviewed your basic response. I think your response and edits did improve the paper somewhat though some of my concerns remain regarding intuition, deep understanding, and empirical validation. I will keep my score at this point.

Dear Reviewer,

Thank you for your detailed and insightful comments. Below, we address your points and elaborate on the revisions made to the paper in response to your feedback.

1. Reliance on Gaussian mixture parameterization. <...> What makes Gaussian mixture parameterization particularly suitable for the proposed loss compared to alternative parameterizations? <...> The discussion of alternative parameterizations...

Dear reviewer, you accurately captured the essence of our parameterization approach in your summary, where you noted that it "allows closed-form expressions for their loss terms". Indeed, the Gaussian Mixture Model (GMM) parameterization offers a significant advantage by enabling the normalization constant $Z_{\theta}$ to be evaluated in closed form. This capability facilitates closed-form expressions for the loss terms, thereby improving both computational efficiency and clarity.

We agree with you that the GMM parameterization may experience issues in scaling to high dimensions. To further elaborate your comment, we made two additional edits in the revised paper.

• First, we have introduced new Theorem 3.1 to demonstrate that our solver, using GMM-based parameterization, acts as a universal approximator. Under mild assumptions, it can approximate the true conditional distribution to an arbitrary degree of accuracy in terms of KL-divergence, provided that a sufficient number of Gaussian components are utilized. For further details, please refer to Section 3.4 and Appendix C.2 in the updated version of our paper.

• Second, we demonstrated the possibility to enable the neural parameterization for both the potential function $f^\theta$ and the cost function $c^\theta$ within our framework. This neural parametrization may potentially allow our framework to handle tasks is higher dimensions as detailed in the newly added Appendix A in our revised paper.

• Third, as we noted in the discussion section, we believe that our ideas have potential to be incorporated into advanced learning techniques, e.g., such as diffusion models. This is due to our framework being based on the well-established KL-divergence optimization principles. This serves as a promising avenue for future research.

2. How does the proposed method compare to existing semi-supervised domain translation techniques in terms of both theoretical properties and empirical performance?

Theoretical comparison. Our paper introduces a novel, unified framework for semi-supervised domain translation that seamlessly integrates both paired and unpaired data while providing strong theoretical guarantees. In contrast, existing methods that often rely on heuristically designed loss functions [1] or manually incorporate paired data into a conditional framework [2], our approach is grounded in solid theoretical principles.

• Our loss function (Equation 12) is derived from a KL-divergence minimization framework. When this loss function converges, it ensures the recovery of the true conditional plan $\pi^\star$, as $\text{KL}(\pi^\star \Vert \pi^\theta) = 0$. To our knowledge, this is the first method to address the problem in this manner, offering a principled theoretical foundation.

• Our loss function is related to inverse entropic optimal transport (OT), as discussed in Section 3.2. This connection enables ML researchers to leverage a rich set of existing methods for solving semi-supervised domain translation tasks using the already well-established techniques from computational OT.

Empirical Properties & Comparative Performance. Following your comment about the experiments and the comments of other reviewers, we added various baselines to the real-world data experiment and revised Section 5.2. In particular, following the suggestion of the other reviewer Nafh, we added two additional semi-supervised baselines (CNF-SS and CGMM-SS, see Appendix B.2.2). Our method consistently outperforms the alternatives in the considered experiment.

At the same, we understand your concern about comparison with the existing semi-supervised methods for image-to-image (I2I) translation [1,3]. Here we highlight two things.

• First, our the goal of paper is to make the first step toward establishing theoretically-grounded methods for semi-supervised domain translation rather than design a large-scale approach for complex I2I tasks. As we already pointed above, there is a feasible potential to scaling our framework to large-scale problems via neural parameterization (Appendix A).

• Second, our experimental evaluation already covers a kind of principal versions of losses which are used for I2I (e.g., Conditional and Unconditional GANs) and clearly demonstrates that they fail to learn the true conditional data distributions even in 2D.

Thank you for taking the time to review our paper and provide us useful feedback. Below are the answers to your questions.

1. Justification of Loss Function: The paper lacks sufficient theoretical or empirical justification for the specific loss function design. The decomposition into paired and unpaired data components seems plausible, but more discussion on its necessity or optimality would strengthen the work. Our loss is essentially the KL divergence with a smart energy-based reparameterization (equation 11) to be able to consider both paired and unpaired data.

• The theoretical justification for our loss function is that it reaches the minimum if and only if the optimized distribution $\pi^{\theta}$ coincides with the true data distribution $\pi^{*}$ that we need to learn ($KL(\pi^{\star},\pi^{\theta})=0$ if and only if $\pi^{\star}=\pi^{\theta}$).
• The theoretical justification for our energy-based reparameterization (11) through $E^{\theta}$ and its decoupling to $c^{\theta}$ and $f^{\theta}$ is that such parameterization is not any restrictive and can represent the true data distribution $\pi^{\star}(y|x)$. To achieve this optimality, term $c^{\theta}(x,y)$ in our parametrization can be set to $-\log \pi^{\star}(y|x)$, $f^{\theta}(y)=0$. In this case when $\varepsilon= 1$, from (9) one immediately gets that $Z_{\theta}(x)=1$. As a result, from (11) one derives $\pi^{\theta}(y|x)=\exp(-E_{\theta}(y|x))/Z_{\theta}=\exp(-\log\pi^{\star}(y|x))=\pi^{\star}(y|x).$ In practice, the considered parametric class for $c^{\theta},f^{\theta}$ may be insufficient to exactly represent $\pi^{\star}$. Nevertheless, in response to your question, we included a new theoretical result showing that our considered GMM-based parameterization can still minimize the loss up to any desired small error, see Section 3.4 in the revision.

2. Additionally, while the paper claims the loss solves inverse entropic OT, the practical implications of this for general semi-supervised tasks are underexplored. Thanks to our established connection, now ML researchers may be able to leverage a rich set of existing methods for computing OT problems for solving semi-supervised domain translation tasks.

• First, our paper already demonstrates how recent ideas in computational OT based on Gaussian mixtures can be exploited to design a theoretically-justified algorithm for semi-supervised domain translation.
• Second, answering the questions of the reviewers, we added a new Appendix A (in the revised version) which provides a discussion and a proof-of-concept experiment showing that other recent advances on solving OT with neural networks can also be used within our framework.
• Third, we believe that our findings may motivate future researches to generalize our ideas proposed here to more general semi-supervised setups. This direction is a promising avenue for follow-up works.

3. The experimental section lacks comparisons with several prominent semi-supervised learning methods that do not rely on OT or Gaussian mixtures. Without comparisons to these methods, it is challenging to gauge whether the proposed approach offers an advantage over non-OT-based techniques. Following your suggestion, we added various GAN and Normalizing Flow-based methods to the real-world data experiment and revised Section 5.2. In particular, following the suggestion of the other reviewer Nafh, we added two additional semi-supervised baselines (CNF-SS and CGMM-SS, see Appendix B.2.2). Our method consistently outperforms the alternatives in the considered experiment.

Concluding remarks. Please respond to our post to let us know if the clarifications above suitably address your concerns about our work. We are happy to address any remaining points during the discussion phase; if the responses above are sufficient, we kindly ask that you consider raising your score.