Title: Solving Discrete (Semi) Unbalanced Optimal Transport with Equivalent Transformation Mechanism and KKT-Multiplier Regularization
Venue: NeurIPS 2025
Accepted: yes
Keywords: Discrete Semi-Unbalanced Optimal Transport; Domain Adaptation with Applications
Abstract: Semi-Unbalanced Optimal Transport (SemiUOT) shows great promise in matching two probability measures by relaxing one of the marginal constraints. Previous solvers often incorporate an entropy regularization term, which can result in inaccurate matching solutions. To address this issue, we focus on determining the marginal probability distribution of SemiUOT with KL divergence using the proposed Equivalent Transformation Mechanism (ETM) approach. Furthermore, we extend the ETM-based method into exploiting the marginal probability distribution of Unbalanced Optimal Transport (UOT) with KL divergence for validating its generalization. Once the marginal probabilities of UOT/SemiUOT are determined, they can be transformed into a classical Optimal Transport (OT) problem. Moreover, we propose a KKT-Multiplier regularization term combined with Multiplier Regularized Optimal Transport (MROT) to achieve more accurate matching results. We conduct several numerical experiments to demonstrate the effectiveness of our proposed methods in addressing UOT/SemiUOT problems.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper tackles the problem of efficiently obtaining exact solutions for the Unbalanced Optimal Transport and Semi-Unbalanced Optimal Transport problems. The proposed method, called the Equivalent Transformation Mechanism (ETM), enables the determination of marginal probabilities on which a classical Optimal Transport optimization can be performed.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths: \n- The article addresses a fundamental problem in the optimal transport community. Rather than relying on popular regularization terms in the UOT and Semi-UOT problems—which allow for efficient computations but yield inexact and dense solutions—the authors propose a novel idea: transforming these problems into classical optimal transport problems by computing the exact marginal probability distributions that participate in the transport process. Moreover, for the Semi-UOT version, the authors propose a way to smooth out the $L_p$ functional, highlighting its usefulness to get a good initialization for the main method of interest. In addition, the proposed ETM approach presented in Section 4.1 allows for the computation of a so-called multiplier \"s\" (Corollary 2), which is used to formulate a new regularized version of the optimal transport problem (Section 4.2).\n- Experiments were conducted with synthetic and real-world datsets. \n\nWeaknesses: \nThe algorithms are provided in the appendix, but the methodology is somewhat convoluted to follow in the main text, even though the authors include all the necessary components. I understand that space limitations can make it difficult to fully detail the steps that a reader would ultimately need to follow.\n\nQUESTIONS:\nQuestions:\n\n1) In the well-known article \"Free boundaries in optimal transport and Monge-Ampère obstacle problems\" by L. Caffarelli and R. J. McCann, an unbalanced formulation of optimal transport is transformed into a classical optimal transport problem. Is there any connection between that approach and the one taken in the article under review?\n\n2) The classical unbalanced problem typically includes constraints such as $\\pi 1_N \\leq a$ and $\\pi^T 1_M \\leq b$. These are not included in the formulation of the UOT problem in the paper. Is there a particular reason for omitting them?\n\n3) Proposition 1 is the first key result. However, the notation is not fully clear at that point in the article. Could the authors clarify the use of the star symbol $\\ast$ in expressions like $f^\\ast$ and $\\zeta^\\ast$? (A similar comment applies to Proposition 3.)\n\n4) Regarding the statement of Proposition 1: I suggest rephrasing it as, \"Given the SemiUOT functional with KL divergence $J_{\\text{SemiUOT}}$...\" A similar revision could be applied to Proposition 3.\n\n5) Do the new marginals $\\alpha$ and $\\beta$ represent vectors encoding the weights of discrete probability measures? That is, do their coefficients add to 1?\n\n6) I did not understand the first sentence on line 180. Could the authors clarify it?\n\nSuggestions:\n\n1) Corollary 1 is not formally stated, so I would suggest presenting it as a Remark instead.\n\n2) While not necessary, a brief discussion of the advantages and disadvantages of SemiUOT compared to UOT would enrich the paper.\n\n3) In the Related Works section, I suggest including a citation of the paper [Y. Bai et al., \"Linear Optimal Partial Transport Embedding,\" ICML, 2023], as it presents an alternative approach to addressing the computational complexity of unbalanced optimal transport problems, without relying on the addition of regularization terms.\n\n4) I suggest referencing Algorithms 1 and 2 explicitly in the main text, even if only briefly, to guide the reader.\n\nTypos and Minor Comments:\n\n1) Line 44: After \"i.e.,\" I suggest using a lowercase letter and not capitalizing the following word.\n\n2) Lines 179, 180, and the paragraph between lines 193 and 199: Add commas after \"Hence,\" \"Then,\" and \"After that,\" respectively.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a new framework for solving discrete Semi-Unbalanced Optimal Transport (SemiUOT) and Unbalanced Optimal Transport (UOT) problems. Existing methods often rely on entropy regularization, which leads to dense and inaccurate transport plans. To address this, the authors introduce the Equivalent Transformation Mechanism (ETM), which reformulates SemiUOT and UOT problems into classic optimal transport problems. The authors develop three variants: ETM-Exact, ETM-Approx, and ETM-Refine.\n\nThe paper also proposes Multiplier Regularized Optimal Transport (MROT), which uses KKT multipliers to guide the transport plan and improve solution quality. \n\nExperiments on synthetic data, domain adaptation benchmarks, and treatment effect estimation show that ETM combined with MROT improves accuracy and efficiency over existing methods.\n\nSTRENGTHS AND WEAKNESSES:\nQuality: The paper presents a solid and well-justified theoretical contribution. The proposed Equivalent Transformation Mechanism (ETM) is rigorously derived, and its variants (ETM-Exact, ETM-Approx, ETM-Refine) are grounded in convex analysis and KKT conditions. The integration with Multiplier Regularized Optimal Transport (MROT) provides a framework to enhance the sparsity and accuracy of transport plans. The empirical evaluation is extensive, covering both synthetic and real-world datasets across domain adaptation and treatment effect estimation tasks. \n\nOriginality: The paper introduces a new perspective by transforming discrete SemiUOT and UOT problems into classic OT through reweighted marginals. The use of KKT multipliers as a regularization term in MROT appears to be an interesting idea that improves transport plan quality. \n\nSignificance: This proposed method has potential impact on applications where interpretability and sparsity are important, such as domain adaptation, generative modeling, and causal inference.\n\nClarity: The paper is generally clear and well-organized. The connection between the ETM variants and standard OT is carefully explained. The derivations are detailed, and algorithmic procedures are summarized in pseudocode, which aids reproducibility.\n\nQUESTIONS:\n1. In line 133, 'We iteratively update $L_U$ to reach the optimal solutions'. It seems that $L_U$ doesn't appear before this sentence. \n\n2. In line 162, 'Since $\\hat{f}^\\star$ is close to $f^\\star$'. Is there a theoretical bound for this?\n\n3. Figure 4 shows that the choice of $\\epsilon$ significantly affects the convergence behavior of ETM-Approx and ETM-Refine. For larger values (e.g., $\\epsilon=0.1$), the algorithm fails to recover accurate marginals. However, the paper does not provide guidance on how to choose $\\epsilon$ in practice. Could the authors clarify whether there is a principled way to select or adapt $\\epsilon$ especially for varying values of $\\tau$ or under different data scales?\n\n4. Since the ETM framework reformulates UOT and SemiUOT problems as classic OT with new marginal weights derived from KKT conditions, could the authors explain more clearly how and why these new weights lead to better performance under noisy or mismatched data? Specifically, in what sense does solving OT with reweighted marginals provide robustness or accuracy improvements over applying classic OT directly to the original noisy distributions?\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-24"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper addresses the challenge of accurately solving discrete SemiUOT and UOT problems. They propose an Equivalent Transformation Mechanism (ETM) approach to determine the marginal probability distributions of SemiUOT with KL divergence. A KKT-Multiplier regularization term is proposed with Multi plier Regularized Optimal Transport (MROT) to achieve more accurate matching results. Extensive experiments demonstrate the effectiveness of our proposed methods in addressing SemiUOT and UOT problems. Overall, this paper contributes a practically effective framework for solving SemiUOT and UOT.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n\n1. The SemiUOT and UOT problems are clarified explicitly as relaxations of classical OT. The ETM formulation allows the exact and approximate solutions and reduces the complexity of solving these problems. It is effective to use the solved marginal distributions for the transfer to standard OT problems.\n2. This method balances the trade-off between computing exact solutions and improving efficiency in ETM, and further leverages the obtained exact multipliers for regularization, resulting in sparser and more accurate transport matching.\n3. The comprehensive experiments demonstrate the effectiveness of ETM+MROT, which consistently yields state-of-the-art results across a diverse set of metrics in many applications.\n\nWeaknesses:\n\n1.\tTheoretically, coherence and completeness of this paper should be improved.  On one hand, the authors should first state assumptions, such as the existence of a c-transform and other implicit conditions used in the proofs, and summarize the conclusion in a theorem. On the other hand, there is a lack of detailed proof and analysis of the convergence rate.\n2.\tThe writing and organization could be improved to enhance clarity. The organization and layout seems disordered, and the inclusion of numerous detailed explanations in the main text hinders the understanding of the conclusions. Moreover, the paper lacks a diagram or concise pseudocode illustrating the overall algorithm overview. Typos, such as the incorrect notation in equation (20) and the misuse of inner product symbols in equations (23) and (25), further affect the readability of the paper.\n3.\tA better approach would be to design a customized optimization algorithm tailored for solving the OT problem, rather than directly applying L-BFGS. On one hand, leveraging the inherent structure of the problem may lead to faster computation. On the other hand, we are also concerned with the convergence behavior of the optimization algorithm for this specific problem.\n4.\tOne potential theoretical weakness is the lack of discussion on robustness to outliers. In real-world datasets, noisy or out-of-distribution samples can degrade OT performance. It would be valuable for the authors to either empirically or theoretically analyze how their approach behaves under such perturbations.\n5.\tThe novelty of the proposed approach may be somewhat limited. ETM resembles the discrete form of Equation (9) in Choi et al. (2023) under KL regularization, and MROT adapts standard regularized OT solvers with an additional multiplier. Clarifying the conceptual gap from prior work would strengthen the perceived contribution.\n\nQUESTIONS:\n1.\tIn (27), $f^*$ seems also relative to the $\\tau$ and it is ignored when setting $\\tau \\to \\infty$.\n2.\tThere are many computations of the form $\\exp(\\frac{\\cdot}{\\epsilon})$ in ETM+MROT-ENT. Could numerical instability arise when $\\epsilon\\to 0$?\n3.\tIn MROT, KKT multipliers are used as a regularization mechanism to induce sparsity. What is the theoretical validity for this design, and how does it affect optimization stability or convergence?\n4.\tWhat is the relationship between the optimal solution and optimal value of ETM-Exact and ETM-Approx? A theoretical explanation is  recommended.\n5.\tIt would be positive if the authors demonstrate the adaptability of this method to outliers.\n6.\tIf $u^*$ and $v^*$ are inexact due to the computational error, will the Corollary 2 and Proposition 4 still hold?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-06-28"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper provides a new technique for solving the unbalanced-OT (UOT) and semi unbalanced-OT (semi-UOT) problems. Both problems allow for partial optimal transport solutions (where not all mass is transported). In addition to the cost of the transport plan, an additional penalty function is added, which contributes to the objective function based on the amount of untransported mass. (In semi-OT, only one distribution is considered in this additional penalty function, while UOT considers both sides, with possibly different coefficients). Different variants of this “untransported mass” penalty function have been considered, but this paper specifically uses KL divergence.\n\nThe key contribution of this technique is a transformation, where any instance of the UOT or semi-UOT problem (under KL-divergence) can be transformed into an instance of the regular OT problem, in such a way that the solution to this regular OT problem instance provides a solution to the original UOT or semi-UOT problem. This transformation is named the “Equivalent Transformation Mechanism” (ETM). Furthermore, multiple variants of ETM are presented, offering tradeoffs in computation time and accuracy – these are named ETMExact, ETMApprox, and ETMRefine.\n\nAs both the semi-UOT problem and UOT problem have been investigated before, this paper seeks to differentiate itself in terms of quality, with a specific emphasis on not using entropy regularization like most previous methods, resulting in sparser solutions. The paper also provides extensive experimental evaluation on real and synthetic data. These include multiple variations of ETM, multiple prior work methods / baselines, and some different penalty functions for the untransported mass.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths : \n-\tVery extensive experimental evaluation, which shows an improvement over existing in terms of accuracy when applied to real-world applications.\n\n-\tThorough treatment of technical details. \n\n-\tThe technique is flexible in that any OT solver can be used along with ETM, and it can be easily applied to UOT and Semi-OT. In some sense, ETM seems to act as a “booster” of sorts when added to some existing approaches, based on the experimental results.\n\nWeaknesses : \n-\tThe experimental evaluations, in terms of accuracies, offer relatively modest improvements over the existing approaches, usually a couple percentage points (although this improvement is quite consistent). I have some concern about the experimental design.\n\n-\tSome parts of the complexity analysis of the ETM approach seem unclear in relation to other work. I would like to clearly see a comparison of Big-O time complexities of the approaches and steps of this work as well as descriptions of Big-O time complexities of other work. I want to be sure that the solution-time tradeoff of your approach is comparable to existing approaches.\n\n-\tI have a few questions about the experiments, which might be easily addressable in the author response.\n\nQUESTIONS:\nQuestions:\n-\tJust a question : Is it not true that Semi-OT is a special case of UOT, since you can always set the hyperparameter coefficient of one distribution to 0? I ask this because, one of your selling points is that your approach works for both. Why can’t existing UOT approaches apply to Semi-OT directly?\n\n-\tUnder table 1, you take the best approach without UOT, called JUMBOT, and you add to it various methods of computing UOT, such as MM-UOT, GEM-UOT, $\\ell_2$ norm solver, and sparse solver. All of these UOT approaches offer a modest improvement over JUMBOT. Then, you show that your ETM + MROT approach improves upon the previous UOT approaches for this task by an additional modest amount. However, I find it odd that, for all of your other similar tables, including Tables 2, 4, 5, 6, 7, you omit all variants of “state of the art + UOT(some other solver that’s not yours)”, and instead only consider your approach for the UOT component. Based on the trend in Table 1, it seems as if all UOT variants, yours or not, seem to improve upon the baseline “state of the art” without UOT. That’s not so important for your paper. Instead, your burden of proof is to illustrate that your variant of UOT specifically outperforms other existing variants of UOT in practice. Otherwise, what’s the role of the experiments here?\n\n-\tPlease detail how your approach compares to the existing UOT approaches with regards to time complexity and practical performance. The only clear discussion I saw of this is Figure 3, which doesn’t quite clear things up for me. The thorough treatment of running time seems lacking in both the theory as well, since the iteration count is not bounded. Furthermore, when I see sizes like N = 50 for the experiments, I start to wonder why you can’t bring N higher. Often, OT is fine with N = 1000 or more without an issue. How fast is your approach?\n\n-\tLine 282 : You performed 5 experiments, but that is not a very large sample. How sure can you be that, if the experiment were repeated, the average would be comparable to what you have here? \n\n-\tFigure 3 : You have a couple of charts / labels swapped here. Please address this. I would also recommend adding a consistent color scheme for the plots across UOT and semi-UOT charts respectively. Furthermore, the list of variants of your algorithm between (b) and (c) are different. What I would really like to see is an accuracy versus runtime comparison of your approach(es) versus existing approaches. Finally, the only running times given here are for $\\tau = 1$, which you stated in your introduction other approaches work well for. So, why not give your efficiency for large $\\tau$?\n\nCorrections : \n\nLine 34 – 35 : Grammar : “While … solutions” is an incomplete sentence. Reword.\n\nLine 58 : Did you mean “MROT”, not “MGOT”?\nAlso, see comments about Figure 3 above.\n\nLIMITATIONS:\nNothing to note.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 2, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-02"}