e-means: エポック数を大幅に削減する k-Means 手法

概要と成果

**e-means（エー・ミーンズ）**は、広範なスケーリング範囲とデータセットにおいて以下の驚異的な性能向上を実現しています。

• テスト時の損失: ロイド法（Lloyd's method）が 25 エポックで達成した値の1/25 以下に低下させることが可能です。
• 必要なエポック数: この高い精度を達成するために必要なトレーニングは、わずか1.2〜3 エポックです。
• 推論時間の削減率:
  • GPU 環境（Flash-KMeans[2] vs e-means）: 10 倍〜45 倍の高速化。
  • CPU 環境（Faiss[1] vs e-means）: 8 倍〜15 倍の高速化。
• 品質と速度の両立: k-means において通常は避けられない「スピード vs 品質」のトレードオフを解消し、両者を同時に最適化しています。

---

k-Means の課題

固定された $n$ 個のポイントからなるデータセット $X = {x_1, ..., x_n} \subset R^d$ と整数 $k$ が与えられたとき、k-means 問題とは以下の最適化です。

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

これは、VQ-VAE 内部で行われる動的なクラスタリングとは異なり、静的なデータセットにおける問題です。

---

e-means プロジェクトのコンセプト

e-means の目的は、k-means 問題に対する現代的な最適化ルールを提供することにあります。以下の洞察に基づいています：

• バッチ処理の必要性: k-means は依然としてバッチ全体（full-batch）を使用する最適化手法が主流です。
• 既存手法の限界: 文献にあるミニバッチ手法[3,4,5]は、1967 年の MacQueen の更新ルールに基づいており、現代の深層学習ツールを十分に活用していません。
• EMA の採用: e-means は、最も効率的な指数平滑移動平均（EMA）[6] を用いたミニバッチ更新を採用しました。

メインの特徴

• 自動ハイパーパラメータ設定: スケーリング則によってバッチサイズとステップサイズを自動決定するため、追加のチューニングは不要です。
• 高度な最適化戦略: 初期化手法、「デッドクラスタのリバイバル」、コサインスケジュールによる学習率減衰、トレーニング終了直前の低コスト・高精度化（フィナライズステージ）を実装。
• 専用カーネル: アルゴリズム上の優位性を GPU カーネルで融合し、ウォールクロック時間の高速化に直接変換します。

主要構成要素詳細

• 移動平均による更新: ミニバッチごとに各クラスタの和とカウントを更新。重心はこれらの比（厳密には指数平滑値）となります。

  • 旧重心

    $\to$

    新重心

    =

    バッチ平均

    （$\alpha$ 重み付き）
• ハイパーパラメータなし: スケーリング則（$\log n, \log k, \log T$）のみで設定が決まります。
• 開始と完了: 正確な初期化と、各クラスタの最終平均に基づいた終了判定（

  start

  =

  exact mean

  ）を採用。
• コサインスケジューリング: ステップサイズ $\alpha$ を $0.1\alpha_0$ までアニーリングします（$\alpha \to 0.1\alpha$）。
• デッドクラスタのリバイバル: 閾値未満のクラスタを、既存の大規模クラスタを分割して復活させます。

  • 閾値未満

    +

    大クラスタ

    $\to$

    2 つに分割

• 専用カーネル（フュージョン）: GPU で「割り当て計算」と「重心更新」を融合し、PyTorch の単純実装よりも高速化を実現します。

---

重心の更新ルール

e-means の核心は各クラスタ $j$ の統計量 $S_j$（和）と $N_j$（カウント）の維持と更新です。重心 $c_j$ は比として算出されます。

[ c_j = S_j / N_j ]

以下に主要な更新ルールの比較を示します（$s_j, n_j$ はバッチ内のポイント）。

1. ロイド法 (Lloyd's Method)

データセット全体を 1 パス分走って統計量再計算します。

• 更新: $S_j \leftarrow s_j, \quad N_j \leftarrow n_j$
• 特徴: その時点で割り当てられたポイントの正確な平均。固定回数（通常 20〜50）を反復します。
• 欠点: データ全パスのコストがかかる。

2. MacQueen の更新

統計量を累積し、決して破棄しません。

• 更新: $S_j \leftarrow S_j + s_j, \quad N_j \leftarrow N_j + n_j$
• 特徴: トレーニング開始からのすべてのポイントの平均（ユニフォーム平均）。初期位置の影響が残ります。

3. e-means (指数平滑移動平均)

ミニバッチを使用し、EMA で更新を行います。

• 更新: $S_j \leftarrow (1-\alpha),S_j + \alpha, s_j, \quad N_j \leftarrow (1-\alpha),N_j + \alpha, n_j$
• 特徴: 古い割り当てを比率 $\alpha$ で忘れ、重心は最近の割り当てを追跡します。静的なデータセットに対しても極めて効率的です。

---

なぜ勾配降下法ではなく EMA なのか？

先行研究（Meta FAIR, Hugging Face）でミニバッチ k-means を勾配降下法で行った場合、ロイド法を上回る可能性が示されましたが、大規模な学習率では不安定であることが判明しました。

• 勾配降下法の弱点: 重心が割り当てられたポイントの凸包（convex hull）を超えて飛び出す可能性があります。
• EMA の強み: 常に凸包内に留まることを保証し、安定しています。
• 命名の由来: 固定されたエポック数で最も低い損失を達成するために EMA を採用したことにちなんで「e-means」と名付けられました。

スケーリング則による自動チューニング

ロイド法に追加するパラメータ設定は不要です。2 つのスケーリング則からバッチサイズ $B$ と初期ステップサイズ $\alpha_0$ を決定します：

[ 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) ]

• デフォルト法: 処理速度と品質のバランスを重視。
• 品質法: より高いクラスタリング品質を得ることを目的。

---

デッドクラスタのリバイバル戦略

ミニバッチ k-means の大きな課題は、割り当てポイントがゼロになり「デッド」状態になったクラスタです。これによりコードブックの実効サイズが縮小します。

• 従来の解決策: 入ってくるバッチからランダムサンプリングして再初期化（EnCodec[10] など）。→ 問題の完全な解消には至らない。
• e-means の解決策: 既存の大規模クラスタを分割することでデッドクラスタを復活させます。
  • 確率に比例してアクティブなクラスタを選択し、それを 2 つに分割。
  • Faiss に適応されたこのアプローチは、デッドクラスタの排除と容量維持において著しく効果的です。

---

網羅的な実験的検証結果

24 つの違いのあるデータセット（埋め込み・合成）および設定で検証済みです（詳細は論文参照）。

• 訓練・テスト MSE: ロイド法の 25 エポック値に対する比率。ランダムベクトルからの初期化で測定。
• H100 GPU 環境での性能比較 (1,000 万ベクトル / 65k クラスタ):
  • PyTorch ロイド法: 25 エポックで約 142〜155 秒。
  • Flash-KMeans: 約 14〜22 秒 (10〜7 倍高速)。
  • e-means: 僅か 1.6〜2.4 エポック でロイド法の同等な精度に到達。

---

ライブラリとコード例

faiss.contrib.e_means

を使用します（NumPy 版も利用可能）。

from faiss.contrib.e_means import Emeans  # NumPy 版: NumpyEmeans

# 初期化 (T エポックはデフォルトで 5)
em = Emeans(d, k, nepochs=T)     

# トレーニング (x: PyTorch テンソル shape [n, d])
centroids, labels = em.train(x)  

# results
# centroids: クラスタ中心
# labels: 最後に割り当てられたクラスタ ID

リポジトリ:

faiss/contrib/e_means

---

局所解からの脱出能力

e-means はロイド法よりも優れた損失値に達する理由の一つは、ミニバッチによるノイズが最適化ランドスケープを探索しやすくし、局所解から脱出しやすいことにあります。

ケーススタディ (k=8, 非常に異なる 2 つのクラスタサイズ)

• シナリオ: ロイド法は小さなくたまりに 2 つ重心を割り当ててしまい、大きなくたまりに不十分な残りとなる。
  • ロイド法: 約 26% の開始試行で停滞（局所解）。
  • e-means: コサインスケジューリングと EMA で初期の罠から数エポックで脱出し、7:1 の理想的な分配へ収束。約 3.5% の停滞のみ。

その他のケーススタディ

• 角に重み付けされた 3x3 グリッド (k=10):
  • ロイド法: 約 43% の性能劣化（中央のクラスタに重心が割らない）。
  • e-means: エポック 10 で脱出し、約 4% の劣化のみ。
• 中心を囲む 4 つのクマダリ (k=5):
  • ロイド法: 約 15.5% の停滞（隙間に重心が配置される）。
  • e-means: 0% の停滞。各クラスタにちょうど 1 つずつ重心を割り当てる。
• 一直線上の 5 つのクマダリ (k=5):
  • ロイド法: 約 46% の停滞。
  • e-means: 0% の停滞。

結論: ロイド法は損失減少に貪欲になり抜け出せないのに対し、e-means は動的な更新により局所解を回避します。

---

引用文献

1. M. Douze et al., "The Faiss library." arXiv:2401.08281, 2024.
2. S. Yang et al., "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," 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 et al., "Neural Discrete Representation Learning," NeurIPS, 2017.
7. S. Lloyd, "Least squares quantization in PCM." IEEE Trans. Info. Theory, 1982.
8. X. Zhu et al., "Staleness-Reduction Mini-Batch K-Means," IEEE TNNLS, 2024.
9. E. Jääsaari et al., "VIBE: Vector Index Benchmark for Embeddings." arXiv:2505.17810, 2025.
10. A. Défossez et al., "High Fidelity Neural Audio Compression," TMLR, 2023.