K-meansのクラスタ数をExpected SSE Behaviorを用いて決定する

TL; DR

Stop using the elbow criterion for k-means (Schubert, 2022) で提案されている「クラスタ内分散SSEを使ったk-meansクラスタ数の決定アルゴリズム」をPythonで実装しました

エルボー法に頼ってはいけないよ

上記の論文では、エルボー法が提示するクラスタ数が多くのケースで最適でないことが、様々なデータを用いて検証されています。論文の中身はかのTJO大先生が既に解説してくださっているので、ここでは論文の3.3節で提案されている代替手法を実装してみます。

後にいくつかのデータで検証しますが、いわゆる"many blob dataset" (塊状のデータセット) では、エルボー法より優れたクラスタ数を提示してくれます。

Expected SSE behavior

クラスタ数を1ずつ増やしていったときに、クラスタ数kがそれまでのクラスタ数1,2,...,k-1よりも優れているかどうかを考えます。
具体的には、クラスタ数kでの最近傍クラスタ中心からの標準偏差の推定量 $\sqrt{SSE/(N-K)}$と、1,2,...,k-1のクラスタ数のうちもっとも良いクラスタ数から推定される期待値 $\sqrt{\hat{SSE}/(N-K)}$との比をとります。

\hat{SSE}_k:= \frac{N-k}{k}\underset{j=1,…,k-1}{\operatorname{min}}\frac{j}{N-j}SSE_j

\frac{\sqrt{SSE/(N-K)}}{\sqrt{\hat{SSE_k}/(N-K)}} = \sqrt{\frac{SSE}{\hat{SSE_k}}}

推定量の比は、クラスタ数kがそれまでのどのクラスタ数よりも良ければ減少するため、横軸にクラスタ数k、縦軸に$\sqrt{\frac{SSE}{\hat{SSE_k}}}$ をとり、その振る舞いを見ていくことで最適なクラスタ数を決定することができます。論文では、比が1を超える手前、または最小値をとるkが最適なクラスタ数として選ばれると書かれています。

実装

sklearn.datasets.make_blobsで生成したデータに対してエルボー法と上記の手法を使い、提案されたクラスタ数を比較します。

実行環境
Windows 10
Python 3.8.10
scikit-learn 1.2.2

demo.py

import pandas as pd 
import seaborn as sns
import matplotlib.pyplot as plt
from sklearn.datasets import make_blobs
from sklearn.cluster import KMeans

# クラスター用のデータを生成（2次元, クラスター数15）
x, y = make_blobs(
    n_samples=2000,
    n_features=2,
    centers=15,
    cluster_std = 0.05,
    center_box=[-1,1],
    shuffle=True,
    random_state=42
)

# 生成したデータをプロット
data = pd.DataFrame(data=np.concatenate((x, y.reshape(-1, 1)), axis=1), columns=["x", "y", "cluster"])
sns.scatterplot(data=data, x="x", y="y", hue="cluster", s=5, legend=None)

# k = 1,2,...,n-1まで探索
n = 30
sse, sse_ratio = np.zeros(n-1), np.zeros(n-1)
for k in range(1, n):
    kmeans = KMeans(n_clusters=k, n_init='auto')
    pred = kmeans.fit_predict(x)
    sse[k-1] = kmeans.inertia_

    # sqrt(SSE)比を計算
    sse_hat = (n-k)/k * min([j+1/(n-j+1) * sse[j] for j in range(k)])
    sse_ratio[k-1] = np.sqrt(sse[k-1] / sse_hat)


# SSEをプロット
plt.plot(sse)
plt.xticks(np.arange(0,n,5))
plt.xlabel('Number of clusters')
plt.ylabel('SSE')
plt.title('Elbow plot')
plt.show()

# sqrt(SSE)比をプロット
plt.plot(sse_ratio)
plt.xticks(np.arange(0,n,5))
plt.yticks(np.arange(0,max(sse_ratio),1))
plt.xlabel('Number of clusters')
plt.ylabel('SSE ratio')
plt.title('Standard Deviation Reduction Plot')
plt.show()

クラスター数15でやや重なりのあるデータが生成されました。このデータに対する最適クラスタ数の推定値は、エルボー法を用いた目検ではk=5付近となるのに対し、上述のクラスタ内標準偏差比を用いた方法ではk=14になります。元論文のFig.3では別の例も紹介されているので、ぜひ見てみてください。

注意

元論文のFig.4にはそもそもK-meansが有効でないデータ分布の例が載っています。今回実装した手法もこれらの分布に対しては無力です。また、高次元空間では次元の呪いをはじめ、2次元空間とはまた異なった振る舞いをするはずで、そこら辺の検証はまだなされていないようです。データはまず可視化して眺めてみるべきだということを今回改めて実感しました。

蛇足

元論文では新手法の名前はつけられていませんでしたが、1970年代にすでに提案されていた Variance Ratio Criterion (VRC) という手法に非常に類似しているらしいです（元論文の3.4節）。「70年代は多くの手法が提案されたが、ここで議論するには狭すぎる（意訳）」とも書いてあり、温故知新って大事なんだなあと思いました。