HN 日本語サマリー

← 一覧へ戻る
AI・機械学習

E-means: 損失同等でk-meansより10倍高速

E-means: 10x faster k-means at equal loss (e-means.github.io)

465 pointsby porridgeraisin68 コメント

要約

E-meansは、k-meansクラスタリングを大幅に高速化する新しいアルゴリズムです。従来のk-meansよりも10倍から45倍速く、同等以上の精度（損失）を達成します。このアルゴリズムは、ミニバッチ更新と指数移動平均を組み合わせ、ハイパーパラメータチューニングを自動化し、ローカルミニマムからの脱出能力を高めることで、速度と品質の両方を向上させています。

全文翻訳

e-means: エポック数を1桁少なくしてk-meansを上回る

エポック数を1桁少なく

様々なスケールとデータセットで、e-meansは25エポックでLloydのテスト損失を下回る結果を、わずか1.2〜3エポックで達成します。エンドツーエンドで測定すると、最先端のLloyd実装（Flash-KMeans[[](https://e-means.github.io/#ref2))に対するGPUでのウォールクロック時間は10〜45倍、CPUでのFaiss[[](https://e-means.github.io/#ref1))に対する時間は8〜15倍短縮されます。通常、k-meansでは速度と品質はトレードオフの関係にありますが、e-meansは両方を向上させています。

［画像: e-meansとLloydのk-meansの比較、エポック数、CPU時間、GPU時間］

k-meansの問題

固定されたn個の点X = {x1, ..., xn} ⊂ Rdと整数kが与えられたとき、k-means問題は、各点とその最も近いセントロイドとの間の平均二乗距離を最小化するk個のセントロイドC = (c1, ..., ck)を見つけることです。

\[ \mathrm{MSE}(X, C) = \frac{1}{n} \sum_{i=1}^{n} |x_i - c extbackslash a(i)\textbar|^2, \qquad\quad a(i) = \arg\min_{j \in {1,\dots,k}} |x_i - c_j|^2 \]

\[ \mathrm{MSE}(X, C) = \frac{1}{n} \sum_{i=1}^{n} |x_i - c extbackslash a(i)\textbar|^2 \]

\[ a(i) = \arg\min_{j \in {1,\dots,k}} |x_i - c_j|^2 \]

これは静的なデータセットであり、分布が常に移動するVQ-VAEs内のコードブック学習とは異なります。

e-meansプロジェクト

e-meansの目標は、k-means問題のための最新の最適化ルールを提供することです。私たちは以下の観察からこのプロジェクトを開始しました。

  • k-meansは、機械学習においてフルバッチ最適化が依然として広く使用されている数少ない問題の1つです。
  • 文献にはミニバッチk-means手法が存在しますが[[](https://e-means.github.io/#ref3),[[](https://e-means.github.io/#ref4),[[](https://e-means.github.io/#ref5))、それらはほとんどが1967年にリリースされたMacQueenの更新ルールに基づいています。
  • 最新のディープラーニングは、ミニバッチ手法を強化できる最適化ツールを提供します。

e-meansでは、k-means問題に最も効率的であることがわかった指数移動平均[[](https://e-means.github.io/#ref6))を使用したミニバッチ更新を選択しました。以前のミニバッチアプローチとは異なり、e-meansはスケーリング則（スケーリング則による自動ハイパーパラメータチューニングを参照）によりバッチサイズとステップサイズを自動的に設定するため、ハイパーパラメータチューニングを必要としません。さらに、e-meansは初期化とデッドクラスタ復旧戦略を改善し、コサインスケジューリングされた減衰率を追加し、トレーニング終了時にわずかなランタイムペナルティでトレーニング損失を削減する最終化ステージを導入しています。e-meansは高度に並列化可能であり、そのアルゴリズム的利点をウォールクロック速度向上に変換する専用カーネルが付属しています。

移動平均更新

各ミニバッチは、各クラスタのカウントと合計を更新します。セントロイドはその比率です。

古いセントロイド新しいセントロイドバッチ平均

ハイパーパラメータチューニングなし

すべてスケーリング則から導出されます。

log n, log k, log T

開始と終了

新しい初期化および最終化手法。実行は各クラスタの正確な最終エポック平均で終了します。

開始正確な平均

コサインスケジューリング

ステップαは、学習率と同様にα0から0.1α0に減衰します。

α₀0.1 α₀

デッドクラスタ復旧

しきい値を下回るサイズのクラスタは、大きいクラスタを分割することによって復旧されます。

しきい値未満大きいクラスタ2つに分割

専用カーネル

GPUカーネルは最も近いセントロイド割り当てを融合し、上記のエポック節約を対応するウォールクロック速度向上に変換します。

通常のPyTorch融合カーネル

セントロイド更新ルール

セントロイド更新ルールはe-meansの中心的なコンポーネントです。このセクションでは、k-meansクラスタリングの更新ルールの概要を提供することを目的としています。すべての手法は、各クラスタjについて、それに割り当てられた点の合計SjとカウントNjを維持し、その比率をセントロイドとします。

\[ c_j = S_j / N_j \]

それらは、SjとNjがどのように更新されるかのみが異なります。以下では、sjとnjは現在のバッチでクラスタjに割り当てられた点の合計とカウント（いずれもゼロでない場合）を示します。

Lloydのアルゴリズム[[](https://e-means.github.io/#cite7))は、すべての点を最も近いセントロイドに割り当て、データセット全体をバッチとして使用して統計を再計算します。

\[ S_j \leftarrow s_j, \qquad\quad N_j \leftarrow n_j \]

したがって、cjは現在それに割り当てられている点の正確な平均です。各更新はデータセット全体を1回通過するコストがかかり、この手順は固定回数（通常20〜50イテレーション）繰り返されます。

MacQueenの更新[[](https://e-means.github.io/#cite3))は、統計を破棄することなく累積します。

\[ S_j \leftarrow S_j + s_j, \qquad\quad N_j \leftarrow N_j + n_j \]

したがって、cjは、以前の、おそらく不十分なセントロイド位置の下で割り当てられた点を含む、トレーニング開始以来クラスタjに割り当てられたすべての点の均一な平均です。これは後にStaleness-Reduction Mini-Batch K-Means[[](https://e-means.github.io/#ref8))によって部分的に対処されました。これはk-meansとMacQueenのハイブリッド手法です。彼らは重要な速度向上を得ましたが、クラスタ数が非常に少ない場合に限られました。

e-meansはミニバッチで動作し、この合計を指数移動平均に置き換えます。

\[ S_j \leftarrow (1-\alpha),S_j + \alpha, s_j, \qquad\quad N_j \leftarrow (1-\alpha),N_j + \alpha, n_j \]

\[ S_j \leftarrow (1-\alpha),S_j + \alpha, s_j \]

\[ N_j \leftarrow (1-\alpha),N_j + \alpha, n_j \]

したがって、古い割り当てはレートαで忘れられ、セントロイドは履歴全体ではなく最近の割り当てを追跡します。この更新ルールは通常、移動する分布に使用されるため、VQ-VAEsのトレーニングに使用されます[[](https://e-means.github.io/#ref6))。私たちが発見したのは、ハイパーパラメータが適切に調整されていれば、静的データセットに対しても非常に効率的であるということです。

指数移動平均更新ではなく勾配降下法を使用するのはなぜですか？

e-meansプロジェクトは、Meta FAIRのHugo TouvronとHervé Jégouによる以前の研究から始まりました。彼らは、ミニバッチk-meansを勾配降下法で行うとLloydのアルゴリズムを上回ることができることを見出しました。これらの結果を再現できましたが、指数移動平均更新とは異なり、勾配降下法ベースのセントロイド更新のパフォーマンスは、大きな学習率で不安定であることがわかりました。EMA更新は、セントロイドが常に割り当てられた点の凸包内にあることを保証するのに対し、勾配降下法更新は、その凸包を通り過ぎてオーバーシュートする可能性があるためだと推測しています。

［画像: 指数移動平均と通常の勾配降下法における、大きなバッチでの有効ステップに対するテスト誤差］

大きなバッチでのセントロイドが取る有効ステップに対するテスト誤差。指数移動平均更新（EMA）と通常の勾配降下法（SGD）について、試行された各ステップ設定で測定。DEEP、100万ベクトル、およびSIFT、1000万ベクトル、16,384クラスタ、5エポック、3シードの中央値。縦軸はパネル内の最良テスト誤差をパーセントで示したもの。最良設定では、勾配降下法は移動平均よりも1.6%（DEEP）および0.4%（SIFT）上回って終了します。

固定エポック数で最低損失を達成するために、指数移動平均更新を選択しました。これが、私たちの手法にe-meansという名前が付いた理由です。

Lloydはe-meansの特殊ケース：ハイパーパラメータ研究

バッチサイズをデータセット全体B = nとし、α = 1と設定すると、NjとSjはクラスタ全体のカウントと合計に正確になります。したがって、1回のe-meansイテレーションは正確に1回のLloydイテレーションになります。Lloydは、学習率のコサインスケジューリングがないe-meansの極端なケースです。しかし、私たちのハイパーパラメータ探索では、この極端な設定は非常に最適ではないことが判明しました。これがe-meansがLloydを大幅に上回る理由を説明しています。以下では、すべてのバッチサイズBとステップαに対する損失を測定しています。Lloydは、損失が最も低い領域の外にあるコーナーに着地します。

［画像: 5エポックでの各バッチサイズと定数ステップに対する検証損失。最も低い損失は対角線状の帯にあり、Lloydはフルバッチとステップ1で、最良値より1.8%上の右上隅に位置しています］

5エポック後、各バッチBとステップαに対する検証損失（最良セルとの相対値）。DEEP、n = 106、k = 1024。ステップはここで一定（コサインスケジューリングなし）であるため、Lloydは正確にセルの一つになります：フルバッチ、α = 1。最も低い損失は対角線状の帯に沿っており、バッチサイズが大きいほどステップサイズも大きくなる必要があります。

このハイパーパラメータ研究は、最適なバッチサイズBとステップサイズαを自動的に選択するスケーリング則を導出する動機となりました。

e-meansはローカルミニマムからより良く脱出する

e-meansが収束時にLloydよりも優れた損失を達成する理由をより深く理解するために、ローカルミニマムからの脱出能力を評価するために、トイデータセットでいくつかのトレーニングをプロットしました。200シードでトレーニングを繰り返した結果、e-meansはLloydと比較して病的な構成に陥る可能性がはるかに低いことがわかりました。

サイズが大きく異なる2つのブロブ、k = 8。Lloydは小さいブロブに2つのセントロイドを配置し、大きいブロブは不足させます。e-meansは同じトラップから開始し、最初の数エポックで脱出します。最終的に大きいブロブに7つ、小さいブロブに1つのセントロイドが配置されます。200回の開始のうち、Lloydは26%でスタックしたままですが、e-meansは3.5%の時間でスタックしたままです。

3x3グリッドに9つのブロブがあり、角が重い、k = 10。Lloydは中央のブロブにセントロイドを配置せず、3つを角に集中させます。e-meansは同じ配置に陥り、その後約10エポックで抜け出します。200回の開始のうち、Lloydは43%で1%以上悪くスタックしたままですが、e-meansは4%で1%以上悪くスタックしたままです。

正方形の周りに重い中央ブロブがある4つのブロブ、k = 5。Lloydは1つのコーナーブロブに2つのセントロイドを配置し、他の2つの間のギャップに1つを配置し、それらが共有します。e-meansは同じ開始から、各ブロブに1つずつ割り当てます。200回の開始のうち、Lloydは15.5%でスタックしたままですが、e-meansは0%でスタックしたままです。

直線上に5つのブロブ、k = 5。Lloydは軽い中央ブロブをセントロイドなしで残し、端のブロブに2つを与えます。e-meansは同じ開始から、各ブロブに1つずつ割り当てます。200回の開始のうち、Lloydは46%でスタックしたままですが、e-meansは0%でスタックしたままです。

その理由として、ミニバッチは最適化ランドスケープをより探索しやすく、ローカルミニマムから脱出しやすい、よりノイズの多い更新を提供することだと考えています。常に損失が減少するLloydは、見つかったローカルミニマムを悪用し、それらから決して脱出しない可能性が高くなります。

スケーリング則による自動ハイパーパラメータチューニング

e-meansの重要な特性は、術者がLloydのアルゴリズムと比較して追加のハイパーパラメータを設定する必要がないことです。私たちは、任意の問題の特性（n、k、T）からバッチサイズBとステップサイズα0を設定するために2つのスケーリング則を適合させ、それらが未知のデータセットと問題スケールに対して強い一般化を示すことを発見しました。

\[ B = c_B , k^{b_k} , n^{b_n} , T^{b_T} \]

\[ \alpha_0 = \min!\Big(c_\alpha , k^{a_k} , B^{a_B} , n^{a_n} , T^{a_T},; 1\Big) \]

法則定数指数（後悔、保持データセット、保持サイズ）R²

k

B

n

T

ステップ α0

0.824

−0.46

0.90

−0.45

−0.11

0.07%

0.07%

0.95

batch B, default

0.035

0.13

n/a

0.88

0.81

0.19%

0.20%

0.96

batch B, quality law

0.0044

0.13

n/a

0.88

0.81

0.07%

0.06%

n/a

後悔 = (法則によるMSE - グリッド全体の最良MSE) / 最良MSE、構成要素の中央値。保持データセット：SIFT、SPACEV、およびVIBE[[](https://e-means.github.io/#ref9))データセットarXivおよびAG-News。保持サイズ：適合から除外され予測されたnの値。

バッチを選択するために2つの式を提案します。「デフォルト法則」と「品質法則」です。「品質法則」は、スループットを犠牲にして、より優れたクラスタリング品質を「デフォルト」法則に対して提供します。以下に、スケーリング則のトレーニングデータをプロットします。検証MSE対ステップサイズα0（左）およびバッチサイズB（右）。

［画像: deep1bに適合した2つのスケーリング則］

デッドクラスタ復旧の重要性

デッドクラスタ復旧の重要性を示すトイ例。

ミニバッチk-meansの主な問題は、「デッド」クラスタ、つまり最適化中に割り当てられた点がゼロになるセントロイドが発生しやすいことです。これはコードブックの実効サイズを縮小させ、特に今日ますます使用されている大規模なコードブックでは問題となります。

一般的な緩和策は、EnCodec[[](https://e-means.github.io/#ref10))で行われているように、特定のサイズしきい値を下回ったクラスタを再初期化することです。そこでは、デッドセントロイドは単に入ってくるバッチからランダムにサンプリングされたサンプルに置き換えられます。これは役立ちますが、問題を完全に解決するわけではありません。バッチから任意の点をサンプリングする代わりに、既存の大きいクラスタを分割することによってデッドクラスタを復旧します。アクティブなクラスタをそのサイズに比例する確率でサンプリングし、それを2つに分割します。Faiss k-meansから適応されたこのアプローチは、デッドクラスタを排除し、コードブック容量を維持する上で大幅に効果的であることがわかりました。

データセット

復旧ポリシー

デッドクラスタ

コードブックの

不均衡

テストMSE

DEEP

比例分割

47

0.3%

1.194

0.38762

batchから再サンプリング（EnCodec）

2,343

14.3%

1.423

0.39124

オフ

4,313

26.3%

1.647

0.39448

SIFT

比例分割

10

0.1%

1.186

0.19623

batchから再サンプリング（EnCodec）

750

4.6%

1.233

0.19684

オフ

2,080

12.7%

1.404

0.19797

100万ベクトル、k = 16,384、5エポック、5シードの中央値。復旧ポリシーのみが異なります。クラスタは、トレーニング終了時に10未満のベクトルが割り当てられた場合、平均クラスタの千分の一としてデッドとカウントされます。不均衡：クラスタサイズの不均一さ（1 =すべて均等）。論文の復旧アブレーションテーブルより。

広範な実験的検証

24種類の異なる埋め込みおよび合成データセットと構成でe-meansを検証しました（論文を参照）。ビデオについては、234のe-means構成をk-meansと比較して実行しました（16データセット、1億〜1億ポイント）。結果マップのすべての曲線を確認してください。

［画像: 25回のイテレーションでのLloydとの収束曲線］

トレーニング（上）とテスト（下）のMSE、25エポック後のLloydの値で割ったもの。ランダムベクトルから初期化。

［画像: GPUでのウォールクロック時間に対するテスト誤差、1000万ベクトル、k=65,536］

H100 1台でのウォールクロック時間に対するテストMSE、1000万ベクトル、65,536クラスタ。通常のPyTorch Lloydの25エポック（GPU上のデータ、bf16）は142〜155秒かかりますが、Flash-KMeans[[](https://e-means.github.io/#ref2))では14〜22秒です。曲線上のラベルはエポック数を示します。

［画像: Lloydの25エポックテスト誤差にe-meansが必要とするエポック数、11 VIBE埋め込みデータセット：1.6〜2.4］

Lloydが25エポックで達成するテスト誤差にe-meansが必要とするエポック数。VIBE埋め込みベンチマークの11データセット（テキスト、コード、画像、マルチモーダル埋め込み、d = 200〜3,072、k = 8,192または16,384）。予算は0.2エポック間隔、5シードの中央値。すべてのデータセットで1.6〜2.4エポックが必要であり、Lloydの25エポックと比較されます。

ライブラリとコード

from faiss.contrib.e_means import Emeans # NumpyEmeans for NumPy
em = Emeans(d, k, nepochs=T) # T epochs (5 by default)
centroids, labels = em.train(x) # x: a PyTorch tensor (n, d)
# labels: each vector's cluster in the last epoch

リポジトリ: faiss/contrib/e_means。

この研究の引用

e-meansを使用する場合は、論文（arXiv:2611.09501）を引用してください。

@misc{bessa2026emeansoutperformingkmeansin,
      title={e-means: outperforming k-means (in one-tenth the epochs)},
      author={Swann Bessa and Pierre-Emmanuel Mazaré and Gergely Szilvasy and Hugo Touvron and David Picard and Matthijs Douze and Hervé Jégou},
      year={2026},
      eprint={2611.09501},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2611.09501},
}

参考文献

  1. M. Douze, A. Guzhva, C. Deng, J. Johnson, G. Szilvasy, P.-E. Mazaré, M. Lomeli, L. Hosseini, H. Jégou. "The Faiss library." arXiv:2401.08281, 2024. ↑
  2. S. Yang, H. Xi, Y. Zhao, M. Li, X. Fan, J. Zhang, H. Cai, Y. Lin, X. Li, K. Keutzer, S. Han, C. Xu, I. Stoica. "Flash-KMeans: Fast and Memory-Efficient Exact K-Means." arXiv:2603.09229, 2026. ↑
  3. J. MacQueen. "Some Methods for Classification and Analysis of Multivariate Observations." Proceedings of the Fifth Berkeley Symposium on Mathematical Statistics and Probability, 1:281–297, 1967. ↑
  4. D. Sculley. "Web-Scale K-Means Clustering." WWW, 2010. ↑
  5. J. Newling, F. Fleuret. "Nested Mini-Batch K-Means." NeurIPS, 2016. ↑
  6. A. van den Oord, O. Vinyals, K. Kavukcuoglu. "Neural Discrete Representation Learning." NeurIPS, 2017. ↑
  7. S. Lloyd. "Least squares quantization in PCM." IEEE Transactions on Information Theory, 28(2):129–137, 1982 (circulated as a Bell Labs technical note since 1957). ↑
  8. X. Zhu, J. Sun, Z. He, J. Jiang, Z. Wang. "Staleness-Reduction Mini-Batch K-Means." IEEE Transactions on Neural Networks and Learning Systems, 35(10):14424–14436, 2024. ↑
  9. E. Jääsaari, V. Hyvönen, M. Ceccarello, T. Roos, M. Aumüller. "VIBE: Vector Index Benchmark for Embeddings." arXiv:2505.17810, 2025. ↑
  10. A. Défossez, J. Copet, G. Synnaeve, Y. Adi. "High Fidelity Neural Audio Compression." Transactions on Machine Learning Research, 2023. arXiv:2210.13438. ↑