Hacker News 2026.03.21 43

Flash-KMeans: K-Means 클러스터링을 메모리 절반으로, 속도는 수십 배 빠르게

들어가며

클러스터링은 머신러닝에서 가장 기본적이면서도 가장 널리 쓰이는 비지도 학습 기법입니다. 그중에서도 K-Means는 1950년대에 제안된 이래 70년이 넘도록 실무에서 살아남은 알고리즘입니다. 추천 시스템에서 사용자 그룹을 나누거나, 이미지에서 대표 색상을 추출하거나, 벡터 데이터베이스에서 양자화(quantization)를 수행할 때 모두 K-Means의 변형이 쓰입니다. 하지만 데이터가 수억 건, 차원이 수백~수천에 달하는 현대 ML 워크로드에서 기존 K-Means 구현은 심각한 메모리 및 속도 병목에 부딪힙니다. 이 문제를 정면으로 해결하려는 새로운 논문 "Flash-KMeans"가 공개되어 소개합니다.

K-Means의 동작 원리와 기존 한계

먼저 K-Means가 어떻게 동작하는지 짚어봅시다. 알고리즘은 크게 두 단계를 반복합니다. 할당(assignment) 단계에서는 각 데이터 포인트를 가장 가까운 중심점(centroid)에 배정하고, 업데이트(update) 단계에서는 각 클러스터에 속한 점들의 평균으로 중심점을 갱신합니다. 이걸 중심점이 더 이상 크게 변하지 않을 때까지 반복하면 됩니다. 원리는 단순하지만, 할당 단계에서 모든 데이터 포인트와 모든 중심점 사이의 거리를 계산해야 하므로 데이터가 커지면 연산량이 O(n × k × d)로 폭증합니다. 여기서 n은 데이터 수, k는 클러스터 수, d는 차원 수입니다.

기존에 이 문제를 완화하기 위한 접근법은 크게 두 갈래였습니다. 하나는 근사(approximate) K-Means로, Mini-Batch K-Means처럼 데이터의 일부만 샘플링해서 처리하는 방식입니다. 속도는 빠르지만 결과의 정확도가 떨어집니다. 다른 하나는 삼각부등식(triangle inequality) 기반 가지치기로, 이전 반복의 정보를 활용해 불필요한 거리 계산을 건너뛰는 Elkan 알고리즘 같은 방법입니다. 정확한 결과를 보장하지만 중심점 개수 k에 비례하는 추가 메모리(n × k 크기의 경계값 테이블)가 필요해서, k가 큰 경우 메모리가 폭발합니다.

Flash-KMeans의 핵심 아이디어

Flash-KMeans는 이 두 가지 한계를 동시에 돌파합니다. 정확한(exact) K-Means 결과를 보장하면서도, 기존 가속 알고리즘 대비 메모리 사용량을 절반 이하로 줄이고 속도를 수십 배 향상시켰습니다.

핵심 기법은 크게 세 가지입니다. 첫째, 반지름 기반 필터링(radius-based filtering)입니다. 각 중심점에 대해 "이 반지름 안에 있는 점은 이 클러스터에 속한다"는 것을 보장하는 경계를 동적으로 계산합니다. 이렇게 하면 대부분의 데이터 포인트에 대해 모든 중심점과의 거리를 계산하지 않고도 올바른 클러스터를 확정할 수 있습니다.

둘째, 포인트별 경계값을 최소한으로 유지합니다. 기존 Elkan 알고리즘이 각 포인트에 대해 k개의 하한(lower bound)을 저장하는 반면, Flash-KMeans는 하한의 수를 1~2개로 줄이면서도 가지치기 효과를 유지합니다. 이것이 메모리 절감의 핵심입니다. k가 1,000이면 Elkan은 포인트당 1,000개의 float를 저장해야 하지만, Flash-KMeans는 2개면 됩니다.

셋째, 중심점 간 거리 정보의 효율적 활용입니다. 중심점끼리의 거리 행렬(k × k)은 비교적 작으므로 이를 적극 활용하되, 데이터 포인트 수준의 대규모 경계값 테이블은 생략하는 전략입니다.

성능은 어느 정도인가

논문에서 보고하는 벤치마크 결과는 상당히 인상적입니다. 다양한 데이터셋에서 기존의 Lloyd 알고리즘(가장 기본적인 K-Means) 대비 10~100배의 속도 향상을 달성했고, Elkan 알고리즘 대비로도 2~10배 빠르면서 메모리는 절반 이하를 사용합니다. 특히 k가 큰 경우(수백~수천 클러스터)에서 격차가 더 벌어집니다.

이 성능 차이가 왜 중요할까요? 실무에서 K-Means는 단독으로 쓰이기보다 더 큰 파이프라인의 일부로 쓰이는 경우가 많습니다. 예를 들어 벡터 검색에서 사용하는 Product Quantization(PQ)은 내부적으로 수천 번의 K-Means를 실행합니다. FAISS 같은 벡터 인덱스 라이브러리에서 IVF 인덱스를 구축할 때도 대규모 K-Means가 필수입니다. 이런 상황에서 K-Means 자체가 10배 빨라진다면 전체 파이프라인의 빌드 시간이 극적으로 줄어듭니다.

경쟁 기술 및 업계 맥락

현재 실무에서 가장 널리 쓰이는 K-Means 구현은 scikit-learn의 기본 구현, FAISS에 내장된 GPU 기반 K-Means, 그리고 cuML(RAPIDS)의 GPU 구현 정도입니다. scikit-learn은 Lloyd와 Elkan 알고리즘을 제공하지만 대규모 데이터에서는 느리고, FAISS와 cuML은 GPU를 활용해 속도를 높이지만 GPU 메모리 제약이 있습니다.

Flash-KMeans의 접근은 알고리즘 수준의 최적화라는 점에서 하드웨어 가속과 직교합니다. 즉, GPU 구현과 결합하면 추가적인 성능 향상을 기대할 수 있습니다. 이름에서 짐작할 수 있듯이 FlashAttention의 철학과 유사하게, 하드웨어를 바꾸는 것이 아니라 알고리즘의 메모리 접근 패턴을 최적화해서 성능을 끌어올리는 전략입니다.

한국 개발자에게 주는 시사점

벡터 검색, 추천 시스템, 이미지/텍스트 임베딩 처리 등 K-Means가 파이프라인에 포함된 프로젝트를 운영하고 있다면 주목할 만합니다. 특히 대규모 벡터 인덱스를 주기적으로 재구축해야 하는 서비스에서 빌드 시간 단축에 직접적인 효과가 있을 수 있습니다. 논문이 공개된 만큼 구현체가 오픈소스로 나올 가능성이 높고, 기존 라이브러리에 통합되는 것도 시간문제일 수 있습니다.

또한 이 논문은 "70년 된 알고리즘도 아직 최적화할 여지가 있다"는 점을 보여줍니다. 최신 딥러닝 아키텍처만 쫓기보다, 기본 알고리즘의 효율적 구현을 깊이 이해하는 것이 실무 성능에 큰 차이를 만들 수 있다는 교훈입니다.