PaperlayerView exact revision on arXiv

Source-linked AI summary

Flash-GMM: A Memory-Efficient Kernel for Scalable Soft Clustering

Gal Bloch, Ariel Gera, Matan Orbach, Ohad Eytan, Assaf Toledo

arXiv:2606.10896v1cs.LGcs.DBcs.IRcs.PF

TL;DR

Large-scale GMM estimation is constrained by GPU memory because materializing responsibilities is impractical. Flash-GMM avoids that matrix with a fused Triton kernel and, in IVF search, achieves up to 1.7× fewer distance computations at fixed recall or +2–12 recall@10 at matched cost.

  • Problem

    GPU memory makes materializing the full GMM responsibility matrix impractical for large datasets, limiting scalable estimation.

  • Method

    Flash-GMM is a fused, tile-based Triton kernel that performs GMM estimation without materializing responsibilities, using O(KD) working memory.

  • Results

    Up to 1.7× fewer distance computations are achieved at fixed recall, or +2–12 recall@10 at matched computational cost in IVF search.

  • Takeaways & Limitations

    Flash-GMM makes soft GMM clustering practical at scale and supports multi-assignment that improves the IVF recall–compute tradeoff beyond k-means.

  • Takeaways & Limitations

    Real billion-scale ANN datasets remain future work because the input data alone may exceed single-device GPU memory.

Abstract

from arXiv · show

We present \textbf{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 \textbf{20$\times$} speedup over existing implementations and enables training on datasets more than \textbf{100$\times$} larger than previously feasible on one device. To demonstrate its impact, we integrate Flash-GMM 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\times$ fewer distance computations, or equivalently, yields $+2$--$12$ recall@10 at matched computational cost. We release the kernel as an open-source project.

1 Introduction

Flash-GMM is a fused, memory-efficient Triton kernel that avoids materializing GMM responsibilities, addressing GPU-memory limits and accelerating large-scale training. Integrated into IVF coarse quantization, it enables soft multi-assignment and improves ANN recall–compute tradeoffs.

  • Motivation: GMM responsibilities form an N × K matrix recomputed each EM iteration, making full materialization impractical on memory-limited GPUs.Existing GPU implementations run out of memory beyond 10 million data points, while CPU solvers are orders of magnitude too slow.
  • Method: Flash-GMM is a fused Triton kernel that performs GPU-accelerated GMM estimation without materializing the responsibility matrix in GPU HBM memory.Its design adapts FlashAttention’s IO-aware tiling strategy to the EM algorithm.
  • Results: 20× speedup over existing implementations and training on datasets more than 100× larger than previously feasible on one device demonstrate Flash-GMM’s scalability.These results are reported for GMM EM training and compare against existing implementations.
  • IVF application: Flash-GMM replaces k-means with soft GMM clustering in the IVF coarse quantizer, enabling larger-scale index construction and multiple cluster assignments.Responsibilities provide posterior probabilities that naturally support assigning vectors near cluster boundaries to multiple clusters.
  • ANN results: 1.7× fewer distance computations at fixed recall and +2–12 recall at matched compute budget show improved ANN search quality–cost tradeoffs.Multiple assignment prevents near-boundary vectors from being prematurely discarded from the search space.

2 Gaussian Mixture Models

Gaussian mixture models represent data as a mixture of K Gaussian components, here restricted to isotropic covariances. Parameter estimation uses iterative EM updates, with responsibilities computed in the E-step and model parameters updated in the M-step until convergence.

  • Model formulation: GMMs model each sample as generated from a mixture of K Gaussian distributions.The data matrix contains samples x_i ∈ R^D.
  • Model formulation: Each covariance matrix is constrained to the isotropic form σ^2_kI to improve statistical stability while remaining computationally tractable.The isotropic restriction is stated as an assumption for the Gaussian components.
  • Parameter estimation: The EM algorithm alternates between an E-step that computes component responsibilities and an M-step that updates model parameters.The responsibility r_ik measures the degree to which sample x_i belongs to cluster k.
  • Parameter estimation: Responsibilities are normalized across components and computed in log-space using the log-sum-exp trick for numerical stability.The E-step first forms unnormalized assignment scores before normalization.
  • Parameter estimation: The M-step updates effective component counts, mixture weights, means, and variances using the responsibilities from the E-step.These updates are applied iteratively after responsibility computation.
  • Parameter estimation: EM repeats the E-step and M-step until the log-likelihood changes sufficiently little or a predefined maximum number of iterations is reached.Convergence can therefore be determined by either a likelihood-change criterion or an iteration limit.

3 Flash-GMM

Flash-GMM addresses GPU GMM scalability by avoiding responsibility-matrix materialization and reducing repeated HBM traffic through tiled, on-chip computation. This design enables substantially faster runtime and much larger datasets than existing CPU and GPU implementations.

  • Kernel design: Flash-GMM divides X into tiles, processes K components in blocks, computes log-normalizers and responsibilities on chip, and accumulates sufficient statistics for the M-step.Per-tile contributions are atomically reduced into global N, M, and Q accumulators after all tiles are processed.
  • Memory efficiency: Flash-GMM never materializes the N×K responsibility matrix in HBM, storing only O(KD) GMM parameters and streaming data tiles.Per-vector log-normalizers and per-tile accumulators remain on chip, allowing the kernel to scale to arbitrarily large datasets.
  • Memory efficiency: O(ND) HBM accesses replace the naive baseline’s 3ND + 4NK accesses, eliminating O(NK) responsibility-matrix traffic.Flash-GMM reads X once and the KD model parameters twice; this reduced traffic is identified as the primary source of speedup.
  • Evaluation: 766–1,740× faster than SciPy and 19–32× faster than TorchGMM, Flash-GMM scales to N = 10^8 while TorchGMM runs out of memory at N > 10^6.On the same hardware, this enables soft GMM training on datasets more than 100× larger than previously feasible on one device.
  • Evaluation: 4,668× larger kernel footprint for TorchGMM than Flash-GMM at N = 10^6, with K = 1024 and D = 128.Flash-GMM allocates 4.5MB at this scale, whereas TorchGMM consumes 21GB by materializing the full responsibility matrix and intermediate tensors.

4 Novel Usage of GMMs for IVF

Flash-GMM enables GMM-based soft clustering as a drop-in IVF coarse quantizer without changing the index or search pipeline. Responsibility-based multi-assignment improves recall–cost trade-offs, while Flash-GMM scales GMM training to previously impractical dataset sizes.

  • Drop-in IVF coarse quantization: Flash-GMM replaces k-means in the IVF coarse quantizer without modifying the index structure or query pipeline.It outputs centroids in the same format as k-means, and standard IVF assignment uses the highest-responsibility cluster.
  • Soft multi-assignment: GMM multi assigns each vector to at most two clusters when the top-2 responsibilities exceed τ = 1/K, enabling responsibility-derived boundary-vector replication.The threshold corresponds to a cluster’s uniform prior probability, selecting clusters whose posterior probability exceeds that prior.
  • Recall–cost trade-offs: +7.0 pp recall@10 over K-Means at nprobe=16 raises recall from 0.85 to 0.92 on GloVe-100, while DCO increases from 18.4K to 32.8K.Across operating points, GMM-multi can outperform single-assignment methods at equivalent DCO; the largest gain occurs on GloVe-100.
  • Scalable training: N ≈ 10^7 on Deep10M was previously impractical on an A100-80GB GPU because existing GMM implementations typically exhausted memory beyond N ≈ 10^6.Flash-GMM scales GMM training to this regime through its memory-efficient design, although it is approximately 2.5×–3.3× slower than K-Means in runtime.
  • Threshold and assignment analysis: 1.49, 1.60, and 1.78 are the average posting-list counts per vector for SIFT1M, Deep10M, and GloVe-100 under τ = 1/K.Alternative thresholds τ = 2/K and τ = 0.5/K are inferior, while hard top-2 requires up to 1.8× more DCO than single-assignment k-means for the same recall.

5 Related Work

Flash-GMM is, to the authors’ knowledge, the first scalable soft-assignment EM method applied to the IVF quantizer. It differs from complementary fine-quantizer and assignment-side methods, while sharing IO-aware kernel motivation with Flash-k-means but targeting GMM-specific E-step computations.

  • Soft-assignment quantizers: No prior work, to the authors’ knowledge, applies scalable soft-assignment EM to the IVF quantizer.
  • Fine-quantizer methods: Product quantization and its extensions improve residual quantization after coarse assignment and are orthogonal to Flash-GMM.
  • Assignment-side methods: RAIRS assigns every vector to two posting lists using the AIR heuristic for the second list, without altering centroids or considering boundary proximity.
  • IO-aware kernels: Flash-k-means reduces HBM bandwidth with a fused distance kernel, whereas Flash-GMM computes stable logsum-exp responsibilities and accumulates weighted sufficient statistics in one pass.

6 Discussion

The discussion presents Flash-GMM’s scalable training and responsibility-based multi-assignment as complementary methods that improve IVF coarse quantization beyond k-means while making GMMs practical on large datasets. It also outlines broader applications, future hardware and quantizer extensions, and the released kernel’s reuse beyond GMM training.

  • Contributions: Flash-GMM uses a fused Triton kernel with O(KD) working memory, while multi-assignment reuses final responsibilities to place boundary vectors in multiple posting lists.Together, these contributions address scalable soft clustering and multi-cluster assignment.
  • Results: The two methods push the recall–compute Pareto frontier beyond either method alone, delivering superior IVF coarse-quantizer results to k-means on previously infeasible large datasets.The discussion attributes this practical utility to Flash-GMM multi-assignment.
  • Broader impact: Flash-GMM removes the scalability barrier limiting GMM adoption, supporting potential drop-in replacement of k-means and applications including medical-image segmentation and genomics.The discussion also identifies possible combinations with fine-quantizer methods such as IVF-PQ and IVF-PQfs for ANN search.
  • Limitations and future work: Because the kernel targets A100 GPUs, redesigning it around H100 primitives such as Tensor Memory Accelerator and asynchronous warpgroup MMA could yield further speedups.This is presented as future work beyond the reported results.
  • Extensions: The released standalone library enables multi-assignment through a threshold on the existing responsibility matrix, while its fused responsibility-and-statistics pattern could accelerate Fisher Vector encoding and kernel density estimation.The authors hope these tools lower barriers to GMM use in research and production ANN systems.

Limitations · A Isotropic vs. Full Covariance · B Effect of K on the Recall Gain

Flash-GMM’s main limitations are higher training and storage costs, unresolved billion-scale deployment, and an unexplored larger-K regime. The method uses isotropic covariance because full covariance is severely underdetermined at standard IVF scales, while GMM multi achieves the lowest nprobe requirement across evaluated K values.

  • Limitations: 2–3× slower training than k-means may constrain applications requiring frequent reindexing.The training-cost limitation is identified as practical for frequently reindexed applications.
  • Limitations: 1.49–1.78× larger stored indexes result from multi-assignment, with overhead concentrated in posting lists.The overhead is proportional to average multiplicity m̄; for most datasets, m̄ < 2, keeping storage overhead below 2×.
  • Limitations: N ≥ 10^9 cannot fit on one device because the input X alone requires ≥512 GB GPU memory.SSD streaming in chunks with accumulated O(KD) buffers could enable fullbatch EM without quality loss, but real billion-scale ANN evaluation remains future work.
  • Limitations: Recall gains are strongest at K = 1024, while the interaction between multi-assignment and larger K remains unexplored.The evaluation covers K ∈ {256, 1024, 4096}.
  • A Isotropic vs. Full Covariance: Flash-GMM uses isotropic scalar covariance rather than full covariance because full covariance is ill-suited for standard IVF scales.Full covariance requires O(D^2) parameters per component, including 8,256 for SIFT with D = 128.
  • A Isotropic vs. Full Covariance: At N = 10^6 and K = 1024, approximately 1,000 points per cluster make full-covariance fitting severely underdetermined.The resulting empirical covariance matrices become ill-conditioned or singular, leading to structural collapse during the fitting process.
  • B Effect of K on the Recall Gain: GMM multi requires the smallest nprobe to reach R@10 of 0.99 for every evaluated K on SIFT1M.Table 5 compares the minimal nprobe across different K values.

C Warm-Start vs. kmeans++ Initialisation

Flash-GMM supports warm-start and kmeans++ initialization strategies. Warm-start is 3–4× faster than kmeans++ while achieving identical recall@10 on SIFT1M and GloVe-100, making it the recommended default.

  • Warm-start runs 10 iterations of FAISS k-means followed by 90 iterations of Flash-GMM soft EM from the resulting centroids.
  • kmeans++ runs 100 iterations of Flash-GMM soft EM from a random kmeans++ seeding.
  • 3–4× faster: warm-start outpaces kmeans++ while achieving identical recall@10 on SIFT1M and GloVe-100.Table 6 compares recall@10 and training time for both initializations.
  • Warm-start is recommended as the default initialization because its k-means warm-start basin is already well-aligned with the soft-EM optimum.

D Full Recall–DCO Results

Table 7 reports Recall@10 and DCO across SIFT1M, Deep10M, and GloVe-100 for K = 1024 and nprobe values from 1 to 48. DCO is defined differently for single-assignment and GMM multi-assignment methods.

  • Evaluation scope: Table 7 covers Recall@10 and DCO (×10^3) for three datasets at K = 1024 across nprobe ∈ {1, 4, 8, 16, 32, 48}.The table presents the full recall and DCO results for every listed nprobe value.
  • DCO definitions: For single-assignment methods, DCO equals nprobe × N/K.This definition uses the number of probed clusters and dataset size relative to K.
  • DCO definitions: For GMM multi, DCO equals nprobe × N m̄/K, with m̄ values of 1.49, 1.60, and 1.78 for SIFT1M, Deep10M, and GloVe-100.The dataset-specific m̄ values account for multi-assignment in the DCO calculation.

E Hard Multi-Assignment Ablation

This ablation separates GMM responsibilities from redundant assignment by comparing GMM multi-assignment with k-means hard top-2. Hard top-2 is less DCO-efficient than FAISS single assignment, showing that distance-based redundancy alone provides marginal recall gains.

  • Method: Kmeans hard top-2 assigns each vector to its two nearest k-means centroids by L2 distance, using fixed m̄ = 2.0 for every vector.The assignment count is independent of boundary proximity.
  • Results: Up to 1.8× more DCO is required by Kmeans hard top-2 than FAISS single to reach the same recall target.Table 8 evaluates recall@10 targets of 0.90 and 0.95 at K = 1024.
  • Interpretation: The second nearest centroid by L2 distance rarely contains the true nearest neighbor when the first does not, so doubling index size yields only marginal recall gains.Hard top-2 is strictly less DCO-efficient than FAISS single at both reported targets.

F IVF-PQ Compatibility and RAIRS Comparison

Flash-GMM integrates with IVF-PQ without changing product-quantization training or encoding, while improving recall through multi-assignment. In the RAIRS-matched IVF-PQfs configuration, GMM-based assignments are competitive with approximately 30% fewer probes and 25% less index inflation, although exact head-to-head comparison is unavailable.

  • IVF-PQ Compatibility: Flash-GMM is fully compatible with IVF-PQ, substituting coarse centroids without modifying PQ training or encoding.The evaluation covers IVF-PQ with M ∈ {16, 8, 4}, 8-bit codes, and IVF-PQfs with M = 64, 4-bit fast-scan codes on SIFT1M.
  • IVF-PQ Compatibility: At IVF-PQ M = 16, multi-assign reaches recall ≥0.87 at nprobe=8 versus FAISS’s 0.830, a +5.8 pp gain.Gains of +2–4 pp persist at high compression, where PQ quantisation error dominates.
  • RAIRS Comparison: In the RAIRS-matched IVF-PQfs configuration, Flash-GMM improves centroids via soft EM, derives assignments from GMM responsibilities, and uses adaptive per-vector redundancy.RAIRS uses a geometric heuristic, while AIR universally assigns a second list; GMM assigns one only when r_ik > 1/K.
  • RAIRS Comparison: Approximately 30% fewer probes reach the same recall target, while index inflation is 25% lower: m̄ = 1.49 versus RAIRS’s universal m̄ = 2.0.Based on published recall–nprobe curves, GMM multi is competitive with RAIRS on the shared IVF-PQfs configuration.
  • RAIRS Comparison: An exact head-to-head comparison with RAIRS is unavailable without its code, making Table 9 the closest available proxy.The proxy uses the shared IVF-PQfs configuration: M = 64, 4-bit fast-scan, nlist=1024.
Loading 2606.10896v1…