Meta FAIRがe-meansを公開、k-meansを10分の1のエポック数で上回るアルゴリズム

k-meansは、データをグループに分割し、各グループをセントロイドと呼ばれる中心で表現するクラスタリング手法である。ベクトルインデックス、量子化器、データ前処理パイプラインを構築する者にとって、この計算はしばしばオフライン段階で最も遅いステップであり、e-meansがまさにこれをターゲットとしている。

完全再計算ではなく移動平均

k-meansのすべての変種は、各クラスターについて合計SとカウントNを保持し、セントロイドはSをNで割った商である。主な違いはSとNの更新方法にある。Lloydのアルゴリズムは各イテレーションで全データからこれらの値を再計算する。1967年のMacQueenの手法は、バッチの統計量を継続的に加算するため、古い割り当てが消えることはない。

e-meansは中間の道を採用する：古い統計量を減衰率αで減衰させ、式S ← (1−α)S + αsに従って更新し、Nについても同様とする。これにより古い割り当ては徐々に消え、セントロイドは最近のデータに追従する。Lloydのケースはα=1のフルバッチに相当する。データセットDEEPでは、このケースは達成可能な最良の損失から1.8%上回る位置にある。

著者らは従来の勾配降下法もテストした。最適な設定でも、DEEPではEMAより1.6%、SIFTでは0.4%高い損失となった。EMA法が優れている理由として、セントロイドが常に割り当てられた点の凸包の内部に留まるためではないかと推測している。

スケーリング則で固定されるパラメータ

2つの経験的に調整されたスケーリング則が、データセットサイズ、クラスター数、エポック数からバッチサイズBと初期ステップα0を決定する。除外されたデータセットでは、ステップ則の後悔の中央値は0.07%に達した。ステップはその後、α0から0.1α0までコサインスケジュールに従う。さらに2つのメカニズムが仕組みを補完する：「死んだ」クラスター（サイズが閾値を下回り、大きなクラスターを分割して再起動される）の復活と、最終エポックの正確な平均による仕上げである。

著者らによると、発表された性能は、GPU上でFlash-KMeansより10〜45倍高速で、CPU上ではFaiss k-meansより8〜15倍高速である。VIBEベンチマークでは、11の埋め込みデータセットに1.6から2.4エポックを要した。大規模化の例として、100億の1024次元埋め込みからなるDINO10Bデータセットが、2.5エポックで1670万のクラスターに分割された。GPUでの結果はH100で得られた。1670万セントロイドの構成では、float32のアキュムレータが64GiBに達し、セントロイドを複数GPUに分散させる必要がある。

発表スレッドで公開されたシングルスレッドCPUでのデモでは、25エポックのLloydに4.39秒、2.2エポックのe-meansに0.43秒かかった。e-meansはMITライセンスの下、FAISS経由で配布されている。