未来を形作るテクノロジーの深掘り記事。
English中文Españolالعربية日本語हिन्दीDeutschFrançaisPortuguês한국어ItalianoBahasa IndonesiaРусский

大規模K-meansの限界と次世代アルゴリズム

「K-means」は一見単純ですが、10億点規模になるとメモリ・速度・初期化の問題が一気に表面化します。最新アルゴリズムによる解決策を解説します。

Sarah Chen 9分で読めます 3,534 文字
暗い夜空の中で、いくつかの明るい灯台の周りを無数の色とりどりの点が群れをなして飛び交う様子

K-meansは、ML入門コースで誰もが学び、そのあとキャリアを通じて誤った使い方をし続けるアルゴリズムです。見かけは単純で、K個のクラスタ中心を選び、各点を最も近い中心に割り当て、中心を割り当てられた点の平均へ動かし、これを繰り返すだけです。擬似コードにすると3行で済みます。問題は、この単純さの裏に、実データを大規模に適用したときに爆発する地雷が潜んでいることです。

1万点ならK-meansはミリ秒で終わり、実装の詳細を気にする人はいません。1,000万点になると、初期化方法の選択だけで実行時間が100倍変わることもあります。10億点になると、標準的なアルゴリズムはメモリに収まらず、根本的に異なるアプローチが必要になります。最近のFlash-KMeansに関する論文はまさにこの問題に取り組み、正確なK-meansの結果を保ったまま、メモリ使用量を大幅に減らし、収束も高速化しています。大規模でK-meansが難しくなる理由と、現代的な派生手法がそれをどう解決するのかを見ていきましょう。

標準K-meansの実際の動作

「k-means」と言うとき一般に指されるのは、Lloydのアルゴリズムです。1回の反復には2つのステップがあります。割り当てステップでは、各データ点について全K個の重心との距離を計算し、最も近い重心に割り当てます。更新ステップでは、各重心を、そこに割り当てられた全点の平均として再計算します。

import numpy as np
def kmeans_lloyd(X, K, max_iter=100):
n, d = X.shape
# Random initialization (bad, but we'll fix this later)
centroids = X[np.random.choice(n, K, replace=False)]
for _ in range(max_iter):
# Assignment: O(n * K * d) — this is the bottleneck
distances = np.linalg.norm(X[:, None] - centroids[None, :], axis=2)
labels = np.argmin(distances, axis=1)
# Update: O(n * d)
new_centroids = np.array([
X[labels == k].mean(axis=0) for k in range(K)
])
if np.allclose(centroids, new_centroids):
break
centroids = new_centroids
return labels, centroids

割り当てステップのコストは、反復1回あたりO(n × K × d)です。nは点の数、Kはクラスタ数、dは次元数です。n = 10億、K = 1000、d = 128（現実的な埋め込みベクトルのクラスタリングのケース）の場合、反復1回あたり128兆回の浮動小数点演算になります。1テラFLOPSの計算能力があっても、反復1回に128秒かかり、K-meansは通常10〜50回の反復を要するので、合計で長時間待つことになります。

初期化問題

K-meansは1回目の反復を実行する前に、初期重心の位置を選ばなければなりません。この選択は多くの人が思っている以上に重要です。初期化が悪いと、最適解より任意に悪い解に収束してしまうことがあります。

ランダム初期化（初期重心としてK個のデータ点をランダムに選ぶ方法）は単純ですが、信頼性に欠けます。2つの初期重心がたまたま同じクラスタに落ちると、そのクラスタは代表されないままになります。アルゴリズム自体は収束しますが、準最適解に落ち着きます。K = 100では、少なくとも1回は衝突が起きる確率が高くなります。

K-means++（Arthur and Vassilvitskii, 2007）が標準的な解決策です。最初の重心はランダムに選び、以降の各重心は、既存の重心のうち最も近いものからの二乗距離に比例した確率で選びます。既存の重心から遠い点ほど選ばれやすいため、初期重心がデータ全体に散らばります。これによりO(log K)の近似保証が得られ、解が最適解のO(log K)倍以内に収まることが理論的に示されています。

def kmeans_plus_plus_init(X, K):
n, d = X.shape
centroids = [X[np.random.randint(n)]]
for _ in range(1, K):
# Distance from each point to nearest existing centroid
dists = np.min([
np.sum((X - c) ** 2, axis=1) for c in centroids
], axis=0)
# Sample proportional to distance squared
probs = dists / dists.sum()
next_idx = np.random.choice(n, p=probs)
centroids.append(X[next_idx])
return np.array(centroids)

ただし、K-means++の初期化そのものがO(n × K × d)です。データセット全体をK回走査する必要があるため、Kもnも大きいと、初期化だけでLloydアルゴリズムの数回分の反復より時間がかかることがあります。k-means||（Bahmani et al., 2012）のようなスケーラブルな派生手法は、候補点を並列にオーバーサンプリングし、それを最終的にK個の重心にまとめることで、走査回数を減らしています。

メモリの壁

標準K-meansでは、データセット全体を同時にメモリに載せる必要があります。割り当てステップでは、すべてのデータ点にアクセスし、すべての重心と比較しなければならないからです。128次元のfloat32ベクトルが10億個あると、データだけで512 GBが必要になり、ほとんどの単体マシンのメモリを超えてしまいます。

単純な解決策はミニバッチK-meansです。反復ごとにデータ点のランダムなサブセットをサンプリングし、そのサンプルに基づいて重心を更新します。これは機能しますが、収束は遅く、正確なK-meansとはやや異なる（通常はやや劣る）解に至ります。多くの用途では品質の低下は許容範囲内です。しかし、ベクトル検索における量子化コードブックのように、品質の差が重要になる用途もあります。

Flash-KMeansは別のアプローチを取ります。データセット全体をメモリに載せて処理するのではなく、賢いキャッシュを使ったストリーミング処理を行います。重要な洞察は、ほとんどの点が反復ごとにクラスタを移らない、ということです。ある点がクラスタ7にしっかり属していれば（重心7に、他のどの重心よりもずっと近い）、K個すべての距離を計算するのは無駄な作業です。距離の上下限を保持することで、Flash-KMeansは大半の点について、大半の反復で完全な距離計算を省略できます。

三角不等式による最適化

Elkanのアルゴリズム（2003）は三角不等式を使って、不要な距離計算を省略します。三角不等式とは、点Pから重心Aまでの距離は、Pから重心Bまでの距離とBからAまでの距離の和以下である、というものです。Pが現在重心Bに割り当てられていて、重心AとBの間の距離がわかっていれば、実際に距離を計算しなくても、AがPから遠すぎることを証明できる場合があります。

実際には、最初の数回の反復の後はほとんどの点がすでに正しい重心の近くにあるため、これによって距離計算の80〜95%が省略されます。残る距離計算は、クラスタの境界付近にある点、つまり実際に割り当てが変わりうる点のためのものだけです。

トレードオフとして、Elkanのアルゴリズムは距離の上下限を保存するために、O(n × K)の追加メモリを必要とします（各点から各重心への下限と、割り当て先の重心への上限）。Kが大きいと、このメモリコストは無視できないものになります。Hamerlyのアルゴリズムは、点ごとに下限を1つだけ保持するように変えることでメモリをO(n)に抑えますが、その代わり省略できる計算は少なくなります。

GPUによる高速化

K-meansの割り当てステップは、典型的な並列化しやすい処理です。各点の距離計算は互いに独立しているからです。そのためGPUによる高速化と相性が良いです。NVIDIAのcuMLライブラリとFacebookのFAISSは、どちらもGPU版のK-means実装を備えており、大規模データセットでCPU実装比10〜50倍の高速化を実現しています。

ただし、GPUのメモリには限りがあります。A100は80 GBのメモリを持ちますが、これは128次元ベクトルで約1億5,000万個分です。それより大きいデータセットでは、複数GPUへの分割か、データをチャンク単位で読み込んでGPUで処理し、結果をCPU上で集約するストリーミング方式が必要になります。

# Using FAISS for GPU-accelerated k-means
import faiss
import numpy as np
# 10 million 128-dimensional vectors
n, d, K = 10_000_000, 128, 1000
X = np.random.randn(n, d).astype('float32')
# CPU k-means (for comparison)
kmeans_cpu = faiss.Kmeans(d, K, niter=20, verbose=True)
kmeans_cpu.train(X)  # ~120 seconds
# GPU k-means (single GPU)
kmeans_gpu = faiss.Kmeans(d, K, niter=20, verbose=True, gpu=True)
kmeans_gpu.train(X)  # ~8 seconds — 15x faster

K-meansが不適切な場合

K-meansの実装を最適化する前に、そもそもK-meansが適切なアルゴリズムかを検討しましょう。多くの実データでは成り立たない強い仮定を置いているためです。

  • 球状のクラスタ。 K-meansは、クラスタがおおむね球状で、サイズも等しいと仮定しています。クラスタが細長かったり、不規則な形だったり、サイズが大きく異なったりすると、K-meansは大きなクラスタを分割し、小さなクラスタを統合してしまいます。ガウス混合モデル（GMM）は楕円形のクラスタを扱えます。DBSCANは任意の形状を扱えます。
  • Kが既知であること。 クラスタ数をあらかじめ指定しなければなりません。Kがわからない場合は、異なるK値でK-meansを複数回実行し、指標（シルエットスコア、エルボー法）で最良のものを選ぶ必要があります。試すK値の数だけ、総計算量が増えます。
  • ユークリッド距離。 K-meansはデフォルトでユークリッド距離を使います。テキスト埋め込みでは、通常コサイン類似度の方が適切です。ベクトルをL2正規化すれば（ユークリッド距離がコサイン距離と等価になる）回避できますが、忘れやすい点です。
  • 外れ値への感度。 どのクラスタからも遠く離れた外れ値1つでも、割り当てられた重心をそちらへ引き寄せてしまいます。平均の代わりに中央値を使うk-medoidsのような頑健な派生手法は外れ値に強いですが、計算コストは高くなります。

実践的なアドバイス

数千点から数十億点までのデータセットでK-meansを使ってきた経験から、特に重要だと感じていることは次のとおりです。

  1. 必ずK-means++で初期化する。 ランダム初期化は、リスクに見合うことはほぼありません。最終的なクラスタ品質の差は10〜30%になることも多く、K-means++は小〜中規模のKであればオーバーヘッドがほぼ無視できます。
  2. 複数回実行する。 K-meansは大域的最適解ではなく局所解を見つけます。異なる乱数シードで5〜10回実行し、最良の結果（クラスタ内距離の総和が最小のもの）を採用しましょう。悪い実行に備えた安価な保険になります。
  3. 特徴量を正規化する。 ある特徴量の値域が[0, 1000000]で、別の特徴量が[0, 1]だと、値域の大きい特徴量が距離計算を支配してしまいます。クラスタリングの前に、標準化（平均を引いて標準偏差で割る）またはmin-max正規化を行いましょう。
  4. 大規模クラスタリングにはFAISSを使う。 100万点を超えるなら、scikit-learnのK-meansは遅くなります。FAISSの実装は高度に最適化されており、GPUアクセラレーションもすぐに使えます。
  5. 探索的な作業では近似手法を検討する。 ミニバッチK-meansは正確なK-meansの10倍ほど高速で、探索には通常十分な結果が得られます。品質が重要な本番のクラスタリングには、正確なK-meansを使いましょう。

K-meansは、使うのは簡単でも上手く使うのは難しく、そして深く理解する価値のあるアルゴリズムです。素朴な実装と最適化された実装のあいだには、速度と結果の品質の両面で大きな差があります。小規模では、これらは何も問題になりません。問題になる規模では、初期化、メモリ管理、距離の枝刈り、GPUアクセラレーションを理解しているかどうかが、数時間かかるクラスタリングと数分で終わるクラスタリングの違いになります。