On Model-Based Clustering With Entropic Optimal Transport

TL;DR

This paper introduces a new entropic optimal transport-based loss function for model-based clustering, which improves optimization landscape and outperforms traditional log-likelihood methods in practical applications.

Contribution

The authors propose a novel loss function based on entropic optimal transport that shares the global optimum with the log-likelihood but has a better landscape for optimization.

Findings

The new loss function avoids spurious local optima common with the log-likelihood.

Sinkhorn-EM algorithm converges at a rate comparable to EM.

Numerical experiments and real-world applications show the new method outperforms traditional approaches.

Abstract

We develop a new methodology for model-based clustering. Optimizing the log-likelihood provides a principled statistical framework for clustering, with solutions found via the EM algorithm. However, because the log-likelihood is nonconvex, only convergence to stationary points can be guaranteed, and practitioners often use multiple starting points in the hope that one will converge to the global solution. We consider a new loss function based on entropic optimal transport that shares the same global optimum as the log-likelihood but has a much better-behaved landscape, thereby avoiding spurious local-optima configurations that are pervasive with the log-likelihood. Similar to the EM algorithm for the log-likelihood, this new loss can be optimized by the Sinkhorn-EM algorithm, which we show converges at a rate comparable to that of EM. By analyzing extensive numerical experiments and two real-world applications in image segmentation in C. elegans microscopy and clustering in spatial transcriptomics, we show that this new loss outperforms log-likelihood optimization, indicating that it represents a valuable clustering methodology for practitioners.