[Paper Review] A Simulated Annealing based Inexact Oracle for Wasserstein Loss Minimization

Jianbo Ye, James Z. Wang+1 more|arXiv (Cornell University)|Aug 12, 2016

Adversarial Robustness in Machine Learning20 references3 citations

TL;DR

This paper proposes Gibbs-OT, a simulated annealing-based inexact oracle using Gibbs sampling to efficiently approximate Wasserstein loss gradients in optimization problems. By enabling numerical stability, warm starts, and low per-iteration complexity, Gibbs-OT accelerates Wasserstein non-negative matrix factorization and optimal transport with Coulomb cost, outperforming entropy-regularized methods in speed and component sharpness while maintaining robustness.

ABSTRACT

Learning under a Wasserstein loss, a.k.a. Wasserstein loss minimization (WLM), is an emerging research topic for gaining insights from a large set of structured objects. Despite being conceptually simple, WLM problems are computationally challenging because they involve minimizing over functions of quantities (i.e. Wasserstein distances) that themselves require numerical algorithms to compute. In this paper, we introduce a stochastic approach based on simulated annealing for solving WLMs. Particularly, we have developed a Gibbs sampler to approximate effectively and efficiently the partial gradients of a sequence of Wasserstein losses. Our new approach has the advantages of numerical stability and readiness for warm starts. These characteristics are valuable for WLM problems that often require multiple levels of iterations in which the oracle for computing the value and gradient of a loss function is embedded. We applied the method to optimal transport with Coulomb cost and the Wasserstein non-negative matrix factorization problem, and made comparisons with the existing method of entropy regularization.

Motivation & Objective

• To address the computational bottleneck in Wasserstein loss minimization (WLM), where exact optimal transport (OT) solves are prohibitively expensive at O(m³ log m) complexity.
• To develop a fast, numerically stable inexact oracle that supports warm starts and early stopping for multi-level iterative WLM problems.
• To enable efficient optimization in WLM tasks such as Wasserstein NMF and optimal transport with Coulomb cost, where repeated OT solves are required.
• To provide an alternative to entropy regularization that avoids smoothing artifacts and improves component sharpness in factorization.

Proposed method

• Uses a time-inhomogeneous Gibbs sampler as a stochastic oracle to approximate partial gradients of Wasserstein losses.
• Applies simulated annealing to gradually lower temperature, enabling convergence to low-energy states resembling optimal transport plans.
• Leverages the dual formulation of OT to simplify the Metropolis-Hastings proposal into a Gibbs update, improving efficiency.
• Integrates the Gibbs-OT oracle into online mirror descent (OMD) and accelerated mirror descent (AMD) for joint optimization of basis vectors and membership weights.
• Employs temperature annealing: T := T(1 − √(1/(m + m̄))), allowing gradual refinement of transport plans.
• Supports warm starts by reusing previous MCMC samples as initial states, reducing mixing time across successive OT problems.

Experimental results

Research questions

• RQ1 Can a Markov chain Monte Carlo approach like Gibbs sampling provide a numerically stable and efficient inexact oracle for Wasserstein loss minimization?
• RQ2 Does the Gibbs-OT method enable faster convergence in WLM problems through effective warm starts and early stopping?
• RQ3 How does Gibbs-OT compare to entropy regularization in terms of solution quality and computational efficiency for Wasserstein NMF?
• RQ4 Can Gibbs-OT preserve sharper components in factorization tasks compared to smoothed entropy-regularized methods?

Key findings

• Gibbs-OT achieved average epoch times of 0.84 seconds on MNIST “5” images and 16.8 seconds on ORL face images, representing a two-order-of-magnitude speedup over commercial LP solvers.
• The learned components using Gibbs-OT were sharper than those from entropy-regularized W-NMF, indicating better preservation of structural details.
• Salt-and-pepper noise was observed in face image components, consistent with the Wasserstein distance’s insensitivity to sub-pixel mass displacements.
• The method demonstrated improved numerical stability and robustness compared to traditional OT solvers, especially in iterative WLM settings.
• Gibbs-OT enabled effective warm starts, significantly reducing per-iteration time in multi-epoch optimization due to reuse of prior MCMC states.