A universal network strategy for lightspeed computation of entropy-regularized optimal transport

Area of Science:

• Machine Learning
• Computational Mathematics
• Optimization

Background:

• Optimal Transport (OT) measures discrepancies between probability distributions.
• Entropy-regularized OT via Sinkhorn algorithm is computationally intensive for neural networks.
• High accuracy demands lead to large computation graphs, consuming significant time and memory.

Purpose of the Study:

• To develop a novel network strategy for estimating the transport matrix in OT.
• To reduce the computational complexity and memory footprint of OT calculations within neural networks.
• To create a method suitable for arbitrary cost functions and varying marginal distributions.

Main Methods:

• A novel network strategy to estimate the transport matrix, bypassing the Sinkhorn algorithm.
• Implementation in the log domain using the dual form to prevent numerical instability.
• Theoretical error bound estimation for approximate inputs.

Main Results:

• Significantly reduced computation graph size compared to Sinkhorn-based methods.
• Demonstrated suitability for arbitrary cost functions and varying marginal distributions.
• Outperformed baseline methods in both computation cost and accuracy across extensive experiments.

Conclusions:

• The proposed network strategy offers a more efficient and accurate alternative for OT computations in machine learning.
• The method is robust and adaptable to various cost functions and data distributions.
• Extensions to robust OT and barycenter computation are feasible and effective.