Modified K-means Algorithm with Local Optimality Guarantees

TL;DR

This paper introduces modifications to the K-means algorithm that guarantee local optimality under certain conditions, improving clustering results while maintaining computational efficiency.

Contribution

The paper provides conditions for local optimality in K-means and proposes simple modifications that ensure local optimality with the same complexity.

Findings

Modified K-means achieves better local optima in experiments.

The proposed method maintains the original algorithm's computational complexity.

Numerical results show reduced clustering loss with modifications.

Abstract

The K-means algorithm is one of the most widely studied clustering algorithms in machine learning. While extensive research has focused on its ability to achieve a globally optimal solution, there still lacks a rigorous analysis of its local optimality guarantees. In this paper, we first present conditions under which the K-means algorithm converges to a locally optimal solution. Based on this, we propose simple modifications to the K-means algorithm which ensure local optimality in both the continuous and discrete sense, with the same computational complexity as the original K-means algorithm. As the dissimilarity measure, we consider a general Bregman divergence, which is an extension of the squared Euclidean distance often used in the K-means algorithm. Numerical experiments confirm that the K-means algorithm does not always find a locally optimal solution in practice, while our proposed methods provide improved locally optimal solutions with reduced clustering loss. Our code is available at https://github.com/lmingyi/LO-K-means.