[Paper Review] Learning Cost Functions for Optimal Transport

Ma, Shaojun, Haodong Sun+3 more|arXiv (Cornell University)|Feb 22, 2020

Sparse and Compressive Sensing Techniques45 references6 citations

TL;DR

This paper proposes a novel convex optimization framework for learning cost functions in optimal transport (OT) by reformulating inverse OT as an unconstrained problem, avoiding the computational bottleneck of bi-level optimization. It introduces two algorithms—matrix scaling via Sinkhorn-Knopp for discrete OT and a deep neural network-based method for continuous OT—achieving superior efficiency and accuracy over prior methods.

ABSTRACT

Inverse optimal transport (OT) refers to the problem of learning the cost function for OT from observed transport plan or its samples. In this paper, we derive an unconstrained convex optimization formulation of the inverse OT problem, which can be further augmented by any customizable regularization. We provide a comprehensive characterization of the properties of inverse OT, including uniqueness of solutions. We also develop two numerical algorithms, one is a fast matrix scaling method based on the Sinkhorn-Knopp algorithm for discrete OT, and the other one is a learning based algorithm that parameterizes the cost function as a deep neural network for continuous OT. The novel framework proposed in the work avoids repeatedly solving a forward OT in each iteration which has been a thorny computational bottleneck for the bi-level optimization in existing inverse OT approaches. Numerical results demonstrate promising efficiency and accuracy advantages of the proposed algorithms over existing state-of-the-art methods.

Motivation & Objective

• Address the challenge of learning the cost function in optimal transport from observed transport plans, which is critical for accurate data representation and inference.
• Overcome the computational inefficiency of existing bi-level optimization approaches in inverse OT, which require repeatedly solving forward OT problems.
• Develop a general, unconstrained convex formulation of inverse OT that supports customizable regularization and enables theoretical analysis.
• Provide theoretical guarantees on solution uniqueness and bounded variation of dual variables in the inverse OT problem.
• Design efficient numerical algorithms that scale to high-dimensional settings without iterative forward OT solves.

Proposed method

• Formulate inverse OT as an unconstrained convex optimization problem using the entropic regularized OT framework, enabling efficient optimization and theoretical analysis.
• Propose a matrix scaling algorithm based on the Sinkhorn-Knopp algorithm to solve the discrete inverse OT problem by alternatingly updating dual variables and the cost matrix.
• Introduce a deep learning-based method that parameterizes the cost function as a neural network for continuous OT, enabling end-to-end learning of complex cost structures.
• Incorporate regularization into the inverse OT formulation to promote desired properties such as smoothness or sparsity in the learned cost function.
• Apply block coordinate descent (BCD) with variable rescaling to maintain bounded dual variables, ensuring numerical stability and convergence.
• Use a rescaling strategy to maintain dual variables within bounded sets, which enables convergence guarantees via known results on BCD for Lipschitz-smooth problems.

Experimental results

Research questions

• RQ1 Can inverse optimal transport be reformulated as an unconstrained convex optimization problem to avoid the computational burden of bi-level optimization?
• RQ2 What theoretical properties—such as solution uniqueness and bounded variation of dual variables—hold for the inverse OT problem under the proposed formulation?
• RQ3 How can efficient numerical algorithms be designed to solve inverse OT without repeatedly solving the forward OT problem?
• RQ4 Can deep neural networks effectively parameterize cost functions in continuous inverse OT, enabling flexible and data-driven cost learning?
• RQ5 What convergence guarantees can be established for the proposed algorithms, particularly in the context of block coordinate descent with bounded variables?

Key findings

• The proposed inverse OT formulation is unconstrained and convex, enabling efficient optimization and theoretical analysis of solution uniqueness and stability.
• The matrix scaling algorithm based on Sinkhorn-Knopp converges with a global sublinear rate of O(1/k), as shown by the convergence bound in Theorem 10.
• The algorithm maintains bounded dual variables through rescaling, which ensures numerical stability and enables convergence guarantees.
• The deep learning-based method successfully parameterizes the cost function as a neural network, allowing for flexible, data-driven cost learning in continuous OT settings.
• Numerical results demonstrate that both proposed algorithms outperform state-of-the-art methods in terms of efficiency and accuracy, particularly in high-dimensional settings.
• The framework avoids the repeated forward OT solves that plague traditional bi-level optimization, significantly reducing computational cost in inverse OT.