Computing $k$-means in mixed precision
For practitioners using k-means on modern hardware with mixed precision capabilities, this work provides guidelines on when low precision can be safely used to improve runtime and energy efficiency.
The paper studies the numerical stability of Lloyd's k-means algorithm and proposes a mixed-precision framework for k-means computation. Through simulations on clustering and image segmentation tasks, they find that normalized data tolerates low-precision distance computation better, while unnormalized data requires care to avoid overflow, demonstrating potential for acceleration.
The k-means algorithm is one of the most popular and critical techniques in data mining and machine learning, and it has achieved significant success in numerous science and engineering domains. Computing k-means to a global optimum is NP-hard in Euclidean space, yet there are a variety of efficient heuristic algorithms, such as Lloyd's algorithm, that converge to a local optimum with superpolynomial complexity in the worst case. Motivated by the emergence and prominence of mixed precision capabilities in hardware, a current trend is to develop low and mixed precision variants of algorithms in order to improve the runtime and energy consumption. In this paper we study the numerical stability of Lloyd's k-means algorithm, and, in particular, we confirm the stability of the widely used distance computation formula. We propose a mixed-precision framework for k-means computation and investigate the effects of low-precision distance computation within the framework. Through extensive simulations on various data clustering and image segmentation tasks, we verify the applicability and robustness of the mixed precision k-means method. We find that, in k-means computation, normalized data is more tolerant to the reduction of precision in the distance computation, while for unnormalized data more care is needed in the use of reduced precision, mainly to avoid overflow. Our study demonstrates the potential for the use of mixed precision distance kernels to accelerate the k-means computation and offers insights into other distance-based machine learning methods.