Flash-GMM: A Memory-Efficient Kernel for Scalable Soft Clustering
arXiv:2606.10896v1 [cs.LG] 9 Jun 2026
Gal Bloch and Ariel Gera and Matan Orbach and Ohad Eytan and Assaf Toledo IBM Research [email protected] https://github.com/IBM/Flash-GMM
1
Abstract
N
Flash-GMM
vs. SciPy
vs. TorchGMM
We present Flash-GMM, a fused Triton kernel for efficient computation of Gaussian Mixture Models (GMMs) over large-scale data in a single GPU pass. By eliminating the need to materialize the full responsibility matrix in GPU memory, Flash-GMM achieves a 20× speedup over existing implementations and enables training on datasets more than 100× larger than previously feasible on one device. To demonstrate its impact, we integrate FlashGMM into the IVF coarse quantizer for approximate nearest-neighbor (ANN) search. We show that soft GMM clustering is now a viable drop-in replacement for k-means, and that GMM responsibilities can be leveraged to assign border vectors to multiple clusters. Our approach reaches fixed recall targets with up to 1.7× fewer distance computations, or equivalently, yields +2–12 recall@10 at matched computational cost. We release the kernel as an open-source project.
10K 50K 100K 250K 500K 1M 10M 50M 100M
85 ms 261 ms 444 ms 1,032 ms 2,069 ms 3,755 ms 7,400 ms 35,510 ms 74,270 ms
766× 1,260× 1,458× 1,597× 1,571× 1,738× 1,740× 1,752× 1,782×
32× 20× 23× 19× 20× 22× OOM OOM OOM
Introduction
Gaussian Mixture Models (GMMs) are versatile probabilistic tools that have found wide application across domains ranging from computer vision to bioinformatics. GMMs fit a statistical model to the data by estimating its underlying probability density as a mixture of Gaussians. GMMs quantify the probabilities of assigning each data point to each of the mixture components. These probabilities, termed responsibilities, are typically realized in a full matrix of size N ×K, where N is the number of data points and K is the number of Gaussian components. During parameter estimation via Expectation Maximization (EM, Dempster et al., 1977), this matrix is recomputed at every iteration. The massive parallelism offered by GPUs allows significantly faster GMM estimation. However, as GPU memory is limited, materializing the full responsibility matrix becomes impractical even for
Table 1: Runtime of 30 GMM EM iterations for different data scales (N ), with K = 1024, D = 128, on an A100-80 GB GPU (OOM = out of GPU memory). Flash-GMM obtains significant speedup in comparison to TorchGMM (CSOgroup, 2023), an existing GPU kernel, and a CPU-based implementation from SciPy.
moderately sized datasets. Existing GPU implementations such as TorchGMM (CSOgroup, 2023) run out of memory beyond 10 million data points, while CPU-based solvers (Virtanen et al., 2020) are orders of magnitude too slow (see Table 1). Thus, large scale applications based on GMM have remained out of reach at production scales. To address this gap, we introduce Flash-GMM, a fused Triton (Tillet et al., 2019) kernel that performs GPU-accelerated GMM estimation without materializing the responsibility matrix in the GPU HBM memory. The design is inspired by the IOaware tiling strategy of FlashAttention (Dao et al., 2022), adapted to the EM algorithm. The resulting kernel requires only O(KD) GPU memory, where D is the data dimensionality. Because memory usage does not grow with dataset size N , Flash-GMM can handle arbitrarily large datasets. Moreover, by performing all tile-local computations in GPU registers rather than round-tripping through GPU main memory, the kernel minimizes main memory access, thus alleviating the primary latency bottleneck. Empirically, our kernel achieves a 20× speedup over existing implementations and enables training on datasets more than 100× larger than
previously feasible on one device (see Table 1). Our contribution of an efficient GMM kernel opens the door to many practical use cases. We demonstrate this by focusing on a prominent application: the IVF index coarse quantizer (Jégou et al., 2011b), commonly used for approximate nearest neighbor (ANN) search. The quantizer is typically implemented via k-means clustering (Lloyd, 1982); instead, we propose Flash-GMM as a practical and performant alternative (§4). Using Flash-GMM confers several advantages in the IVF setting. First, it enables fast IVF index construction, for data scales that are too compute-intensive for a CPU and were previously too memory-demanding for a single GPU. Second, with GMM, the estimated responsibilities (posterior probabilities of data points given Gaussian components) naturally yield soft assignments of vectors to clusters. This enables a multiple assignment scheme (§4.3), such that vectors near cluster boundaries can be associated with several clusters. This contrasts with the hard assignment of k-means, which forces such ambiguous boundary cases into a single cluster. For ANN search, this scheme directly translates to improved recall, as near-boundary vectors are no longer prematurely discarded from the search space. We demonstrate this improvement through experiments on standard search benchmarks. In summary, our contributions are: • We introduce the Flash-GMM kernel: a memory-efficient, IO-aware GPU kernel for GMM, enabling training on arbitrarily large datasets. • We demonstrate the impact of Flash-GMM for a practical IVF application. In the IVF setup, the introduced kernel enables larger data scales, and, with the multi-assignment scheme, improves search quality and cost tradeoffs. Flash-GMM with multi-assignment delivers up to 1.7× fewer distance computations at fixed recall and +2–12 recall at any matched compute budget (§4.5).
2
Gaussian Mixture Models
Formally, given a data matrix X = {x1 , x2 , . . . , xN } with xi ∈ RD , GMMs model the data as generated from a mixture of K Gaussian distributions (Bishop, 2006).1 The 1 Here we restrict ourselves to isotropic Gaussian components, i.e., each covariance matrix is constrained to the form σk2 I. This assumption ensures statistical stability while re-
probability density function of a sample xi is defined as K X p(xi | Θ) = πk N (xi | µk , σk2 I), k=1
where πk denotes the mixture weight of component k, satisfying K X πk = 1, πk ≥ 0, k=1 N (xi | µk , σk2 I) is the Gaussian distribution parameterized by mean vector µk ∈ RD and isotropic variance σk2 ∈ R, and Θ = {πk , µk , σk2 }K k=1 de-
notes all model parameters. Parameter estimation is performed by maximizing the log-likelihood of the observed dataset: ! N K X X log p(X | Θ) = log πk Nik , i=1
k=1
where Nik = N (xi | µk , σk2 I). Typically, GMM parameter estimation is carried out using the Expectation-Maximization (EM) algorithm (Dempster et al., 1977), which alternates between an Expectation (E) step and a Maximization (M) step. The E-step computes the posterior probability rik of component k and each sample xi , also referred to as the responsibility. First, the unnormalized assignment score for k is defined as zik = πk N (xi | µk , σk2 I). The responsibility rik is then obtained by normalizing these scores across all components: zik . rik = p(k | xi , Θ) = PK j=1 zij Intuitively, rik measures the degree to which sample xi belongs to cluster k, where K X 0 ≤ rik ≤ 1, rik = 1. k=1
To improve numerical stability, responsibilities are computed in log-space using the log-sum-exp trick (Bishop, 2006; Blanchard et al., 2021). In the M-step, the model parameters are updated using the responsibilities computed in the E-step. First, the effective number of samples assigned to component k is computed as N X Nk = rik , (1) i=1
and the mixture weights are set to Nk πk = . N maining computationally tractable (see Appendix A).
The updated mean vector of k is computed from the responsibilities by N 1 X µk = rik xi , (2) Nk i=1
followed by an updated variance: N 1 X σk2 = rik ∥xi − µk ∥22 . DNk
(3)
i=1
The E-step and M-step are repeated iteratively until convergence, typically determined by a sufficiently small change in the log-likelihood between successive iterations or by reaching a predefined maximum number of iterations.
3
Flash-GMM
In principle, GMM is well suited for GPU acceleration due to the massive parallelism and computational throughput offered by modern GPUs. Since the responsibility rik for each data point xi depends only on the shared model parameters {πk , µk , σk2 } and not on other data points, the E-step computations across all N vectors are fully independent and can be executed in parallel. However, efficiently mapping a CPU-oriented implementation to GPU is non-trivial, as memory efficiency can quickly become the primary bottleneck and limit the benefits of the available compute (Dao et al., 2022). Consider a moderate setup, with N = 10M, D = 128, K = 2,048: at 4 bytes precision, the responsibility matrix R ∈ RN ×K alone occupies 107 × 2,048 × 4 ≈ 80GB of memory. Storing the input X ∈ RN ×D adds roughly 5GB. Together, this exceeds the capacity of many GPUs. While the memory footprint grows linearly with N , D, and K, in practice, N dominates: datasets often contain millions of vectors, making naive GPU implementations infeasible. Secondly, consider the memory bandwidth requirements of the EM procedure described in Section 2. The E-step requires a full read of X (N D reads) and a full write of the responsibility matrix R (N K writes). The M-step then performs three additional full reads of R in (1), (2), and (3), as well as two full reads of X in (2) and (3), respectively. Overall, the number of accesses is ∼ 3N D+4N K. These repeated passes generate substantial memory traffic and significantly increase the pressure on HBM bandwidth. Thus, a straightforward GPU implementation of GMM is fundamentally constrained by memory capacity and memory bandwidth. This limits
scalability and prevents efficient utilization of the available compute power. 3.1
Flash-GMM kernel
To enable GPU-based GMM estimation at scale, we introduce a new tile-based memory-efficient kernel: Flash-GMM. The new kernel is aimed at efficiently utilizing the parallel computational capacity of a GPU, while minimizing GPU memory access, and keeping peak memory use constrained. The implementation is inspired by the work of Dao et al. (2022) and Yang et al. (2026). We divide X into contiguous tiles X1 , . . . , XT of BN rows each. For each tile Xt , we perform two steps: 1. Compute the log-likelihood log zik for the tile vectors against the K components, and the logP normalizer log Zi = log k zik via the numerically stable online log-sum-exp (Blanchard et al., 2021; Dao et al., 2022). 2. Compute the responsibilities rik (using the log Zi normalizers from above), and accumulate per-tile sufficient P statistics Nk (t) = P r , M (t) = k i rik xi , and Qk (t) = Pi ik 2 i rik ∥xi − µk ∥ . Within each step, the K components are processed sequentially in blocks of size BK , since loading all component parameters at once exceeds on-chip memory. The log normalizers log Zi accumulated in Step 1, and the per-tile accumulators of Step 2 (Nk (t), M k (t) and Qk (t)) are maintained in on-chip memory across the inner sweep. After all tiles are processed, the per-tile contributions are atomically reduced into global accumulators N, M , Q, from which the M-step recovers πk , µk , σk2 via Eqs. (1)–(3). The complete two-loop structure is given in Algorithm 1. 3.2
Memory Efficiency
Peak memory use A significant advantage of Flash-GMM is that it never materializes the N × K responsibility matrix in HBM memory, thus keeping peak memory consumption low. Instead, the kernel maintains two kinds of onchip state. First, within each tile Xt , the per-vector log-normalizers log Zi are kept on chip, and are never written back to HBM. That makes them immediately available for computing the rik responsibilities in Step 2. Then, the per-tile accumulators Nk (t), M k (t), and Qk (t) are directly computed from the responsibilities. Overall, the HBM only
Algorithm 1 F LASH -GMM (Single EM Iteration) Require: Data matrix X ∈ RN ×D in HBM, GMM parameters {πk , µk , σk2 }K k=1 in HBM. 1: Set block sizes BN (tile size, e.g. 64), BK (component block, e.g. 16). 2: Initialize accumulators N = 0 ∈ RK , M = 0 ∈ RK×D , Q = 0 ∈ RK in HBM. 3: Divide X into T = ⌈N/BN ⌉ tiles X1 , . . . , XT of size BN × D each. 4: for 1 ≤ t ≤ T do ▷ Parallel across GPU blocks 5: Load Xt from HBM to on-chip registers. Initialize log Z = −∞ ∈ RBN on chip. 6: for j = 1 to ⌈K/BK ⌉ do ▷ Step 1: Log-likelihoods and log-normalizers 7: Load {πk , µk , σk2 } for k ∈ [(j−1)BK +1, jBK ] from HBM to SRAM. 8: On chip, compute log zik = log πk N (xi | µk , σk2 I) ∈ RBN ×BK . P 9: On chip, accumulate log Zi ← log exp(log Zi ) + k exp(log zik ) ∈ RBN . 10: end for 11: for j = 1 to ⌈K/BK ⌉ do ▷ Step 2: Responsibilities and sufficient statistics (log Z remains in registers) 12: Load {µk , σk2 } for k ∈ [(j−1)BK +1, jBK ] from HBM to SRAM. 13: On chip, recompute log zik (same as line 10). ×BK 14: On chip, compute rik = exp(logP zik − log Zi ) ∈ RBNP . P 15: On chip, accumulate Nk (t) += i rik , M k (t) += i rik xi , Qk (t) += i rik ∥xi − µk ∥2 . 16: end for 17: Atomically add Nk (t), M k (t), Qk (t) to global accumulators N, M , Q in HBM. 18: end for P 19: Compute πknew = Nk / k′ Nk′ , µnew = M k /Nk , (σk2 )new = Qk /(D · Nk ) for all k. k new new 2 new K 20: return {πk , µk , (σk ) }k=1 .
stores the GMM parameters (O(KD) elements), and streaming tiles of the data matrix X. Thus, the kernel scales to arbitrarily large datasets. HBM bandwidth. A second advantage of the described kernel is the reduced number of HBM memory accesses. The data X, with N D elements, is read once from the HBM. The KD parameters {πk , µk , σk2 } are read twice, yielding N D + 2KD reads in total. Compared to the naive baseline of Section 3 (3N D + 4N K memory accesses), Flash-GMM eliminates the O(N K) responsibilitymatrix traffic entirely, reducing total HBM accesses to O(N D). That is the primary source of FlashGMM’s speedup over existing kernels. 3.3
Implementation Details
We implemented the Flash-GMM kernel using Triton (Tillet et al., 2019) and validated it on an NVIDIA A100 (80GB) GPU against a SciPy CPU reference. While the development was done on A100, the kernel has been validated to produce correct results on H100 and RTX5080 GPUs as well; the core algorithmic ideas are hardware-agnostic, and the implementation is straightforward to adapt to new architectures. The kernel uses a 1-D grid with ⌈N/BN ⌉ blocks, where the tile size is BN = 64, so each block processes 64 input vectors. Within a block, the K components are processed in chunks of BK = 16, and input vectors of dimension D are padded to BD = 128. Each block contains 4 warps, and blocks are scheduled independently across stream-
ing multiprocessors (SMs). For N = 106 , this gives ⌈N/BN ⌉ = 15,625 parallel blocks, providing enough work to fully occupy the 108 SMs of the A100. 3.4
Evaluation
To empirically validate the benefits of the new kernel, we test its runtime and memory costs across multiple scales of data. We compare FlashGMM to two contemporary baselines: the CPU implementation (on an AMD EPYC 7763 processor) from SciPy (Virtanen et al., 2020), and TorchGMM (CSOgroup, 2023), a GPU GMM kernel. Runtime and Scale Table 1 depicts the runtime measurements. Overall, Flash-GMM is 766– 1,740× faster than SciPy and 19–32× faster than TorchGMM across dataset sizes. Critically, TorchGMM runs out of memory at N > 106 , while Flash-GMM scales up to N = 108 on the same hardware, enabling soft GMM training on datasets more than 100× larger than previously feasible on a single device. Flash-GMM thus unlocks new use cases while delivering substantial speedups over existing implementations. Peak memory footprint Table 2 reports the GPU memory allocated by the kernel itself for FlashGMM and TorchGMM, across dataset sizes at K = 1024, D = 128. Flash-GMM’s kernel allocation grows as O(N ) via the log Zi buffer (N × 4 bytes) plus O(KD) accumulators (≈ 0.5 MB fixed), totalling 4.5 MB at N = 106 . TorchGMM materi-
N
Flash-GMM
TorchGMM
10K 50K 100K 250K 500K 1M
0.6 MB 0.7 MB 0.9 MB 1.5 MB 2.5 MB 4.5 MB
229 MB 1,067 MB 2,113 MB 5,262 MB 10,514 MB 21,006 MB
est centroids, and only the vectors contained in the associated posting lists are compared against q using the exact distance metric. The top-r nearest vectors among these candidates are then returned as the final retrieval results. 4.2
Table 2: Kernel GPU memory (excluding input data X) for Flash-GMM vs. TorchGMM (CSOgroup, 2023), K = 1024, D = 128, A100-80 GB. Flash-GMM allocates only log Zi (N ×4 bytes) plus O(KD) accumulators; TorchGMM materializes the full N ×K responsibility matrix, and exhausts memory for N > 1M .
alizes the full N ×K responsibility matrix and intermediate tensors, consuming 21 GB at the same scale — a 4,668× larger kernel footprint.
4
Novel Usage of GMMs for IVF
In this section, we introduce a novel use case for GMMs that is enabled by the Flash-GMM kernel. We apply Flash-GMM to the IVF index coarse quantizer (Jégou et al., 2011a), which typically relies on K-Means clustering, as a drop-in (§4.2) and with the addition of soft assignment (§4.3). Prior GMM implementations, whether CPUor GPU-based, are generally too compute- and memory-intensive to scale to the regimes required for practical IVF training. By introducing an optimized and scalable implementation, Flash-GMM allows for soft clustering within the IVF pipeline. This unlocks several advantages rooted in probabilistic soft assignments. First, the soft training procedure provides additional flexibility during clustering, allowing vectors to contribute to multiple clusters according to their posterior probabilities rather than enforcing hard assignments. Second, the probabilistic formulation naturally yields a principled multi-assignment strategy derived directly from the final responsibilities, which can improve recall by assigning vectors to multiple coarse partitions in a statistically-grounded manner. 4.1
IVF Indexing
An Inverted File (IVF) index partitions X into K cells using a coarse quantizer. This is almost universally implemented with k-means clustering (Johnson et al., 2021), where the cells correspond to clusters induced by the learned centroids. Each vector xi is assigned to a single cell and stored in the corresponding posting list. At query time, a query vector q is assigned to its nprobe near-
GMM as a Drop-in Coarse Quantizer
Replacing k-means with Flash-GMM in the IVF coarse quantizer requires no modifications to either the index structure or the query pipeline. The IVF index consumes only the cluster centroids produced by the quantizer, and Flash-GMM outputs centroids in the same format as k-means. Consequently, the search algorithm remains entirely unchanged. During index construction, each vector is assigned to the cluster with the highest responsibility, thereby collapsing the soft partition induced by GMM into the standard IVF posting-list structure. 4.3
GMM Multi-Assignment
Even with improved centroids from soft clustering, a vector near a Voronoi boundary is still stored in only one posting list. If a query falls on the other side of the boundary, the vector is invisible. The GMM training provides an immediate remedy: the final-iteration responsibilities rik directly quantify how much each vector “belongs to” each cluster. A vector with rik = 0.45, rij = 0.42 near two cluster boundaries clearly deserves to be indexed in both. We assign each vector xi to at most two clusters: the top-2 clusters whose responsibilities satisfy rik > τ with τ = 1/K. The threshold τ = 1/K corresponds to the uniform prior probability of a cluster, so a cluster is selected only if observing xi increases its posterior probability beyond the prior. 4.4
Experimental Setup
We compare 3 IVF coarse quantizers: • K-Means: 100 iterations of k-means from the FAISS library (Douze et al., 2026), serving as industry standard baseline. Single assignment. • GMM single: Flash-GMM with single assignment of each vector to a cluster. Model estimation warm-starts with 10 iterations of FAISS k-means, followed by 90 Flash-GMM iterations. • GMM multi: Flash-GMM with multiassignment using the τ = 1/K responsibility threshold. The model estimation process follows the procedure described for GMM single. The used datasets are:
SIFT1M (Jégou et al., 2011a): 106 SIFT descriptors, D = 128, 10K queries, standard ground truth. Deep10M: first 107 vectors of the deep-image96 dataset (Babenko and Lempitsky, 2016), D = 96, 10K queries, ground truth computed via bruteforce on the full 10M subset. GloVe-100: 1.18×106 GloVe word embeddings, D = 100, 10K queries, angular distance (Pennington et al., 2014). Vectors are L2-normalised before clustering (as is standard for angular benchmarks); ground truth recomputed on normalised vectors. Index construction uses FAISS IndexIVFFlat; centroids are injected directly without re-training. For multi-assignment, vectors are inserted into multiple posting lists in the standard FAISS index, without modifying the search routine. All methods use K = 1024 and the same random seeds. Quality metric We evaluate the search quality with Recall@10 (R@10): the fraction of queries for which the true nearest neighbor (determined by brute-force) appears in the top-10 returned results. Search cost metric We use the number of distance computation operations (DCO) per query as the primary cost metric. DCO directly measures the amount of search work performed, and is largely independent of hardware and the particular choice of K. Comparing methods using only nprobe can be misleading, since nprobe does not account for posting-list length. This is particularly important under multi-assignment, where improved recall may partially arise from scanning more vectors per probe due to longer posting lists. We therefore report R@10 as a function of both nprobe and DCO. 4.5
Results - Quality vs. Cost
Table 3 depicts a representative result from the GloVe-100 dataset, illustrating the recallcomputation trade-off for the different approaches. We see that for a given choice of nprobe , GMMsingle slightly outperforms K-Means at matched DCO; GMM-multi achieves substantially higher recall, but at the cost of higher DCO: at nprobe =16, GMM multi gains +7.0 pp over K-Means (0.92 vs. 0.85) while DCO rises from 18.4K to 32.8K. Importantly, the DCO increase is more than offset when comparing across operating points. In the GloVe-100 example, GMM-multi’s recall of 0.92 at nprobe =16 outperforms K-Means at nprobe =32 in terms of both recall and DCO. GMM-multi thus simultaneously delivers higher recall and lower com-
K-Means
GMM single
GMM multi
np R@10 DCO R@10 DCO R@10 DCO 16 32
0.85 0.90
18.4 36.9
0.86 0.90
18.4 36.9
0.92 0.95
32.8 65.6
Table 3: Recall@10 and DCO (×103 ) for GloVe-100 at selected nprobe values, K = 1024. For singleassignment methods DCO = nprobe × N/K; for GMM multi, DCO = nprobe × N m̄/K (m̄ = 1.78 for GloVe100). Full results across all datasets and nprobe values are in Appendix Table 7.
putation than the next reachable single-assignment operating point. Full results are in App. Table 7. Figure 1 depicts the full recall-DCO pareto curves, across all 3 datasets. Each curve depicts six points, corresponding to the different nprobe values. The figure demonstrates the consistent pareto-improvement pattern described above: because multi-assignment inflates list lengths, each GMM-multi point sits to the right of the singleassignment point at the same nprobe – but it also sits above the single-assignment point at an equivalent DCO, which corresponds to a single-assignment point with a higher nprobe . At every such comparison, GMM multi delivers strictly higher recall for the same compute budget. The gain is largest on GloVe (Figure 1(c)), potentially since word-embedding spaces have high semantic density: many vectors cluster near boundaries between topically related word clusters, yielding a high assignment average and more boundary vectors that benefit from multi-assignment. Experiments were executed with 3 random seeds. Above we report the results for the first seed; we find that Recall@10 at any fixed nprobe varies by less than 0.003, and the multi-assignment recall gains exhibit similarly low variance. All reported results use K = 1024, though we verified that the recall improvements persist across K ∈ {256, 1024, 4096} (Appendix Table 5). All experiments use a warm-start initialized with 10 k-means iterations. We also evaluated k-means++ initialization (Arthur and Vassilvitskii, 2007), but warm-start achieved comparable or better recall on most datasets while providing faster training (Appendix C). Index build runtime Table 4 reports the index build wall-clock time for K-Means and FlashGMM across all datasets, on a single Nvidia A10080GB GPU. Importantly, training a GMM coarse
1.00
0.99
Recall@10
0.98 0.96
0.95
0.94
k-means (1-assign)
0.92
0.90 Flash-GMM single-assign Flash-GMM multi-assign avg 1.49 assign./vec( (==1/K 1/K) )
0.90 0.88
10K
20K
30K
40K
Distance computations per query (DCO)
0.99
Recall@10
0.98
0.90
0.90
k-means (1-assign)
0.88
Flash-GMM single-assign Flash-GMM multi-assign avg 1.60 assign./vec( (==1/K 1/K) )
0.85 40K
60K
80K
100K
Distance computations per query (DCO)
120K
(b) Deep10M (D = 96) 0.95
0.95
Recall@10
TorchGMM
SciPy
11.8s 123.1s 21.8s
247.8s OOM 493.86s
≈5.6h ≈58h ≈10h
Table 4: Index build runtime for K-Means, Flash-GMM, TorchGMM, and SciPy. Due to prohibitively long wallclock times, SciPy runtimes are approximated based on the speedup factors observed in Table 1.
multi-assignment reaches the same recall targets at substantially lower per-query DCO, so the savings compound across the lifetime of the index. 4.6
0.93
0.90
0.90 0.85
k-means (1-assign)
0.80 0.75
Flash-GMM
3.5s 53.3s 6.4s
0.95
0.95
0.83
K-Means
50K
(a) SIFT1M (D = 128) 1.00
SIFT1M Deep10M GloVe-100
Flash-GMM single-assign Flash-GMM multi-assign avg 1.78 assign./vec( (==1/K 1/K) ) 10K
20K
30K
40K
50K
Distance computations per query (DCO)
60K
(c) GloVe-100 (D = 100)
Figure 1: Recall@10 vs. DCO (K = 1024). Each point corresponds to one nprobe value. GMM multi dominates both baselines across the full latency budget. Dashed reference lines mark recall targets 0.90, 0.95, 0.99.
quantizer at N = 107 (Deep10M) was previously impractical on an A100-80GB GPU: existing GMM implementations typically exhaust GPU memory beyond N ≈ 106 . Scaling GMM training to this regime is thus a key contribution of FlashGMM, enabled by its memory-efficient design. In terms of runtime, Flash-GMM is approximately 2.5×–3.3× slower than K-Means. However, the absolute overhead amounts only to additional one-time offline work.2 Moreover, the additional training cost is amortized over query time: 2 Multi-assignment introduces no additional training cost beyond a single scan of the N ×K responsibility matrix, which is negligible at N = 106 , K = 1024.
Multi-Assignment Analysis
We analyze the practical behavior of the proposed multi-assignment strategy. Under the chosen responsibility threshold τ = 1/K, the average number of posting lists associated with a vector is 1.49 for SIFT1M, 1.60 for Deep10M, and 1.78 for GloVe-100. In practice, most vectors have one dominant assignment, or at most two clusters whose responsibility exceeds τ . We additionally evaluated τ = 2/K and τ = 0.5/K, and found both inferior to τ = 1/K. The former reduces recall by 1–4 percentage points at every nprobe , while yielding only marginal searchtime savings. The latter increases search time due to larger posting lists, without improving recall. To isolate the benefit of GMM responsibilities from multi-assignment alone, we compare GMM multi against K-Means hard top-2. This baseline assigns each vector to the posting lists of its two nearest k-means centroids according to L2 distance, regardless of proximity to the cluster boundary. Despite doubling the posting-list size, hard top2 requires up to 1.8× more DCO than singleassignment k-means to achieve the same recall (Appendix E). This indicates that the gains of GMM multi are not explained by multi-assignment alone. Prior work also observes this limitation of hard top-2 assignment. To obtain higher-quality multiassignments for k-means-based indexing, RAIRS (Yang and Chen, 2026) propose the AIR heuristic for determining the posting lists. Since they did not release public code, we cannot directly compare to their approach; for completeness, we adapt GMM multi-assignment to their chosen index configuration, and show that GMM multi achieves competitive recall with less index inflation, compared to their reported results on a shared dataset (Appendix F).
5
Related Work
Soft-assignment quantizers. To our knowledge, no prior work applies scalable soft-assignment EM to the IVF quantizer. Fine-quantizer methods. Product quantization (PQ, Jégou et al., 2011a) and its extensions improve the residual quantizer applied after coarse assignment, and are orthogonal to Flash-GMM. Assignment-side methods. RAIRS (Yang and Chen, 2026) corrects boundary assignment errors by storing each vector in two posting lists using the AIR geometric heuristic to select the second list, without altering centroids — assigning a second list to every vector regardless of boundary proximity. IO-aware kernels. Flash-k-means (Yang et al., 2026) applies a fused distance kernel to reduce HBM bandwidth in hard-assignment clustering, sharing the IO-aware tiling motivation of our work. Our kernel targets the GMM E-step: it computes soft responsibilities via a numerically stable logsum-exp and accumulates responsibility-weighted sufficient statistics in the same pass — operations with no analogue in hard-assignment k-means.
6
Discussion
In this work we presented two complementary contributions to scalable soft clustering with GMMs. Flash-GMM makes GMM training practical at scale, via a fused Triton kernel with O(KD) working memory, eliminating the memory barrier that previously confined soft EM to small datasets. GMM multi-assignment reuses the final FlashGMM responsibilities to store boundary vectors in multiple posting lists. Together, the two methods push the recall– compute Pareto frontier significantly beyond what either achieves alone. In our application on the IVF coarse quantizer, we demonstrate the practical utility unlocked by Flash-GMM multi-assign — reaching superior results compared to k-means, while operating on large datasets for which GMM was not previously feasible. The benefits we demonstrate in the context of IVF coarse quantization can potentially be coupled with other advances, most notably fine-quantizer methods such as IVF-PQ (Jégou et al., 2011a) and IVF-PQfs (André et al., 2016). Future work can explore the potential for such combinations to further promote ANN search applications.
The IVF use case, however, is merely one example; GMMs have varied applications across multiple domains, which could similarly benefit from this efficient Flash-GMM implementation. In medical imaging, for instance, a probabilistic GMM approach has proven useful for image segmentation problems (Song et al., 2014; Riaz et al., 2020). Similarly, GMMs are explored for diverse use cases in genomics, such as clustering gene expression patterns (Liu et al., 2022) or modeling relationships between the genomes of different organisms (Clarke et al., 2018). Critically, many of the existing use cases for GMMs operate at particularly vast data scales, and are currently constrained to applying inefficient CPU-based implementations. Furthermore, this unlocks the potential for increased use of GMMs. Prior work has shown GMMs can act as a drop-in replacement for kmeans with improved partition quality across diverse workloads (Patel and Kushwaha, 2020; Liang et al., 2022); Flash-GMM removes the scalability barrier that has prevented this potential from being realized in practice. The Flash-GMM kernel was optimized for the A100 GPU. The H100 architecture introduces new hardware primitives - most notably the Tensor Memory Accelerator and asynchronous warpgroup MMA instructions. A Flash-GMM kernel redesigned around these primitives could exploit H100-specific features to deliver speedups beyond those reported here. This mirrors the transition from the IO-aware tiling of FlashAttention-2 (Dao et al., 2022) to the H100-enabled speedups demonstrated by FlashAttention-3 (Shah et al., 2024). We leave this as a natural direction for future work. We release the Flash-GMM kernel as a standalone library; multi-assignment requires no additional code beyond a threshold on the existing responsibility matrix. We hope these tools lower the barrier to GMM usage in diverse research applications as well as production ANN systems. Lastly, beyond GMM-specific applications, the computational pattern at the heart of Flash-GMM – a fused log-sum-exp over K weighted Gaussian terms, followed by responsibility-weighted accumulation of sufficient statistics – is not unique to GMM training. The same pattern — per-point GMM responsibilities combined into weighted statistics — appears in Fisher Vector encoding (Perronnin et al., 2010) and kernel density estimation. Flash-GMM’s IO-aware tiling strategy has the potential to accelerate these settings as well.
Limitations Training cost Flash-GMM is 2–3× slower than k-means. For applications requiring frequent reindexing this may be a practical constraint. Index size Multi-assignment increases the stored index size by m̄ on average (1.49–1.78× in our experiments). For memory-constrained deployments this may be a consideration. The storage overhead is concentrated in posting lists and is proportional to m̄; for most datasets m̄ < 2, keeping the overhead below 2×. We demonstrate IVF results at N = 107 ,
Scale with the kernel validated at 108 using O(KD) working memory. For billion-scale training (N ≥ 109 ), the input data X itself requires ≥512 GB GPU memory and cannot fit on a single device. SSD streaming processing X in chunks loaded from disk with the kernel accumulating into the same O(KD) buffers across chunks enables fullbatch EM over arbitrary N with no quality loss. Applying this to real billion-scale ANN datasets is left to future work.
Conference on Very Large Data Bases, volume 9, page 12. David Arthur and Sergei Vassilvitskii. 2007. kmeans++: the advantages of careful seeding. In Proceedings of the Eighteenth Annual ACM-SIAM Symposium on Discrete Algorithms, SODA ’07, pages 1027–1035, Philadelphia, PA, USA. Society for Industrial and Applied Mathematics. Artem Babenko and Victor Lempitsky. 2016. Efficient indexing of billion-scale datasets of deep descriptors. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 2055–2063. Christopher M. Bishop. 2006. Pattern Recognition and Machine Learning. Springer. Pierre Blanchard, Desmond J. Higham, and Nicholas J. Higham. 2021. Accurately computing the log-sumexp and softmax functions. IMA Journal of Numerical Analysis, 41(4):2311–2330. Thomas H Clarke, Lauren M Brinkac, Granger Sutton, and Derrick E Fouts. 2018. Ggrasp: a r-package for selecting representative genomes using gaussian mixture models. Bioinformatics, 34(17):3032–3034. CSOgroup. 2023. TorchGMM: Gaussian mixture models in PyTorch. https://github.com/CSOgroup/ torchgmm. GitHub repository.
K regime We evaluate K ∈ {256, 1024, 4096}; the recall gain is strongest at K = 1024. The interaction of multi-assignment with larger K is unexplored.
Tri Dao, Daniel Y Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. 2022. FlashAttention: Fast and memory-efficient exact attention with IO-awareness. In Advances in Neural Information Processing Systems, volume 35, pages 16344–16359.
Isotropic vs. Full Covariance The GMM formulation explored here and implemented in the Flash-GMM kernel uses isotropic (scalar) covariance matrices rather than full covariance. As detailed in Appendix A, full covariance is ill-suited for standard IVF scales.
A. P. Dempster, N. M. Laird, and D. B. Rubin. 1977. Maximum likelihood from incomplete data via the em algorithm. Journal of the Royal Statistical Society: Series B, 39(1):1–38.
Combining Flash-GMM centroids with the AIR assignment heuristic RAIRS’s AIR metric selects second-list assignments based on a geometric residual criterion, applied on top of k-means centroids. An interesting direction is whether applying AIR on Flash-GMM centroids (which already improve recall by 25–33% in single-assignment mode) yields additive gains. This experiment is not currently possible without re-implementing RAIRS from scratch, as no public code is available.
Matthijs Douze, Alexandr Guzhva, Chengqi Deng, Jeff Johnson, Gergely Szilvasy, Pierre-Emmanuel Mazaré, Maria Lomeli, Lucas Hosseini, and Hervé Jégou. 2026. The faiss library. IEEE Transactions on Big Data, 12(2):346–361. Hervé Jégou, Matthijs Douze, and Cordelia Schmid. 2011a. Product quantization for nearest neighbor search. IEEE Transactions on Pattern Analysis and Machine Intelligence, 33:117–128. Hervé Jégou, Romain Tavenard, Matthijs Douze, and Laurent Amsalem. 2011b. Searching in one billion vectors: re-rank with source coding. In Proceedings of the IEEE International Conference on Acoustics, Speech and Signal Processing, pages 861–864.
References
Jeff Johnson, Matthijs Douze, and Hervé Jégou. 2021. Billion-scale similarity search with GPUs. IEEE Transactions on Big Data, 7(3):535–547.
Fabien André, Anne-Marie Kermarrec, and Nicolas Le Scouarnec. 2016. Cache locality is not enough: High-performance nearest neighbor search with product quantization fast scan. In 42nd International
Chen Liang, Wenguan Wang, Jiaxu Miao, and Yi Yang. 2022. Gmmseg: Gaussian mixture based generative semantic segmentation models. Advances in Neural Information Processing Systems, 35:31360–31375.
Ta-Chun Liu, Peter N Kalugin, Jennifer L Wilding, and Walter F Bodmer. 2022. GMMchi: gene expression clustering using gaussian mixture modeling. BMC bioinformatics, 23(1):457. Stuart Lloyd. 1982. Least squares quantization in PCM. IEEE Transactions on Information Theory, 28(2):129–137. Eva Patel and Dharmender Singh Kushwaha. 2020. Clustering cloud workloads: K-means vs gaussian mixture model. Procedia Computer Science, 171:158–167. Third International Conference on Computing and Network Communications (CoCoNet’19). Jeffrey Pennington, Richard Socher, and Christopher D Manning. 2014. GloVe: Global vectors for word representation. In Proceedings of the 2014 Conference on Empirical Methods in Natural Language Processing, pages 1532–1543. Florent Perronnin, Jorge Sánchez, and Thomas Mensink. 2010. Improving the Fisher kernel for large-scale image classification. In Proceedings of the European Conference on Computer Vision, pages 143– 156. Springer. Farhan Riaz, Saad Rehman, Muhammad Ajmal, Rehan Hafiz, Ali Hassan, Naif Radi Aljohani, Raheel Nawaz, Rupert Young, and Miguel Coimbra. 2020. Gaussian mixture model based probabilistic modeling of images for medical image segmentation. IEEE Access, 8:16846–16856. Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, and Tri Dao. 2024. FlashAttention-3: Fast and accurate attention with asynchrony and low-precision. arXiv preprint arXiv:2407.08608. Yantao Song, Zexuan Ji, and Quansen Sun. 2014. An extension gaussian mixture model for brain mri segmentation. In 2014 36th Annual International Conference of the IEEE Engineering in Medicine and Biology Society, pages 4711–4714. IEEE. Philippe Tillet, Hsiang-Tsung Kung, and David Cox. 2019. Triton: an intermediate language and compiler for tiled neural network computations. In Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages, pages 10–19. Pauli Virtanen, Ralf Gommers, Travis E. Oliphant, Matt Haberland, Tyler Reddy, David Cournapeau, Evgeni Burovski, Pearu Peterson, Warren Weckesser, Jonathan Bright, Stéfan J. van der Walt, Matthew Brett, Joshua Wilson, K. Jarrod Millman, Nikolay Mayorov, Andrew R. J. Nelson, Eric Jones, Robert Kern, Eric Larson, and 16 others. 2020. SciPy 1.0: Fundamental Algorithms for Scientific Computing in Python. Nature Methods, 17:261–272.
Shuo Yang, Haocheng Xi, Yilong Zhao, Muyang Li, Xiaoze Fan, Jintao Zhang, Han Cai, Yujun Lin, Xiuyu Li, Kurt Keutzer, Song Han, Chenfeng Xu, and Ion Stoica. 2026. Flash-KMeans: Fast and memoryefficient exact k-means. Preprint, arXiv:2603.09229. Zehai Yang and Shimin Chen. 2026. RAIRS: Optimizing redundant assignment and list layout for IVFbased ANN search. Proceedings of the ACM on Management of Data, 4(1).
A
Isotropic vs. Full Covariance
Flash-GMM relies on isotropic (scalar) covariance matrices rather than full covariance. While full matrices capture feature correlations, they require estimating O(D2 ) parameters per component (e.g., 8,256 parameters for SIFT’s D = 128). At standard IVF scales (N = 106 , K = 1024), each cluster receives only ≈1,000 points, making a full-covariance fit severely underdetermined. This causes the empirical covariance matrices to become ill-conditioned or singular, leading to structural collapse during EM. Making full covariance mathematically viable would require drastically reducing K to ensure enough points per cluster. However, a much smaller K produces massive Voronoi cells, which would fundamentally break the IVF search efficiency and require a complete redesign of the querying heuristic (such as hierarchical indexing or aggressive dimensionality reduction), rather than serving as a drop-in replacement for k-means.
B
Effect of K on the Recall Gain K
FAISS
Flash-GMM warm
GMM multi
256 1024 4096
16 48 96
16 32 96
12 24 64
Table 5: The minimal nprobe required to reach R@10 of 0.99, over SIFT1M, for different values of K. For all K values, GMM multi requires the smallest value of nprobe to reach 0.99 R@10.
C
Warm-Start vs. kmeans++ Initialisation
Flash-GMM supports two initialisation strategies: • Warm-start: 10 iterations of FAISS k-means, then 90 iterations of Flash-GMM soft EM from the resulting centroids. • kmeans++: 100 iterations of Flash-GMM soft EM from a random kmeans++ seeding.
Table 6 compares recall@10 and training time for both initialisations on SIFT1M and GloVe-100. Dataset
Init
np=16 np=32 np=48 Time (s)
SIFT1M
Warm-start 0.962 0.991 0.997 kmeans++ 0.962 0.991 0.997
11.8 48.6
GloVe-100
Warm-start 0.856 0.902 0.924 kmeans++ 0.856 0.902 0.924
21.8 89.4
Table 6: Recall@10 at selected nprobe values and training time for warm-start vs. kmeans++ initialisation, K = 1024, single assignment.
Warm-start is 3–4× faster than kmeans++ and achieves identical recall on both datasets, confirming that the k-means warm-start basin is already well-aligned with the soft-EM optimum. We recommend warm-start as the default initialisation.
D
hard top-2: each vector is assigned to its two nearest k-means centroids by L2 distance, giving a fixed m̄ = 2.0 for every vector regardless of boundary proximity. Table 8 reports the DCO required to reach recall targets of 0.90 and 0.95. Kmeans hard top-2 requires up to 1.8× more DCO than FAISS single to reach the same recall target: doubling the index size yields only marginal recall gains, because the second nearest centroid by L2 distance rarely contains the true nearest neighbor when the first does not. Hard top-2 is therefore less DCO-efficient than plain FAISS single assignment, confirming that redundant assignment without a principled boundary signal is counterproductive. Method
SIFT1M
single hard top-2
9.4 12.6
14.7 17.2
Deep1M
single hard top-2
7.3 12.2
12.2 18.7
GloVe-100
single hard top-2
37.6 47.7
– –
Full Recall–DCO Results
Table 7 depicts the full results of recall and DCO for every nprobe value. FAISS
GMM single
GMM multi
Dataset
np R@10 DCO R@10 DCO R@10 DCO
SIFT1M
1 4 8 16 32 48
0.439 0.767 0.885 0.958 0.988 0.995
1.0 3.9 7.8 15.6 31.3 46.9
0.434 0.767 0.890 0.962 0.991 0.997
1.0 3.9 7.8 15.6 31.3 46.9
0.530 0.851 0.942 0.985 0.998 0.999
1.5 5.8 11.6 23.3 46.6 69.8
Deep10M
1 4 6 8 12 16 32 48
0.549 0.853 0.910 0.932 0.958 0.973 0.992 0.996
9.8 39.0 58.5 78.1 117.1 156.2 312.2 468.4
0.553 0.860 0.916 0.938 0.963 0.976 0.993 0.996
9.8 39.0 58.5 78.1 117.1 156.2 312.2 468.4
0.668 0.930 0.963 0.978 0.990 0.994 0.999 0.999
15.6 62.4 93.7 124.9 187.3 249.9 499.5 749.5
1 4 8 GloVe-100 16 32 48
0.453 0.698 0.782 0.847 0.899 0.927
1.2 4.6 9.2 18.4 36.9 55.3
0.471 0.718 0.795 0.856 0.902 0.924
1.2 4.6 9.2 18.4 36.9 55.3
0.587 0.810 0.872 0.917 0.949 0.963
2.1 8.2 16.4 32.8 65.6 98.5
Table 7: Recall@10 and DCO (×103 ) for all three datasets, K = 1024, nprobe ∈ {1, 4, 8, 16, 32, 48}. For single-assignment methods DCO = nprobe × N/K; for GMM multi, DCO = nprobe × N m̄/K (m̄: SIFT1M 1.49, Deep10M 1.60, GloVe-100 1.78). Best recall per row in bold.
E
DCO at DCO at R@10 = 0.90 R@10 = 0.95
Dataset
Hard Multi-Assignment Ablation
To isolate the contribution of GMM responsibilities from the benefit of redundant assignment alone, we compare GMM multi-assignment against Kmeans
Table 8: DCO (×103 ) required to reach recall@10 targets of 0.90 and 0.95, at K = 1024. Kmeans hard top-2 (m̄ = 2.0, nearest-centroid L2) is strictly less DCOefficient than FAISS single at both targets. “–” = target not reached within nprobe = 48.
F
IVF-PQ Compatibility and RAIRS Comparison
Flash-GMM is fully compatible with IVF-PQ: the coarse centroids are substituted without modifying PQ training or encoding. Table 9 reports results across IVF-PQ (M ∈ {16, 8, 4}, 8-bit) and IVF-PQfs (M = 64, 4-bit fast-scan) on SIFT1M. The coarse-quantizer benefit is orthogonal to PQ fidelity: at IVF-PQ M = 16, multi-assign reaches recall ≥ 0.87 at nprobe =8 vs. FAISS’s 0.830 a +5.8 pp gain and gains of +2–4 pp persist even at high compression where PQ quantisation error dominates. The IVF-PQfs row (M = 64, 4-bit, nlist=1024) matches the configuration evaluated by RAIRS (Yang and Chen, 2026). Our approach differs in three ways: we improve the centroids themselves via soft EM; assignments are derived from GMM responsibilities rather than a geometric heuristic; and redundancy is adaptive per-vector
Index
M
FAISS
GMM single GMM multi
np=8 np=12 np=8 np=12 np=8 np=12 IVF-PQ (8-bit)
16 0.830 0.867 0.831 0.871 0.888 0.909 8 0.671 0.691 0.669 0.688 0.703 0.715 4 0.450 0.456 0.450 0.457 0.472 0.475
IVF-PQfs (4-bit) 64 0.844 0.885 0.850 0.891 0.926 0.949
Table 9: IVF-PQ and IVF-PQfs recall@10 on SIFT1M, K = 1024. IVF-PQfs (M = 64, 4-bit fastscan, nlist=1024) matches the configuration used by RAIRS (Yang and Chen, 2026). IVF-PQ uses 8-bit codes. Bold = best per row.
rather than universal. Both methods likely select similar second-list candidates for genuine boundary vectors, but diverge on interior vectors, where AIR always assigns a second list while GMM assigns only when rik > 1/K. Based on their published recall–nprobe curves, GMM multi is competitive with RAIRS on this configuration, while requiring approximately 30% fewer probes to reach the same recall target and achieving 25% less index inflation (m̄ = 1.49 vs. RAIRS’s universal m̄ = 2.0). An exact head-to-head is not possible without RAIRS’s code, but the shared configuration makes Table 9 the closest available proxy.