대규모 K-Means: 왜 무너지고 그다음은 무엇인가

K-means는 입문 ML 강의에서 누구나 배우고, 그다음 커리어 내내 잘못 쓰게 되는 알고리즘입니다. 겉보기엔 아주 단순합니다. K개의 클러스터 중심을 고르고, 각 점을 가장 가까운 중심에 할당하고, 중심을 할당된 점들의 평균 위치로 옮기고, 이를 반복하면 됩니다. 의사코드로 세 줄이면 끝나죠. 문제는 이 단순함 뒤에 실제 데이터를 대규모로 돌리는 순간 터지는 지뢰가 여럿 숨어 있다는 것입니다.

점이 1만 개라면 k-means는 밀리초 단위로 끝나고 구현 세부 사항은 아무도 신경 쓰지 않습니다. 점이 1,000만 개가 되면 초기화 방법을 무엇으로 고르느냐에 따라 실행 시간이 100배까지 차이 날 수 있습니다. 10억 개가 되면 표준 알고리즘은 아예 메모리에 올라가지 않아서 근본적으로 다른 접근이 필요합니다. 최근 Flash-KMeans 논문은 바로 이 문제를 다룹니다. 메모리 사용량을 크게 줄이면서도 정확한 k-means 결과를 더 빠르게 수렴시키는 것이 목표입니다. 대규모에서 k-means가 왜 어려운지, 그리고 최신 변형들이 이를 어떻게 해결하는지 자세히 살펴보겠습니다.

표준 K-Means가 실제로 하는 일

흔히 'k-means'라고 할 때 가리키는 Lloyd 알고리즘은 반복마다 두 단계를 거칩니다. 할당 단계에서는 각 데이터 포인트에 대해 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

할당 단계의 비용은 반복당 O(n × K × d)입니다. 여기서 n은 점의 개수, K는 클러스터 개수, d는 차원 수입니다. n = 10억, K = 1,000, d = 128(현실적인 임베딩 클러스터링 시나리오)이라면 반복 한 번에 128조 번의 부동소수점 연산이 필요합니다. 1테라플롭스의 연산 성능으로도 반복 한 번에 128초가 걸리고, k-means는 보통 10~50번의 반복이 있어야 수렴합니다.

초기화 문제

k-means가 단 한 번의 반복도 돌기 전에 초기 중심점 위치를 골라야 합니다. 이 선택은 대부분의 사람이 생각하는 것보다 훨씬 중요합니다. 초기화가 나쁘면 최적해와 임의로 큰 차이가 나는 해에 수렴할 수 있습니다.

무작위 초기화(데이터 점 K개를 무작위로 골라 초기 중심점으로 사용)는 간단하지만 믿을 만하지 않습니다. 초기 중심점 두 개가 우연히 같은 클러스터에 떨어지면 다른 클러스터 하나는 대표되지 못합니다. 알고리즘은 수렴하지만 최적에 못 미치는 해에 도달하게 됩니다. K = 100이면 이런 충돌이 한 번 이상 생길 확률이 꽤 높습니다.

K-means++(Arthur와 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 외, 2012) 같은 확장 가능한 변형이 나왔습니다. 이 방식은 후보 점을 병렬로 과샘플링해서 통과 횟수를 줄인 뒤, 이를 K개의 중심점으로 압축합니다.

메모리의 벽

표준 k-means는 전체 데이터셋을 동시에 메모리에 올려야 합니다. 할당 단계에서는 모든 데이터 포인트에 접근해서 모든 중심점과 비교해야 하기 때문입니다. 128차원 float32 벡터 10억 개면 데이터만으로 512GB가 필요합니다. 대부분의 단일 머신이 가진 메모리보다 많은 양입니다.

가장 단순한 해결책은 미니배치 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 알고리즘은 점마다 하한 하나만 유지해서 메모리를 O(n)으로 줄입니다. 대신 가지치기로 건너뛸 수 있는 계산은 줄어듭니다.

GPU 가속

K-means의 할당 단계는 본질적으로 병렬 처리에 적합합니다. 각 점의 거리 계산이 서로 독립적이기 때문입니다. 그래서 GPU 가속과 잘 맞습니다. NVIDIA의 cuML 라이브러리와 Facebook의 FAISS 모두 GPU k-means 구현을 제공하며, 대규모 데이터셋에서 CPU 구현 대비 10~50배 빠른 속도를 냅니다.

문제는 GPU 메모리가 제한적이라는 점입니다. A100은 메모리가 80GB라서 128차원 벡터를 약 1억 5천만 개 정도 담을 수 있습니다. 그보다 큰 데이터셋은 여러 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는 클러스터가 대략 구형이고 크기가 비슷하다고 가정합니다. 클러스터가 길쭉하거나 모양이 불규칙하거나 크기 차이가 크면, k-means는 큰 클러스터를 쪼개고 작은 클러스터를 합쳐 버립니다. 가우시안 혼합 모델(GMM)은 타원형 클러스터를 다룰 수 있고, DBSCAN은 임의의 모양을 다룰 수 있습니다.
K를 미리 안다. 클러스터 개수를 미리 지정해야 합니다. K를 모른다면 서로 다른 값으로 k-means를 여러 번 돌리고, 실루엣 점수나 엘보우 방법 같은 지표로 최적값을 골라야 합니다. 시도하는 K 값의 개수만큼 전체 연산량이 곱절로 늘어납니다.
유클리드 거리. K-means는 기본적으로 유클리드 거리를 씁니다. 텍스트 임베딩에는 보통 코사인 유사도가 더 적합합니다. 벡터를 L2 정규화하면 유클리드 거리가 코사인 거리와 동등해지므로 이 문제를 우회할 수 있지만, 깜빡하기 쉽습니다.
이상치에 대한 민감성. 어떤 클러스터에서도 멀리 떨어진 이상치 하나가 배정된 중심점을 자기 쪽으로 끌어당깁니다. 평균 대신 중앙값을 쓰는 k-medoids 같은 강건한 변형이 이상치를 더 잘 다루지만, 비용이 더 큽니다.

실전 조언

수천에서 수십억 개의 점까지 데이터셋에 k-means를 써 보면서 가장 중요하다고 느낀 점들을 정리하면 다음과 같습니다.

항상 K-means++ 초기화를 쓰세요. 무작위 초기화는 위험을 감수할 가치가 없습니다. 최종 클러스터 품질 차이가 종종 10~30%에 달하고, K-means++는 작은 K와 중간 K에서 추가 비용이 거의 없습니다.
여러 번 실행하세요. K-means는 전역 최적이 아니라 지역 최적을 찾습니다. 서로 다른 무작위 시드로 5~10번 돌리고, 클러스터 내 거리 합이 가장 작은 결과를 고르세요. 나쁜 실행에 대비하는 값싼 보험입니다.
특성값을 정규화하세요. 한 특성의 범위가 [0, 1000000]이고 다른 특성의 범위가 [0, 1]이면, 큰 범위의 특성이 거리 계산을 지배하게 됩니다. 클러스터링 전에 표준화(평균 빼기, 표준편차로 나누기)나 min-max 정규화를 하세요.
대규모 클러스터링에는 FAISS를 쓰세요. 점이 100만 개를 넘는다면 scikit-learn의 k-means는 느릴 것입니다. FAISS 구현은 많이 최적화되어 있고 GPU 가속도 기본으로 지원합니다.
탐색적 작업에는 근사 방법을 고려하세요. 미니배치 k-means는 정확한 k-means보다 10배 빠르고, 탐색 목적으로는 대개 충분히 좋은 결과를 줍니다. 품질이 중요한 프로덕션 클러스터링에는 정확한 k-means를 쓰세요.

K-means는 쓰기는 쉽지만 잘 쓰기는 어렵고, 깊이 이해할 가치가 있는 알고리즘입니다. 순진한 구현과 최적화된 구현 사이의 격차는 속도와 결과 품질 양쪽 모두에서 엄청납니다. 작은 규모에서는 이런 것들이 아무 의미도 없습니다. 하지만 정말 중요해지는 규모에서는 초기화, 메모리 관리, 거리 가지치기, GPU 가속을 이해하고 있느냐에 따라 몇 시간 걸리는 클러스터링과 몇 분 만에 끝나는 클러스터링이 갈립니다.