Euclidean distance, k-means, mixed-precision algorithm, low precision, GPU computing, clustering

1 Introduction↩︎

The k-means algorithm has been one of the most studied algorithms in data mining and machine learning for decades, and it remains widely used due to its simplicity. Given a prescribed number of clusters, it seeks a partition of the data in which points are assigned to representative centers according to their similarity in a certain metric. The k-means algorithm plays a critical role in vector quantization [1], bioinfomatics [2], computer vision [3], anomaly detection [4], database management [5], and nearest neighbor search [6]. As stated in [7], with superpolynomial runtime, k-means is often very slow in practice. Many techniques have therefore been proposed to enhance the k-means algorithm regarding speed and scalability, e.g., parallelization and batch computations; associated variants include k-means++ [8], k-means\(||\) [9], mini-batch k-means [10], and parallel k-means [11].

The increasing availability of lower-precision floating-point arithmetic, extending beyond the IEEE Standard 64-bit double precision (FP64) and 32-bit single precision (FP32) formats [12], has motivated extensive research on low-precision number formats and arithmetic operations. In particular, 16-bit formats such as FP16 and BF16 [12] are now widely supported in modern hardware and software environments; see Table 1 for the parameters of the floating-point formats used in this paper. By providing higher arithmetic throughput, reducing data movement, lowering energy consumption, and decreasing storage requirements, low-precision arithmetic has become an important tool for improving computational efficiency. It has been successfully exploited in a wide range of applications, including numerical linear algebra [13], [14], unsupervised learning [15], deep neural networks [16], [17], and other computational science and machine learning tasks. However, the use of low precision also introduces larger rounding errors, and rounding error analysis is therefore essential for identifying which parts of an algorithm can safely be performed in low precision [18]. Such analysis helps determine where mixed precision can be used judiciously, how computational resources can be exploited most effectively, and which precisions are needed to preserve the quality of the computed outputs.

Despite the prevalence of mixed-precision computing, it appears that its potential has largely remained unexploited in Euclidean distance computation and k-means clustering. In this work, we develop a mixed-precision variant of Lloyd’s algorithm for \(k\)-means clustering and provide CUDA implementations of several algorithmic variants. Guided by rounding error analysis, we propose a mixed-precision approach to Euclidean distance computation. The resulting distance kernel performs most computations in low precision and invokes higher-precision correction only when indicated by an inexpensive reliability test, thereby combining the efficiency of low-precision arithmetic with improved numerical stability. To the best of our knowledge, this work is the first to systematically analyze the numerical stability of two common formulas for computing Euclidean distances, their evaluation in mixed precision, and their incorporation into a mixed-precision Lloyd’s k-means algorithm. The resulting mixed-precision k-means has already been integrated into the symbolic time series aggregation library fABBA [19] that supports large-scale time series and signal processing, as well as downstream applications in large language models [20].

The paper is organized as follows. In Section 2 we review some recent work on k-means clustering with extra computing resources such as GPU and parallel computing. Lloyd’s k-means algorithm is then presented in Section 3, where we briefly discuss two Euclidean distance computation formulas. In Section 4 we discuss the numerical stability of Lloyd’s k-means method, including the stages of distance computation and center update. In particular, we perform rounding error analyses for two distance computation formulas and derive the new mixed-precision method. The resulting mixed-precision k-means method is presented in Section 5. Numerical experiments on variants of the mixed-precision distance computation method and k-means clustering are presented in Section 6. Conclusions are drawn in Section 7.

Throughout the paper, inequalities between vectors hold componentwise, and the norm \(\|\cdot\|\) denotes the \(\ell^2\) norm \(\|x\|=(\sum_i|x_i|^2)^{1/2}=\sqrt{x^Tx}\). We use \(|S|\) to represent the cardinality of a set \(S\).

2 Related work↩︎

In this section, we briefly review several works on accelerating k-means clustering that are closely related to our work, with particular emphasis on approaches that improve performance by exploiting additional hardware or parallel computing resources. The k-means algorithm as well as the trending seeding of \(D^2\) weighting [8] are known to be inherently sequential, which makes it tricky to implement the algorithm in a parallel way. Numerous novel approaches have been proposed to enhance scalability and speed. The authors of [21] proposed a distributed computing scheme for coreset construction that enables an acceptable low communication complexity and distributed clustering. Moreover, an initialization approach of \(D^2\) weighting with MapReduce is proposed in [22] to enable the use of only one job for choosing the \(k\) centers, which can significantly reduce the communication and I/O costs. A k-means algorithm based on GPUs using single instruction multiple data architectures is presented in [23]; the algorithm performs the assignment of data points as well as cluster center updates on the GPU, which offers acceleration by orders of magnitude. The GPU power is also harnessed in [24] to update centroids efficiently; the k-means algorithm therein also executes a single-pass strategy to remove data transfers—up to 2 and 18.5 times higher throughput is achieved compared to multi-pass and cross-processing strategies, respectively. More recently, Li et al. [25] proposed a batched GPU-based k-means algorithm, which achieves substantially faster performance than a standard GPU-based implementation while generating outputs of comparable quality.

3 The k-means method↩︎

As a classical problem in machine learning and computational geometry, k-means clustering seeks a codebook of \(k\) vectors, i.e., \(C=[c_1, \ldots, c_k] \in \mathbb{R}^{r \times k}\) in \(n\) vectors \(P=[p_1, \ldots, p_n] \in \mathbb{R}^{r\times n}\), such that \(k \ll n\), where each \(c_i\) is associated with a unique cluster \(S_i \subseteq P\). Letting \(S_1, S_2, \ldots, S_k \subseteq P\) be the \(k\) computed clusters, k-means aims to minimize the sum of squared errors (SSE), given by \[\label{eq:sse} \texttt{SSE}= \sum_{i=1}^{k} \phi (c_i, S_i),\tag{1}\] where \(\phi\) denotes some energy function and \(c_i\) denotes the center of cluster \(S_i\). In this paper, \(\phi (c_i, S_i) = \sum_{p\in S_i}\mathrm{dist}(p, c_i)^2\), where \(\mathrm{dist}(p_i, p_j)\) denotes the Euclidean distance between \(p_i\) and \(p_j\), measured by \(\|p_i - p_j\| =\sqrt{(p_i-p_j)^T(p_i-p_j)}\). In Euclidean space, Lemma 1 ensures that choosing the mean center \(\mu_i = \sum_{p \in S_i} p / |S_i|\) as \(c_i\) always leads the iterations to monotonically decrease SSE in 1 . Lloyd’s algorithm [26] (also known as the k-means algorithm) is a suboptimal solution of vector quantization to minimize SSE. For simplicity, we will henceforth denote the Euclidean norm as \(\mathrm{dist}(\cdot)\).

3.1 Lloyd’s k-means algorithm↩︎

It is well known that computing a globally optimal k-means clustering is NP-hard. In practice, one therefore usually resorts to heuristic algorithms, the most classical of which is Lloyd’s algorithm [26], often also referred to as the k-means algorithm [27]. Lloyd’s k-means algorithm (given as Algorithm [alg:kmeans]), is a local improvement heuristic: starting from an initial set of cluster centers, it alternates between assigning each data point to its nearest center and updating each center as the mean of the points assigned to it. Although this procedure is simple and widely used, it may converge to a local minimum, has no general worst-case approximation guarantee [28], and can have superpolynomial complexity in the worst case.

The quality and convergence behavior of Lloyd’s algorithm depend strongly on the choice of the initial centers. A particularly successful initialization scheme is seeding by \(D^2\) weighting, introduced by Arthur and Vassilvitskii [8]. The resulting method, commonly known as k-means++, is \(O(\log k)\)-competitive in expectation with the optimal clustering and is the default choice for the k-means implementation of scikit-learn [29]—the popular, comprehensive machine learning library. A parallel alternative is k-means\(\|\) [9], which replaces the sequential \(D^2\)-sampling step in k-means++ by an oversampling-based parallel procedure followed by a reclustering step.

We do not study initialization schemes in this paper, since our objective is to study the effect of mixed-precision arithmetic on the computationally dominant steps of Lloyd’s iteration, in particular, the Euclidean distance computations. We therefore use a simple randomized initialization, selecting \(k\) initial centers uniformly from the data, and apply the same initialization across the methods being compared. More sophisticated seeding strategies, such as k-means++ or k-means\(\|\), are orthogonal to the mixed-precision techniques developed here and could be combined with them in a complete implementation.

3.2 Distance computation↩︎

The computationally dominant step in the k-means algorithm, even including cluster initialization, is the computation of pairwise distances. Two natural approaches to forming the squared Euclidean distance \(d(x, y)^2\) between two points \(x,y\in\mathbb{R}^r\) are \[\label{eq:dist-direct} d(x, y)^2 = (x-y)^T (x-y),\tag{2}\] and \[\label{eq:dist-expanded} d(x, y)^2 = x^Tx - 2x^T y + y^T y.\tag{3}\] The second scheme 3 is of particular computational interest and is the prevailing choice in algorithmic implementations, for example, in the k-means algorithm of the machine learning library [29]. In PyTorch [30], torch.cdist⁵ by default uses 2 to compute pairwise Euclidean distances for small inputs but switches to a matrix-multiplication-based formulation related to 3 for larger inputs; this behavior can be controlled explicitly through the compute_mode argument.

If the distance formula 2 is used, it requires about \(kn\) distance computations in each iteration of executing Lines [alg:k-means:line:forloop-start]–[alg:k-means:line:forloop-end] of Algorithm [alg:kmeans] since each of the \(n\) points \(p_i\) must be compared with each of the \(k\) centers \(c_j\). Each distance evaluation costs about \(3r\) floating-point operations (flops), and so the distance computation in one iteration costs about \(3r kn\) flops, summing to an overall cost of \(3Tr kn\) flops for \(T\) iterations.

In contrast, the expanded formula 3 allows the \(n\) inner products \(p_i^Tp_i\), \(i=1\colon n\), to be precomputed and stored. This costs approximately \(2r n\) flops and these values can then be reused across all iterations. The center norms \(c_j^Tc_j\) needs be recomputed whenever the centers are updated, costing approximately \(2rk\) flops per iteration. The dominant cost is the is the computation of all inner products \(p_i^Tc_j\), which can be performed at a cost of \(2rkn\) flops per iteration. Forming the final distance matrix from these terms costs only \(O(kn)\) additional flops. Therefore, the overall cost of distance computing via 3 in a total of \(T\) iterations is about \(2rn + T(2rk+2rkn+3kn)\) flops. The leading term is approximately \(2Tr kn\) flops if \(n\gg 1\) and \(r\gg 1\), which is the case in practical applications.

More importantly, the expanded formula 3 avoids duplicated computations by precomputing the squared norm values and enables formulation of the main computation (i.e., the inner products between data points and centres) as a level-3 basic linear algebra subprograms (BLAS) matrix–matrix multiplication [31]. Let \[X = \begin{bmatrix} x_1 & x_2 & \ldots & x_m \end{bmatrix}^T \in \mathbb{R}^{m\times r}, \qquad Y = \begin{bmatrix} y_1 & y_2 & \ldots & y_n \end{bmatrix}^T \in \mathbb{R}^{n\times r},\] where \(x_i,y_j\in\mathbb{R}^{r}\). Then the squared pairwise Euclidean distance matrix \(D\in\mathbb{R}^{m\times n}\), defined by \[D_{ij} = \|x_i-y_j\|_2^2 = x_i^T x_i - 2x_i^T y_j + y_j^T y_j, \qquad i=1,\ldots,m,\quad j=1,\ldots,n,\] can be written as \[\label{eq:dist-blas3} D = p\mathbf{1}_n^T - 2XY^T + \mathbf{1}_m q^T,\tag{4}\] where \(p_i=x_i^T x_i\), \(q_j=y_j^T y_j\), and \(\mathbf{1}_m\in\mathbb{R}^{m}\), \(\mathbf{1}_n\in\mathbb{R}^{n}\) denote the vectors of all-ones. This formulation evaluates the squared norms of \(p_i\) and \(q_j\) only once and reuses the values across all corresponding pairs. The computationally dominant term \(-2XY^T\) can be evaluated by a GEMM, \[D \leftarrow \alpha AB+\beta C,\] with \(A=X, B=Y^T, C=p \mathbf{1}_n^T + \mathbf{1}_m q^T, \alpha=-2\), and \(\beta=1\), so the formulation 4 can give substantially better performance on modern hardware than implementing 2 via level-2 BLAS.

To give a brief illustration of the efficiency of the two distance formulas, we computed distances between two feature matrices of dimension \(128 \times 5,000\) (so \(r=128\) and \(n=5,000\)), which requires a total of 25 million distance entries. Table 2 presents the results, where we have tested two different data-distribution settings: randomly generated data and nearly located data; see Section 6 for full details on the experiment environment. The expanded formula 3 is clearly much faster than the direct formula 2 in both FP32 and FP64.

4 Numerical stability of the k-means method↩︎

In this work, numerical stability is understood in the standard sense of backward and forward error analysis [18]. To study the numerical stability of the k-means method, we need to look at its two main computational steps: distance computation and center update. We will use the standard model of floating-point arithmetic introduced in [18] for our stability analysis: \[\label{eq:fpmodel} \mathop{\mathrm{\operatorname{f\kern.2ptl}}}(x \texttt{op} y) = (x \texttt{op} y) (1 + \delta), \quad |\delta| \le u,\tag{5}\] where \(u\) denotes the unit roundoff associated with the floating-point number system, \(x\) and \(y\) are floating-point numbers, and op denotes addition, subtraction, multiplication, or division. For inner products of two vectors, we have [18], \[\label{eq:fpmode-vecprod} \mathop{\mathrm{\operatorname{f\kern.2ptl}}}(x^Ty) = x^Ty + s,\quad |s|\le \gamma_r|x|^T|y|,\tag{6}\] where \(\gamma_r:= r u/(1-r u)\) and \(r u<1\). In addition to the working precision, with unit roundoff \(u\), we consider two floating-point arithmetics with unit roundoffs \(u_{\ell}\) and \(u_{h}\) that satisfy \[0 < u_{h}\le u < u_{\ell}< 1.\] We refer to these two arithmetics as the low and high precisions, respectively. By “computing in precision \(u_*\)”, we mean that \(u_*\) is the unit roundoff of the current precision of computation, where \(u_*=u_{h}\), \(u\), or \(u_{\ell}\). The operator \(\mathop{\mathrm{\operatorname{f\kern.2ptl}}}(\cdot)\) with a subscript denotes the computation in the associated precision; when the subscript is omitted, the computation is understood to be performed in the working precision. Quantities computed in floating-point arithmetic wear a hat.

The process of computing vector inner products is backward stable [18] but the forward accuracy depends also on the conditioning. Consider the inner product \(f(x):=x^Ty\), where \(y\) is a fixed vector. This function has the normwise relative condition number [32] (in the \(2\)-norm) \[\mathop{\mathrm{cond}}(f,x) := \lim_{\varepsilon\to 0} \sup_{\|\Delta x\|\le \varepsilon \|x\|} \frac{|f(x+\Delta x) - f(x)|}{\varepsilon |f(x)|},\] which is given explicitly by \[\label{eq:x94Ty-cond} \mathop{\mathrm{cond}}(f,x) = \frac{\|\nabla f \|\|x\|}{|f(x)|} = \frac{\|y\|\|x\|}{|x^Ty|} = \frac{1}{|\cos\omega|},\tag{7}\] where \(\omega\in[0,\pi]\) is the angle between \(x\) and \(y\). This shows that the vector dot product is more sensitive (has larger condition number) if the two vectors are close to being orthogonal. Therefore, we can expect that the relative forward error in \(\mathop{\mathrm{\operatorname{f\kern.2ptl}}}(x^Ty)\) will be large when \(x\) and \(y\) are close to being orthogonal. On the other hand, if \(y=x\), then \(\mathop{\mathrm{cond}}(f,x) = 1\) and \(|\mathop{\mathrm{\operatorname{f\kern.2ptl}}}(x^Tx)-x^Tx|\le \gamma_r|x|^T|x|=\gamma_r x^Tx\), which shows high relative accuracy is guaranteed.

4.1 Distance formulas↩︎

We begin with the direct distance formula 2 . Its proof is a standard textbook exercise based on [18]; a proof is also given in [33].

Theorem 1 shows that the formula 2 computes the distance to high relative forward accuracy in floating-point arithmetic. But the same cannot be guaranteed for the other formula 3 , as shown by the following result.

The bound from Theorem 2 implies that the distance formula 3 is not always stable. Indeed, for fixed \(y\), consider the distance function \[g(x) = d(x,y)^2 = x^Tx - 2x^Ty + y^Ty.\] The gradient is given by \[\nabla g(x) = 2(x-y),\] and hence the normwise relative condition number is \[\mathop{\mathrm{cond}}(g,x) = \frac{\|\nabla g \|\|x\|}{|g(x)|} = \frac{2\|x-y\|\|x\|}{\|x-y\|^2} = \frac{2\|x\|}{\|x-y\|}.\] Thus the stability of the distance formula 3 is not always guaranteed, especially when \(x\approx y\).

4.2 Distance computation with reliability test↩︎

Our discussion in Section 4.1 shows that the commonly used distance formula 3 can be unreliable when the two vectors are close, due to severe cancellation. This potential instability can be particularly problematic in low-precision arithmetic, where rounding errors are larger and can accumulate more significantly. Motivated by the error analysis, we design a numerical safeguard that compares the computed distance with an easily computable error floor, and, when necessary, switches to the distance formula 2 , which is less efficient but more reliable in the cancellation-prone regime.

The quantity \[\label{eq:error-computable} E(x, y) := 2\gamma_{r + 2} \bigl( \|x\|^2 + \|y\|^2 \bigr)\ge 0\tag{8}\] in ?? easily computable and so provides a convenient way of checking the absolute error in the computed squared distance \(\widehat{d}_{r,\mathsf{exp}}\). If \(\widehat{d}_{r,\mathsf{exp}} \ge \rho E(x, y)\) for some \(\rho> 1\), we have \[d(x,y)^2 \ge \widehat{d}_{r,\mathsf{exp}} - |\widehat{d}_{r,\mathsf{exp}} - d(x,y)^2| \ge \widehat{d}_{r,\mathsf{exp}} - E(x,y) \ge (\rho-1)E(x,y).\] It follows that \[\label{eq:expanded-rel-bound} \frac{| \widehat{d}_{r,\mathsf{exp}} - d(x,y)^2|}{|d(x,y)^2|} \le \frac{E(x, y)}{(\rho-1)E(x,y)} = \frac{1}{\rho-1}.\tag{9}\] The parameter \(\rho>1\) can be viewed as a safety factor that controls the relative error in \(\widehat{d}_{r,\mathsf{exp}}\).

Therefore, we propose to exploit the more efficient distance formula 3 with the fallback rule \[\label{eq:fallback-rule} \widehat d(x,y)^2 = \begin{cases} \widehat{d}_{r,\mathsf{exp}}, & \text{if}\; \widehat{d}_{r,\mathsf{exp}} > \rho E(x, y),\\ \widehat{d}_{r,\mathsf{dir}}, & \text{otherwise}, \end{cases}\tag{10}\] where \(\widehat d(x,y)^2\) presents the value we finally take as the computed squared distance between \(x\) and \(y\). When the reliability test \(\widehat{d}_{r,\mathsf{exp}} > \rho E(x, y)\) fails (so there is no control over the relative error in \(\widehat{d}_{r,\mathsf{exp}}\), signifying possible occurrence of significant cancellation), the fallback rule distrusts the computed distance and switches to computing the squared distance via the less efficient but numerically stable formula 2 .

4.3 Effect of input quantization↩︎

Thus far, we have been considering the case where the arithmetic of computation is uniform with precision \(u\). In our mixed-precision setting of interest, the original data will be converted into a lower precision so that the subsequent pairwise distance computation is performed in the low precision \(u_{\ell}> u\). In this case we have to quantify an additional source of error from the input quantization, i.e., converting the high-precision input to the distance formulas into lower precision.

We denote \(\widetilde{x}:= \mathop{\mathrm{\operatorname{f\kern.2ptl}_{\ell}}}(x)\) and \(\widetilde{y}:= \mathop{\mathrm{\operatorname{f\kern.2ptl}_{\ell}}}(y)\) as the low-precision data points converted from \(x\) and \(y\), respectively, so the differences \(\Delta x:= \widetilde{x}- x\) and \(\Delta y:= \widetilde{y}- y\) satisfy the componentwise bounds \[\label{eq:quantis-compon-bound} |\Delta x| \le u_{\ell}|x|, \quad |\Delta y| \le u_{\ell}|y|.\tag{11}\] Now define \[e :=(\widetilde{x}-x) - (\widetilde{y}-y) = \Delta x- \Delta y.\] Then \[\widetilde{x}- \widetilde{y}= (x - y) + e, \quad |e| \le u_{\ell}(|x| + |y|).\] We have \(d(\widetilde{x}, \widetilde{y})^2 - d(x, y)^2 = 2(x-y)^T e+e^T e\), and so \[|d(\widetilde{x}, \widetilde{y})^2 - d(x, y)^2 | \le C(x,y),\] where \(C(x,y)\) is the conversion error bound defined by \[\label{eq:quantis-convert-error} C(x,y) := 2u_{\ell}|x-y|^T (|x|+|y|) + u_{\ell}^2 (|x|+|y|)^T(|x|+|y|).\tag{12}\]

Invoking Theorem 1, the computed distance (in precision \(u_{\ell}\)) between the quantized data via formula 2 , which we denote as \(\widehat{d}^{\ell}_{r,\mathsf{dif}}\), satisfies \[|\widehat{d}^{\ell}_{r,\mathsf{dif}} - d(\widetilde{x}, \widetilde{y})^2| \le \gamma^{\ell}_{r+2} d(\widetilde{x}, \widetilde{y})^2.\] Using the triangular inequality \(|\widehat{d}^{\ell}_{r,\mathsf{dif}} - d(x, y)^2|\le |\widehat{d}^{\ell}_{r,\mathsf{dif}} - d(\widetilde{x}, \widetilde{y})^2| + |d(\widetilde{x}, \widetilde{y})^2 - d(x, y)^2|\), the relative discrepency to the original distance can be bounded above by \[\frac{|\widehat{d}^{\ell}_{r,\mathsf{dif}} - d(x, y)^2|}{d(x, y)^2} \le (1 + \gamma^{\ell}_{r+2})\frac{C(x,y)}{d(x, y)^2} + \gamma^{\ell}_{r+2}.\] Define \[\eta:= \frac{u_{\ell}(\|x\| + \|y\|)}{\|x-y\|}, \quad x\ne y.\] Since \[C(x,y) \le 2u_{\ell}\|x-y\| (\|x\| + \|y\|) + u_{\ell}^2 (\|x\| + \|y\|)^2 = (2\eta + \eta^2)d(x, y)^2,\] we arrive at the relative error bound \[\label{eq:quantis-rel-bound-dif} \frac{|\widehat{d}^{\ell}_{r,\mathsf{dif}} - d(x, y)^2|}{d(x, y)^2} \le (2\eta + \eta^2)(1 + \gamma^{\ell}_{r+2}) + \gamma^{\ell}_{r+2}.\tag{13}\] This implies that a necessary condition for the formula 2 to remain accurate for the original distance (so the input quantisation is harmless) is \(\eta\ll 1\), namely, \[\label{eq:quantis-necess-cond} \|x-y\| \gg u_{\ell}(\|x\| + \|y\|).\tag{14}\] However, this condition will brazenly fail when the original data points \(x\) and \(y\) lie sufficiently away from the origin but stay close enough; for example, taking \(x\) to be a unit vector and \(y=x+\delta\) with \(\|\delta\| = 10 u \ll u_{\ell}\). The conclusion is that the distance formula 2 is not suitable for use together with input quantization.

Now consider, on the other hand, the computed distance (in precision \(u_{\ell}\)) between the quantized data via formula 3 . By Theorem 2, we have \[|\widehat{d}^{\ell}_{r,\mathsf{exp}} - d(\widetilde{x}, \widetilde{y})^2 | \le \gamma^{\ell}_{r+2} (\widetilde{x}^T \widetilde{x}+ 2|\widetilde{x}|^T |\widetilde{y}| + \widetilde{y}^T \widetilde{y}).\] Define \[E^{\ell}(\widetilde{x},\widetilde{y}) := 2\gamma^{\ell}_{r+2} (\|\widetilde{x}\|^2 + \|\widetilde{y}\|^2) \ge |\widehat{d}^{\ell}_{r,\mathsf{exp}} - d(\widetilde{x}, \widetilde{y})^2|,\] corresponding to the computable 8 in the uniform precision. We can obtain from 11 that \(\|x\|^2 + \|y\|^2 \le (1+u_{\ell})^2 (\|\widetilde{x}\|^2 + \|\widetilde{y}\|^2)\), and so \[\label{eq:quantis-error-computable} E^{\ell}(\widetilde{x},\widetilde{y}) \ge \frac{2\gamma^{\ell}_{r+2}}{(1+u_{\ell})^2} (\|x\|^2 + \|y\|^2).\tag{15}\]

Suppose \[\label{eq:quantis-reliability-test} \widehat{d}^{\ell}_{r,\mathsf{exp}} > \widetilde{\rho} E^{\ell}(\widetilde{x},\widetilde{y}), \quad \widetilde{\rho}>2.\tag{16}\] Similarly to the derivation of 9 in Section 4.2, we can obtain \[d(\widetilde{x},\widetilde{y})^2 \ge (\widetilde{\rho}-1)E^{\ell}(\widetilde{x},\widetilde{y})\] and \[\label{eq:quantis-expanded-rel-bound} \frac{| \widehat{d}^{\ell}_{r,\mathsf{exp}} - d(\widetilde{x},\widetilde{y})^2|}{|d(\widetilde{x},\widetilde{y})^2|} \le \frac{1}{\widetilde{\rho}-1}.\tag{17}\] Using \(\||x-y|\|\le \||x| + |y|\| \le \|x\| + \|y\|\) and then the bound 15 , we now bound the conversion error 12 by \[C(x,y) \le 2(2u_{\ell}+ u_{\ell}^2)(\|x\|^2 + \|y\|^2) \le \frac{(2u_{\ell}+ u_{\ell}^2)(1+u_{\ell})^2}{\gamma^{\ell}_{r+2}}E^{\ell}(\widetilde{x},\widetilde{y}).\] Since \(\gamma^{\ell}_{r+2} \ge (r+2)u_{\ell}\) and \(r>1\), this bound can be weakened to \[\label{eq:quantis-convert-error-bound} C(x,y) \le \frac{(2 + u_{\ell})(1+u_{\ell})^2}{r+2}E^{\ell}(\widetilde{x},\widetilde{y}) \le E^{\ell}(\widetilde{x},\widetilde{y}),\tag{18}\] which implies \[d(x, y)^2\ge d(\widetilde{x}, \widetilde{y})^2 - C(x,y) \ge (\widetilde{\rho}-2)E^{\ell}(\widetilde{x},\widetilde{y}).\] Finally, using the triangular inequality \(|\widehat{d}^{\ell}_{r,\mathsf{exp}} - d(x, y)^2| \le |\widehat{d}^{\ell}_{r,\mathsf{exp}} - d(\widetilde{x}, \widetilde{y})^2| + |d(\widetilde{x}, \widetilde{y})^2 - d(x, y)^2|\), we can quantify the relative error of the computed distance in low precision with respect to the original distance: \[\label{eq:quantis-rel-bound-expand} \frac{|\widehat{d}^{\ell}_{r,\mathsf{exp}} - d(x, y)^2|}{d(x, y)^2} \le \frac{E^{\ell}(\widetilde{x},\widetilde{y}) + C(x,y)}{d(x, y)^2} \le \frac{2}{\widetilde{\rho}-2}.\tag{19}\]

The bound 19 implies that the formula 3 can be safely used in combination with input quantization to compute the distance within a relative accuracy of \(2/(\widetilde{\rho}-2)\), given that the reliability test 16 holds with safety factor \(\widetilde{\rho}>2\). In contrast, satisfying the same bound 15 does not guarantee a small relative error for the formula 2 , since the bound does not prevent \(\|x-y\|\) from being arbitrarily small. Therefore, if the computed distance between \(\widetilde{x}\) and \(\widetilde{y}\) via 3 fails the reliability test, then the fallback to the formula 2 should not reuse the low-precision-converted operands. It should recompute \(x-y\) from the original data, or at least from a precision high enough that 14 is satisfied, so the quantisation itself does not destroy true difference.

The discussion motivates the mixed-precision fallback rule \[\label{eq:fallback-rule-mp} \widehat d(\widetilde{x}, \widetilde{y})^2 = \begin{cases} \widehat{d}^{\ell}_{r,\mathsf{exp}}\; \text{in precision}\; u_{\ell}, & \text{if}\; \widehat{d}^{\ell}_{r,\mathsf{exp}} > \widetilde{\rho} E^{\ell}(\widetilde{x},\widetilde{y}),\\ \widehat{d}_{r,\mathsf{dir}}\; \text{in precision}\; u, & \text{otherwise}, \end{cases}\tag{20}\] where \(\widehat d(\widetilde{x}, \widetilde{y})^2\) denotes the value we finally take as the computed squared distance between \(x\) and \(y\), with input quantization \(\widetilde{x}= \mathop{\mathrm{\operatorname{f\kern.2ptl}_{\ell}}}(x)\) and \(\widetilde{y}= \mathop{\mathrm{\operatorname{f\kern.2ptl}_{\ell}}}(y)\). The resulting mixed-precision distance computing scheme is given as Algorithm [alg:mp-dist].

4.4 Cluster center update↩︎

After the selection of the \(k\) initial centers, the k-means method (Algorithm [alg:kmeans]) alternates between assigning each data point to its closest center and updating each center by the mean of the points assigned to it. For a nonempty cluster \(S_i\), with \(|S_i| = m_i\), the exact center update and a componentwise error bound of it are \[\label{eq:center} \mu_i=\frac{1}{m_i}\sum_{p\in S_i}p,\quad \bar a_i: = \frac{1}{m_i}\sum_{p\in S_i} |p| \ge \mu_i.\tag{21}\] The update is structurally simple, but its accuracy is governed by the accuracy of the coordinate summations used to form the cluster mean.

The result on center perturbation can be translated into one that concerns perturbation of the cluster energy.

The bound ?? tells us that the increase in the cluster energy caused by using the rounded center \(\widehat\mu_i\) depends on both the within-cluster variation and the squared norm of the exact center \(\mu_i\).

As mentioned in the previous section, the convergence of the k-means algorithm relies on its property as a local improvement heuristic, that is, reassigning the data points to their closest center and then updating the cluster centers by the mean can only decrease the SSEof 1 . We now give a descent-preservation result for one center update on a fixed cluster. The local statement shows when the rounded center update still decreases the cluster contribution to the SSE.

If the center is computed by the recursive summation of Lemma 2, such that \(|\widehat\mu - \mu| \le \gamma_{m} \bar a\) for some \(\bar a\), then Theorem 4 implies \[\|\widehat c-\widehat\mu\|^2 > 2\gamma_m |\widehat c-\widehat\mu|^T\bar a.\] Furthermore, if \(|\widehat c-\widehat\mu|^T\bar a \ne 0\), then a stronger sufficient condition is \(mu/ (1-mu) < \|\widehat c-\widehat\mu\|^2/ (2|\widehat c-\widehat\mu|^T\bar a)\), or equivalently by \[u < \frac{\|\widehat c-\widehat\mu\|^2}{m(\|\widehat c-\widehat\mu\|^2+2|\widehat c-\widehat\mu|^T\bar a)}.\]

Theorem 4 is a statement for a fixed cluster. Summing the same argument over all clusters gives a corresponding sufficient condition for descent of the total SSEafter one center-update step, given as Corollary 2.

The vector \(|\widehat c-\widehat\mu|\) measures the center movement, which is expected to decrease as the iterations proceed. Therefore, the update may require either higher precision or a more accurate summation scheme to preserve the descent property at the level of the computed centers. For the latter, one possibility is to employ pairwise summation [18], and then the factor \(\gamma_m\) in the bounds above can be replaced by a factor depending on the summation depth, typically of order \(\gamma_{\lceil \log_2 m\rceil}\).

Finally, we note that the analysis is not independent of preprocessing, which changes the data distribution. Indeed, the error bounds for the center update depend on quantities such as \(\bar a_i = \sum_{p\in S_i}|p|/m_i\) and \(\sum_{p\in S_i}\|p\|^2\), which may change when the data are rescaled, centered, or normalized.

5 Mixed-precision k-means↩︎

Based on our discussion in the previous section, we present our mixed-precision k-means algorithm in Algorithm [alg:mp-kmeans].

The distance computations used for cluster assignment are performed in mixed precision using Algorithm [alg:mp-dist]. In contrast, the centroid updates are carried out in the working precision throughout. This choice is motivated by the analysis in Section 4.4, which shows that sufficient precision in the centroid update is important for preserving convergence, especially in the later stages of the k-means iteration.

Reduced-precision center updates may be viable early in the iteration, but we do not pursue this option here because a reliable switching criterion is not yet available.

6 Numerical experiments↩︎

All tests were performed using four independent runs. The first run was used as a warm-up and excluded from timing statistics; we report averages over the remaining three runs. The compute environment is a single compute node⁶ equipped with a Dell PowerEdge R750xa server, featuring Intel Xeon Gold 6330 CPUs (56 cores/112 threads at 2.00 GHz), 1 TB of RAM, and four NVIDIA A100 80GB PCIe GPUs. We utilized a single GPU for computing. The software environment is based on CUDA 12.8 and PyTorch 2.11.0. Before running the algorithms, we standardize each feature using z-score normalization, i.e., by subtracting its mean and dividing by its standard deviation [34].

6.1 Euclidean distance computation↩︎

First, we compare the Euclidean distances computed by 3 in a working precision of BF16, FP16, FP32, and FP64 (see Table 1), respectively. The results for direct formula 2 in uniform FP32 and FP64 (without fallback) are also presented for comparison, and the latter is used as the reference value for accuracy. For the uniform FP32 and FP64 kernels (without fallback) by 3 , we use the native pairwise distance interface cdist in PyTorch. PyTorch’s cdist operator is implemented through ATen, PyTorch’s C++ tensor library, and dispatches to the corresponding CUDA kernel when the input tensors reside on a CUDA device. The uniform BF16 and FP16 kernels (without fallback) are based on our own CUDA implementations.

We test the fallback rules in these four uniform-precision environments described by 10 , and in the input-quantization setting described by 20 , with FP32 or FP64 used as the working precision. For the latter, we report in the experiments the fallback ratio \[\label{eq:eta} \eta = \frac{\text{Number of computations that fell back to high precision}}{\text{Total number of computations}}\tag{22}\] to gauge the proportion of high-precision computations. All kernels with fallback are based on our own CUDA implementations, including the input-quantization variants of Algorithm [alg:mp-dist]. Unless otherwise specified, the safety factors in the fallback rules 10 and 20 are set to \(\rho = \widetilde{\rho}= 5\).

6.1.1 Distance on far and near points↩︎

The first experiment is a comprehensive benchmarking on the performance and accuracy of various precision strategies for pairwise Euclidean distance computation on large-scale feature matrices. We use two feature matrices of size \(128\times 5,000\), yielding 25 million distance entries. The dimension 128 of data points is commonly encountered in modern machine learning and data science, particularly in embedding-based tasks such as face recognition, image retrieval, and contrastive learning; for example, Google’s FaceNet represents faces using 128-dimensional embeddings [35].

To assess robustness under different numerical conditions, we consider two scenarios. The random-points scenario represents typical machine-learning workloads with independently distributed feature vectors, whereas the near-points scenario models a more challenging setting in which many point pairs are nearly identical, differing only by noise of order \(10^{-6}\). The latter is relevant to duplicate detection, fine-grained similarity search, and clustering with tight clusters, where cancellation in floating-point subtraction can significantly degrade accuracy [17], [36], [37].

The results are reported in Table [tab:per95random95near]. For each method, we report the execution time, speedup relative to the CPU baseline by scikit-learn⁷, maximum relative error, and fallback percentage whenever it is applicable. In the first section, the four uniform-precision methods without fallback are delivering excellent accuracy to the respective unit roundoff when the data points are randomly distributed, but they become unstable in the near-points setting and produce maximum relative error much greater than one. This phenomenon is consistent with our analyses in Section 4.1. The middle section reports three uniform-precision methods with fallback (to 2 in the same precision). We see that the half precisions are not accurate enough for the fallback mechanism to work in the near-point setting, and that FP32 with fallback works in both random-points and near-points scenarios, albeit with deteriorate accuracy in the latter. In the last section, we observe that the mixed-precision methods with fallback offer satisfying accuracy in both settings. They automatically detect pairwise distance computation in which low-precision arithmetic may be unreliable and selectively apply high-precision correction, curing instability without sacrificing overall performance.

In terms of runtime, our GPU implementations achieve substantial speedups over the CPU implementation in scikit-learn, reaching up to several hundredfold acceleration when low-precision computations are exploited to their full potential. In particular, our mixed-precision FP16\(\rightarrow\)FP32 and FP16\(\rightarrow\)FP64 kernels are about \(2.7\times\) faster than the cdist routine of PyTorch (FP32 Expanded in the table) and are more stable in the near-points scenario. A caveat is that these speedups may come at the cost of reduced accuracy, with errors typically limited to the level of the unit roundoff of the low-precision format. Finally, we note that the fallback mechanism has to be used sparingly for efficiency. In the cases of BF16\(\rightarrow\)BF16, BF16\(\rightarrow\)FP32, and BF16\(\rightarrow\)FP64, the fallback criterion was triggered for all pairwise distance computations, so the direct formula 2 was used throughout. This formula consists primarily of level-\(1\) BLAS operations and has low arithmetic intensity; it is therefore largely memory-bound. Consequently, replacing FP32 or FP64 arithmetic by BF16 arithmetic does not necessarily improve the runtime. If the input data are stored in FP32 and converted to BF16 inside the kernel, the global-memory traffic is essentially unchanged. Even with BF16 storage, the runtime may still be dominated by memory access and data-movement costs rather than by floating-point throughput.

6.1.2 Sensitivity tests on fallback↩︎

We next study the sensitivity of the fallback mechanism in mixed near-far regimes. The tests are performed again on real feature matrices \(C\) and \(P\) of size \(128\times 5,000\). We construct hybrid datasets where various fraction of points in \(C = P + \varepsilon \mathcal{N}(0, I)\) are near-duplicates of points in \(P\), with noise \(\varepsilon\), while the remaining points are sampled randomly. This setting tests robustness under realistic data distributions containing both clustered structures and outliers, as commonly arise in iterative clustering algorithms.

We measure the GPU kernel execution time, the maximum relative error with respect to an FP64 reference, the fallback rate \(\eta\), and separate maximum relative errors for near (distance less than \(10^{-3}\)) and far pairs. We tested different noise levels \(\varepsilon \in \{10^{-2}, 10^{-3}, \ldots, 10^{-6}\}\) to simulate varying degrees of point proximity and found similar results in these cases. Table 3 reports the results with noise \(\varepsilon = 10^{-4}\). We first observe that uniform low-precision methods are highly sensitive to the presence of near-point pairs. When only \(10\%\) of the data are near points, the maximum relative errors of uniform FP16 and BF16 explode to much greater than one, respectively. Uniform FP32 also becomes unreliable for these near pairs, with errors above one. The far-point errors remain small, indicating that the instability is caused mainly by cancellation in the nearly coincident distance computations. This behavior causes issues for k-means clustering, where each iteration produces a mixture of converged (near) and unconverged (far) centroid-point pairs.

The mixed-precision methods with fallback substantially improve robustness. For FP16\(\rightarrow\)FP32 and FP16\(\rightarrow\)FP64, the fallback rate remains very small, at most \(0.02\%\), while the near-point errors stay at the level of about \(10^{-7}\) to \(10^{-6}\). The overall maximum error is then dominated by the far-point FP16 computation and remains around \(10^{-3}\). In contrast, BF16-based fallback is triggered for essentially all entries in this experiment. This gives stable accuracy when the correction precision is FP32 or FP64, but at a much higher runtime cost. The execution times follow a trend similar as that in the previous experiments. The mixed-precision methods FP16\(\rightarrow\)FP32 and FP16\(\rightarrow\)FP64 consistently achieve approximately \(2.7\times\) speedup than the cdist routine of PyTorch.

These results suggest that the low-precision distance computation with fallback mechanism remains effective even in the presence of numerically delicate distance entries. In such cases, the FP16\(\rightarrow\)FP32 and FP16\(\rightarrow\)FP64 kernels preserve most of the low-precision throughput while correcting the unstable near-point computations. For clustering applications, this is particularly relevant near convergence, where many centroid–point distances may become small and the expanded distance formula becomes more vulnerable to cancellation.

6.1.2.1 Implications for clustering algorithms

These results provide practical guidance for precision selection in k-means and related algorithms. In early iterations, when centroid-point distances are typically large, uniform FP16 or BF16 computations may be sufficient and can provide maximum throughput. Near convergence, however, shrinking centroid-point distances make the expanded distance formula more susceptible to cancellation. In this regime, mixed-precision methods with fallback, such as FP16\(\rightarrow\)FP32 or FP16\(\rightarrow\)FP64, are preferable to prevent divergence. The adaptive fallback mechanism makes the performance-accuracy tradeoff data-dependent: high-precision correction is applied only when the low-precision computation is deemed unreliable. Thus, the computational cost is governed by the numerical difficulty of the input rather than fixed a priori. This feature is particularly useful in clustering applications, where the distribution of centroid-point distances changes during the iterations.

6.2 K-means clustering↩︎

Our k-means implementation here follows a hybrid CPU-GPU architecture where the lightweight main loop and convergence checking run on CPU, while all computationally expensive operations are GPU-accelerated—custom CUDA kernels (same as above) to handle distance computation with multi-precision support and fallback 20 , and the center updates are parallelized across all \(r \times k\) dimensions simultaneously with automatic empty cluster reinitialization. We test in working precision of FP32 the five mixed-precision distance kernels BF16\(\rightarrow\)FP32, BF16\(\rightarrow\)FP64, FP16\(\rightarrow\)FP32, FP16\(\rightarrow\)FP64, and FP32\(\rightarrow\)FP64 with adaptive fallback mechanism, which achieved the best performance in our experiment of Section 6.1. For comparison, we also include standard uniform-precision kernels using FP16, BF16, FP32, and FP64 (without fallback).

The quality of clustering is evaluated in terms of SSEin 1 , Adjusted Rand Index (ARI), and Adjusted Mutual Information (AMI) [38]. The convergence tolerance of Algorithm [alg:mp-kmeans] is set to \(\tau = 10^{-8}\) (unit roundoff level of FP32) with maximum number of iterations \(t_{\max} = 300\).

6.2.1 Clustering on Gaussian Blobs↩︎

We evaluate the performance of distinct distance kernels with different precision configurations in k-means clustering on synthetic Gaussian blob datasets. Tables 4, 5, 6, and 7 report end-to-end k-means performance with \(n=100{,}000\) samples, feature dimensions \(d\in\{10,128\}\), and \(k\in\{100,500\}\) clusters. The reported speedup is measured relative to our FP64 implementation in CUDA, so values below one indicate a slowdown relative to the FP64 baseline. Recall that the experiment results are the average of four independent runs, so the iteration count can be non-integer values.

The uniform low-precision kernels, FP16 and BF16, are not reliable in these experiments. Although the individual distance computations are cheaper in low precision, the resulting loss of accuracy increases the number of k-means iterations substantially. In several cases, uniform FP16 reaches the maximum number of \(300\) iterations, and BF16 also requires many more iterations than FP32 or FP64. Consequently, the total runtime of the uniform low-precision methods is often much larger than that of FP32 or FP64. In contrast, the mixed-precision variants generally recover iteration counts and clustering quality close to those of uniform FP32 and FP64. The AMI and ARI values are typically very close to the FP64 reference values, and the final SSEvalues are also comparable. Among the mixed-precision methods, the FP16\(\rightarrow\)FP32 and FP16\(\rightarrow\)FP64 variants with \(\rho=5\) give the most favorable trade-offs in the reported tests. For instance, FP16\(\rightarrow\)FP64 with \(\rho=5\) achieves speedups of \(1.30\), \(1.14\), \(1.24\), and \(1.36\) in Tables 4, 5, 6, and 7, respectively, while preserving clustering quality close to the FP64 baseline. BF16\(\rightarrow\)FP64 is highly effective in Table 4, where it reaches a speedup of \(1.91\), but it is less competitive in the higher-dimensional cases, where the fallback overhead becomes more pronounced. The uniform FP32 kernel remains a strong baseline without hyperparameter tuning; consequently, adding fallback offers little additional benefit. It converges in essentially the same number of iterations as FP64 and gives nearly identical AMI, ARI, and SSE values, while being faster than FP64 in all cases.

The choice of the threshold parameter \(\rho\) has a visible effect on performance. In the low-dimensional cases \(d=10\), changing \(\rho\) from \(5\) to \(10\) or \(50\) has only a mild impact for most mixed-precision variants. For \(d=128\), however, the larger thresholds often lead to substantially higher runtimes. This is particularly clear in Tables 6 and 7, where FP16\(\rightarrow\)FP32 and FP16\(\rightarrow\)FP64 with \(\rho=10\) or \(\rho=50\) are much slower than the corresponding runs with \(\rho=5\). Thus, although the method is not highly sensitive to \(\rho\) in all regimes, the most aggressive choice, \(\rho=5\), appears to be robust among the tested values. This choice switches to high precision only when the relative error in the pairwise distance computation is not guaranteed to be below 0.67; see 19 and 20 .

Overall, these experiments show that uniform FP16 and BF16 should be used with care in k-means clustering, since their lower arithmetic cost can be outweighed by slower or unstable convergence, and, on the other hand, uniform FP64 is preferred over uniform FP32 reserved for cases where maximum numerical precision is required and longer runtime is acceptable. Mixed-precision distance kernels with fallback provide a more reliable alternative: they retain much of the speed of low-precision arithmetic while correcting computations that would otherwise degrade the k-means iterations. In the present tests, the FP16\(\rightarrow\)FP64 and FP16\(\rightarrow\)FP32 variants with \(\rho=5\) offer the best overall balance between speed, stability, and clustering quality.

6.2.2 Image segmentation↩︎

In this section we test the k-means methods for image segmentation. Given an input image of height \(H\) and width \(W\), we treat each pixel as one sample and construct a feature matrix \(P \in \mathbb{R}^{N \times 5}\) with \(N = H W\). Each row of \(P\) concatenates color and spatial information of one pixel, namely, \((R,G,B,x,y)\), where \((R,G,B)\) are the normalized color channels and \((x,y)\) are normalized pixel coordinates. Unless stated otherwise, each feature dimension is rescaled using min-max normalization so that the color and spatial coordinates have comparable magnitudes in the distance computations.

4.2pt

We run k-means with cluster number \(k \in \{8, 32\}\) to capture both coarse and fine-grained segmentations. The k-means method iteratively alternates between assigning each pixel feature to its nearest centroid in the normalized feature space and updating centroids by averaging the features of assigned samples. After convergence, we reshape the resulting assignment vector into a \(H \times W\) label map. For each k-means method, we execute a short warm-up followed by four different executions and report the average runtime and other metrics of clustering. Table 8 summarizes the results for two images from ImageNet [39]. The method in uniform FP64 serves as the reference for fidelity comparisons, and Diff measures the permutation-invariant percentage of pixel labels that differ from this reference after optimal one-to-one alignment of cluster IDs (to account for the label-permutation symmetry of k-means).

The uniform low-precision kernels are unreliable for this task. Both BF16 and FP16 reach the maximum number of iterations in all four test cases. The degradation is especially pronounced for BF16 at \(k=32\), where the disagreement rates reach \(18.47\%\) and \(32.52\%\) for the two images. This indicates that using reduced precision alone can substantially change the segmentation when the number of clusters is increased. On the other hand, uniform FP32 closely aligns with the FP64 reference in all metrics. The mixed-precision variants show more varied behavior. This indicates that, under the present feature scaling, FP32 is already sufficient to reproduce the FP64 segmentation to high accuracy.

The mixed-precision methods exhibit different behavior depending on the base precision. The FP16-based variants improve the segmentation quality relative to uniform FP16, but they still often require many iterations. In particular, FP16\(\rightarrow\)FP32 reaches the maximum iteration count in three of the four cases, and FP16\(\rightarrow\)FP64 reaches it in two cases. Hence, although FP64 fallback reduces the disagreement rate, it does not fully remove the convergence difficulty caused by the FP16 base computation. The BF16-based variants show a more nuanced trade-off. BF16\(\rightarrow\)FP32 closely matches the FP64 reference with comparable iteration counts and runtimes to FP32 and FP64. In contrast, BF16\(\rightarrow\)FP64 is often the fastest method because it converges in substantially fewer iterations, but this speed can come at the cost of noticeably different segmentations. The most pronounced case is ILSVRC2012_val_00011129 with \(k=32\), where BF16\(\rightarrow\)FP64 is about \(2.2\times\) faster than FP64 but has a disagreement rate of \(14.26\%\). The FP32\(\rightarrow\)FP64 variant is the most conservative mixed-precision method. It essentially reproduces the FP64 reference, with disagreement rates no larger than \(0.04\%\), while achieving runtimes comparable to or slightly smaller than those of FP32 and FP64. Since uniform FP32 is already highly accurate in this experiment, the benefit of adding FP64 fallback to FP32 is modest.

Overall, our experiments show that the image segmentation by k-means is sensitive to uniform low-precision (FP16 and BF16) arithmetic, especially as the number of clusters \(k\) increases, and the k-means using the mixed-precision strategy substantially reduced the Diff for image segmentation compared to the low-precision arithmetic. Therefore, FP16 and BF16 should not be used as uniform precisions for this task, since their lower arithmetic cost is outweighed by poor convergence and larger segmentation discrepancies. The lower-precision mixed-precision variants can improve runtime in some cases, but their accuracy depends on the image and the number of clusters.

7 Conclusions↩︎

In this paper, we studied the use of mixed-precision arithmetic in Lloyd’s algorithm for \(k\)-means clustering, with a particular focus on the two main numerical components of the iteration—the Euclidean distance computation used in the assignment step and the center update. Guided by rounding error analysis, we identified when reduced-precision arithmetic can be used safely and when a higher-precision correction is needed. This led to a mixed-precision framework in which the computationally dominant distance computations are performed primarily in low precision, with an inexpensive reliability test used to detect significant numerical cancellation in the computed distance.

The analysis confirms that the commonly used expanded formula 3 is computationally attractive, but that its accuracy depends on the relative size of the distance being computed. In particular, near-point configurations may lead to significant cancellation and hence unreliable low-precision results. Our mixed-precision distance kernel addresses this issue by using the expanded formula in low precision when the error can be controlled and falling back to the direct formula 2 in higher precision otherwise. The center update was treated more conservatively—our analysis indicates that sufficient precision is needed in this step to preserve convergence behavior, especially in later iterations of Lloyd’s algorithm.

We provided CUDA implementations of the resulting mixed-precision \(k\)-means variants and compared them with uniform-precision implementations and standard library routines such as the cdist routine in PyTorch. The numerical experiments show that the proposed mixed-precision distance kernels can achieve accuracy comparable with uniform IEEE double or single computations while providing speedups of up to \(4\times\) in the tested settings. They are also robust in near-point scenarios, where uniform low-precision methods may suffer substantial accuracy loss. These results demonstrate that mixed precision can offer a favorable balance between performance and numerical reliability for the distance computations in \(k\)-means clustering. Experiments on synthetic data clustering and image segmentation further show that carefully designed reduced-precision computations do not significantly degrade the clustering quality. In many cases, the increase in SSEis small, and the resulting cluster assignments remain comparable with those obtained using uniform higher precision. These observations suggest that mixed precision is a viable strategy for accelerating k-means, provided that the use of low precision is guided by appropriate numerical safeguards, such as 20 .

The ideas and methods developed in this work are expected to benefit a broad class of Euclidean-distance-based algorithms, including, but not limited to, those used in information retrieval and retrieval-augmented generation (RAG) [6], [33], [40], data clustering [36], [41], manifold learning [42], and time-series analysis [20].

References↩︎