기술 요약: 풀 정밀도에서의 인덱싱 중단: 벡터 임베딩을 위한 클러스터링 재고

문제 정의

근사 벡터 유사도 검색(VSS)은 대규모 고차원 벡터 컬렉션을 클러스터로 분할하기 위해 $k$-means와 같은 클러스터링 알고리즘에 크게 의존합니다. 이러한 클러스터는 파티션 기반 인덱스(예: Inverted File Index)의 진입점 역할을 하여, 전체 데이터셋을 스캔할 필요 없이 검색 지연 시간을 크게 줄여줍니다. 그러나 이러한 인덱스의 구축 과정은 중요한 병목 구간입니다. VSS 쿼리와 달리, 클러스터링은 전체 원시 벡터 컬렉션에 여러 번 접근해야 합니다. 이 과정은 메모리 집약적일 뿐만 아니라 연산 집약적이어서, 인덱스가 완전히 구축될 때까지 사용자가 컬렉션을 쿼리하는 것을 방해하곤 합니다.

현재 시스템들은 일반적으로 클러스터링 단계에서 벡터를 풀 정밀도(float32 또는 float16)로 인덱싱한 후, 센트로이드(centroid)가 계산된 다음에 양자화나 차원 축소를 적용합니다. 본 논문은 클러스터링을 위해 풀 정밀도 벡터를 사용하는 것이 과도하다고 주장합니다. 또한, 벡터 검색에서 널리 사용되는 근사 기법들을 클러스터링 이전에 적용함으로써, 결과적인 센트로이드의 품질을 희생하지 않으면서도 스토리지 풋프린트를 획기적으로 줄이고 인덱싱 파이프라인을 가속화할 수 있다고 제안합니다.

방법론

저자들은 차원 축소, 양자화, 차원 프루닝(dimension pruning)을 클러스터링 프로세스 이전에 원시 벡터에 적용하는 새로운 인덱싱 파이프라인을 제안합니다. 이 파이프라인은 다음 단계들로 구성됩니다:

1. 전처리 (근사 우선):

   • 차원 축소: 기법으로는 Johnson-Lindenstrauss 변환(JLT), 주성분 분석(PCA), 그리고 에너지(energy)가 앞쪽 차원에 자연스럽게 집중되는 Matryoshka 임베딩 활용 등이 있습니다.
   • 양자화: 논문은 전처리 클러스터링용으로 몇 가지 양자화 스킴을 적응시켰습니다:
     • SQ8/SQ4: 글로벌 바이어스(bias)와 스케일링을 통해 직접적인 정수 산술 거리 계산을 허용합니다.
     • LVQ4 (Local Vector Quantization): 저자들은 재구성(reconstruction)과 거리 계산을 융합한 L2 거리의 수학적 재정식화를 도출하여, 개별 벡터의 바이어스와 스케일을 처리하면서도 4비트 GEMM 커널을 사용할 수 있게 했습니다.
     • RabitQ: 무작위 회전과 부호 보존에 기반한 1비트 양자화 방법입니다. 저자들은 L2 거리에 대한 편향되지 않은 추정치(unbiased estimator)를 클러스터링 루프 내에서 작동하도록 적응시켜, 데이터 포인트를 1비트 코드로 인코딩하고 센트로이드를 SQ4로 인코딩합니다.
   • 차원 프루닝: 저자들은 일반 행렬 곱셈(GEMM) 루틴과 점진적 프루닝 커널을 인터리빙하는 SuperKMeans를 통합합니다. 이 기법은 앞쪽 차원을 사용하여 후보군을 필터링함으로써, 전체 정밀도 L2 계산에 필요한 횟수를 줄입니다.

2. 클러스터링:

   • 근사된 벡터를 대상으로 $k$-means 알고리즘(Lloyd's algorithm)을 실행합니다.
   • 센트로이드 초기화: 임베딩 데이터셋에서는 $k$-means++가 높은 실행 비용 대비 미미한 개선만을 제공한다는 점이 발견되었기에, 랜덤 샘플링을 사용합니다.
   • 할당(Assignment): 근사된 표현(예: 4-bit GEMM, RabitQ를 위한 FastScan)을 사용하여 거리 계산을 수행합니다.
   • 업데이트(Update): 할당된 벡터들의 평균을 구하여 센트로이드를 업데이트합니다. 양자화된 스킴의 경우, 디코딩과 평균 계산을 융합한 후 새로운 센트로이드를 재인코딩하는 과정을 포함합니다.

3. 인덱스 실체화 (Index Materialization):

   • 최종 센트로이드는 IVF 인덱스를 위해 float32 벡터로 실체화됩니다. 논문은 이들이 원시 벡터로부터 재계산되거나(사용 가능한 경우), 양자화된 코드로부터 유도될 수 있음을 명시합니다.

주요 기여

• 파이프라인 재설계: 저자들은 "클러스터링 후 양자화"에서 "양자화 후 클러스터링"으로 패러다임을 전환하여, 근사 기법이 인덱싱 파이프라인의 초기에 적용될 수 있음을 입증했습니다.
• 알고리즘 적응:
  • 클러스터링의 다대다(many-to-many) 거리 계산에 적합하도록 LVQ 및 RabitQ를 적응시켰으며, 여기에는 개별 벡터 파라미터와 1비트 코드를 처리하기 위한 특정 유도가 포함됩니다.
  • 풀 정밀도 L2 거리가 필요했던 기존의 한계를 극려하며, 차원 프루닝(dimension pruning)을 다양한 양자화 스킴(SQ, LVQ, RabitQ)과 결합한 SuperKMeans를 통합했습니다.
• 실증적 평가: 클러스터링 품질(Recall@k, WCSS), 스토리지 감소, 빌드 시간을 측정하기 위해 6가지 다양한 벡터 임베딩 데이터셋(Cohere, OpenAI, arXiv, Jina, MXBAI, ImageNet)에 대해 종합적인 평가를 수행했습니다.

결과

실험 결과, 클러스터링은 근사 기법에 매우 탄력적임이 드러났습니다:

• 품질 보존: 공격적인 근사 기법을 적용하더라도 최적에 가까운 클로스터링 품질을 유지합니다.
  • SQ8은 풀 정밀도 클러스터링과 동등한 성능을 보입니다.
  • LVQ4 및 RabitQ(1-비트 코드)는 이상적인(풀 정밀도) 결과의 1% 이내의 클러스터링 품질을 달고 있습니다.
  • PCA는 가장 효과적인 차원 축소 기법으로 확인되었습니다. 분산(variance)의 60~70%를 보존하는 것만으로도 최적에 가까운 품질을 유지하기에 충분합니다.
• 스토리지 감소: 제안된 파이프라인은 클러스터링 반복에 필요한 스토리지 풋프린트를 최대 60배(RabitQ와 PCA 사용 시)까지 줄입니다.
• 성능 향상:
  • 속도 향상: 클러스터링 속도가 크게 개선되었습니다. RabitQ와 SuperKMeans를 결합하면 Cohere 데이터셋에서 최대 17배의 속도 향상을 달성합니다.
  • 계층적 $k$-means: RabitQ와 결합했을 때, 계층적 클러스터링은 풀 정밀도 데이터에 대한 vanilla $k$-means 대비 최대 91배 빠른 속도를 보여줍니다.
  • 확장성: 파이프라인은 클러스터 수($k$)에 대해 하위 선형(sublinearly)으로 확장되며, 5,000만 개의 벡터를 20만 개의 클러스터로 1분 이내에 인덱싱할 수 있습니다.
• 트레이드오프(Trade-offs):
  • **PQ (Product Quantization)**는 잔차(residual)가 아닌 원시 벡터에 글로벌 코드북을 구축하기 때문에, 클러스터링 전에 적용할 경우 품질이 낮아지는 경 tendency가 있습니다.
  • RabitQ는 주변부 클러스터에서의 체계적인 재구성 편향(reconstruction bias)으로 인해 클러스터 크기의 불균형을 초래할 수 있으나, 이는 마지막 반복 단계에서 원시 벡터를 사용함으로써 완화할 수 있습니다.

의의 및 주장

본 논문은 풀 정밀도로 인덱싱하는 표준 관행이 불필요하며 비효율적이라고 주장합니다. 근사 기법을 파이프라인의 전면에 배치함으로써 시스템은 다음을 달 수 있습니다:

1. 획기적인 자원 절감: 메모리 풋프린트를 60배 줄이고 연산 시간을 최대 91배까지 단축합니다.
2. 인덱스 품질 유지: 생성된 센트로이드는 무시할 만한 수준의 성능 저하(보통 Recall 손실 1% 미만)로 벡터 검색을 안내합니다.
3. 통합된 표현 방식: 동일한 근사 표현을 클러스터링 단계와 최종 벡터 검색 인덱스 모두에 사용할 수 있어, 데이터 수집 과정을 간소화하고 인덱스 실체화 과정에서의 원시 데이터로의 왕복(round-trip) 필요성을 제거합니다.

저자들은 이러한 접근 방식이 컴퓨팅 시간이 비용에 직접적인 영향을 미치고 인입(ingestion) 지연 시간이 핵심적인 병목인 대규모 클라우드 벡터 시스템에 특히 중요하다고 결론짓습니다. 또한 향-후 연구로서 특화된 하드웨어(예: Intel AMX) 및 그래프 기반 할당 전략을 탐구할 것을 제안합니다.