초고속 K 평균으로 벡터 임베딩 색인 혁신

초록

SuperKMeans는 고차원 벡터 임베딩을 위한 특화된 K‑means 변형으로, 차원 프루닝과 적응형 샘플링을 결합해 CPU에서는 기존 FAISS·Scikit‑Learn 대비 최대 7배, GPU에서는 cuVS 대비 최대 4배 빠른 클러스터링을 구현한다. 또한 검색 리콜이 개선되지 않을 때 조기에 종료하는 “Recall 기반 조기 종료(ETR)” 기법을 도입해 품질 손실 없이 학습 시간을 단축한다.

상세 분석

본 논문은 대규모 고차원 임베딩 컬렉션을 색인하기 위한 K‑means 알고리즘의 구조적 한계를 면밀히 분석하고, 두 가지 핵심 혁신을 제시한다. 첫 번째는 차원 프루닝이다. 기존 K‑means는 전체 차원 d 에 대해 행렬곱(GEMM)으로 거리 계산을 수행하지만, 고차원에서는 대부분의 차원이 클러스터 할당 판단에 크게 기여하지 않는다. SuperKMeans는 전체 차원을 앞쪽 d′ (전체 차원의 약 12 %)와 뒤쪽 d″ (나머지)로 분리한다. 앞쪽 d′ 차원에 대해서는 고성능 BLAS의 SGEMM을 이용해 부분 거리 행렬을 빠르게 계산하고, 이 값을 기반으로 프루닝 단계에서 각 벡터‑센터 쌍에 대해 뒤쪽 d″ 차원의 계산 필요성을 판단한다. 여기서 사용된 적응형 샘플링(ADS) 기법은 초기 d′ 거리값이 충분히 큰 경우 뒤쪽 차원을 완전히 무시하도록 설계돼, 데이터 접근과 연산량을 크게 감소시킨다.

두 번째 혁신은 **Recall 기반 조기 종료(Early Termination by Recall, ETR)**이다. 전통적인 K‑means는 정해진 반복 횟수나 중심 변화량을 기준으로 종료하지만, 색인 품질은 실제 검색 리콜에 의해 결정된다. ETR은 클러스터링 과정 중 일정 주기로 현재 중심을 사용해 샘플 쿼리(사용자 제공 또는 데이터에서 추출)로 벡터 검색을 수행하고, 리콜이 일정 수준 이하로 향상되지 않으면 학습을 중단한다. 이 메커니즘은 클러스터링 품질과 검색 성능 사이의 직접적인 연관성을 활용해 불필요한 반복을 제거한다.

실험에서는 다양한 CPU 마이크로아키텍처(x86 Intel Granite Rapids, AMD Zen 5/3, ARM Apple M4, Graviton 4)와 GPU(NVIDIA CUDA) 환경에서 SuperKMeans를 FAISS와 Scikit‑Learn, cuVS와 비교하였다. 1 M 개의 1536‑차원 벡터에 대해 k = 10 k~100 k 범위에서 CPU에서는 최대 7배, GPU에서는 최대 4배의 속도 향상을 기록했으며, 리콜(Recall@10) 차원에서는 기존 방법과 통계적으로 유의미한 차이가 없었다. 또한, 계층적 K‑means와 결합했을 때 10 M 개의 1024‑차원 데이터셋을 250배 빠르게 색인할 수 있음을 보였다.

추가적으로 논문은 k‑means++ 초기화가 고차원 임베딩에서는 오히려 성능을 저하시킬 수 있다는 실험적 증거를 제시한다. 차원의 저주로 인해 Voronoi 경계가 흐릿해져 초기 중심의 품질 차이가 크게 작용하지 않으며, 무작위 초기화가 충분히 좋은 결과를 만든다. 또한, 전체 데이터의 20~30 %만 샘플링해도 고품질 중심을 얻을 수 있음을 확인했다. 이러한 통찰은 색인 파이프라인 설계 시 데이터 샘플링, 초기화, 반복 횟수 설정에 대한 가이드라인을 제공한다.

전반적으로 SuperKMeans는 알고리즘‑하드웨어 공동 설계 접근을 통해 기존 BLAS 기반 GEMM의 장점을 유지하면서 불필요한 연산을 제거하고, 검색 품질을 직접 모니터링하는 조기 종료 전략을 도입함으로써 실무에서 요구되는 고속 색인 구축을 실현한다.