From df314c245f11a65baa590eacec38297fdc0ff6e4 Mon Sep 17 00:00:00 2001 From: bits-bytes-nn Date: Tue, 10 Feb 2026 00:29:09 +0900 Subject: [PATCH] feat: Add paper review for 'Billion-scale similarity search with GPUs' --- ...llion-scale-similarity-search-with-gpus.md | 783 ++++++++++++++++++ 1 file changed, 783 insertions(+) create mode 100644 _posts/2026-02-10-billion-scale-similarity-search-with-gpus.md diff --git a/_posts/2026-02-10-billion-scale-similarity-search-with-gpus.md b/_posts/2026-02-10-billion-scale-similarity-search-with-gpus.md new file mode 100644 index 0000000..dd026ba --- /dev/null +++ b/_posts/2026-02-10-billion-scale-similarity-search-with-gpus.md @@ -0,0 +1,783 @@ +--- +layout: post +title: "Billion-scale similarity search with GPUs" +date: 2017-02-28 10:42:31 +author: "Facebook AI Research" +categories: ["Paper Reviews", "Training-&-Inference-Optimization"] +tags: ["GPU-k-Selection-Algorithm", "WarpSelect", "Billion-Scale-Similarity-Search", "Product-Quantization-on-GPU", "IVFADC-GPU-Implementation", "In-Register-Sorting-Networks", "Odd-Size-Merging-Networks", "Fused-Kernel-Design-for-k-NN", "Multi-GPU-Index-Sharding-and-Replication", "Register-File-Based-Data-Structures"] +cover: /assets/images/default.jpg +use_math: true +--- +### TL;DR +#### 이 연구를 시작하게 된 배경과 동기는 무엇입니까? + +현대 인공지능 시스템에서 이미지, 비디오, 텍스트와 같은 복잡한 데이터는 딥러닝 모델을 통해 고차원 실수 벡터로 변환되며, 이러한 벡터 표현들 사이의 유사도를 효율적으로 검색하는 것이 핵심 연산이 되었습니다. 예를 들어 CNN을 통해 생성된 이미지 임베딩은 50차원에서 1,000차원 이상의 벡터로 표현되며, 두 벡터 간의 거리가 가까울수록 원본 데이터가 유사하다고 해석할 수 있습니다. 그러나 십억 규모의 벡터 데이터셋에서 최근접 이웃을 찾는 작업은 연산 집약적이며, 기존의 CPU 기반 방법이나 근사 알고리즘들은 실시간 응용에 필요한 성능을 제공하지 못했습니다. 특히 $k$-NN 그래프 구축과 같은 대규모 유사도 검색 작업은 수백 대의 CPU 서버 클러스터에서도 수일 이상이 소요되는 비실용적인 상황이었습니다. + +GPU는 높은 산술 처리량과 메모리 대역폭을 제공하지만, 유사도 검색의 핵심 연산인 $k$-선택(k-selection)을 GPU에서 효율적으로 구현하는 것은 매우 어려운 문제였습니다. 기존의 GPU 기반 $k$-선택 알고리즘들은 다중 패스 처리로 인한 메모리 오버헤드, 데이터 의존적 메모리 이동, 워프 발산(warp divergence)으로 인한 성능 저하, 또는 과도한 동기화 오버헤드를 겪고 있었습니다. 더욱이 GPU의 메모리 계층 구조에서 레지스터 파일의 대역폭이 글로벌 메모리의 250배에 달함에도 불구하고, 이러한 극단적인 대역폭 우위를 활용하는 알고리즘 설계가 부재했습니다. Facebook AI Research는 이러한 문제를 인식하고, GPU의 하드웨어 특성을 깊이 있게 이해하면서도 십억 규모 데이터셋에 실용적으로 적용 가능한 유사도 검색 방법을 개발하고자 했습니다. + +#### 이 연구에서 제시하는 새로운 해결 방법은 무엇입니까? + +이 논문의 핵심 기여는 **WarpSelect**라는 새로운 GPU $k$-선택 알고리즘으로, 모든 중간 상태를 GPU의 가장 빠른 메모리인 레지스터 파일에 유지하면서 단일 패스로 $k$-선택을 수행합니다. WarpSelect는 각 워프(32개 CUDA 스레드의 벡터 단위)가 독립적으로 $k$-선택을 담당하며, 두 가지 핵심 데이터 구조를 사용합니다. 첫째, 각 레인(워프 내 개별 스레드)이 **스레드 큐**를 레지스터에 유지하여 새로 들어오는 값에 대한 1차 필터 역할을 하고, 둘째, 워프 전체가 공유하는 **워프 큐**를 레인-스트라이드 레지스터 배열로 구성하여 지금까지 관찰된 $k$개의 최소 원소를 유지합니다. 이 알고리즘은 **홀수 크기 병합 네트워크(odd-size merging network)**를 도입하여 임의 길이의 정렬된 배열을 효율적으로 병합하며, 바이토닉 정렬 네트워크의 변형을 활용하여 레지스터 내에서 병렬 정렬을 수행합니다. + +이 논문은 또한 정확 검색과 근사 검색 모두에 대한 최적화된 GPU 계산 레이아웃을 제시합니다. 정확 검색에서는 L2 거리 분해를 통해 거리 계산을 행렬 곱셈으로 환원하고, **커널 융합** 전략으로 거리 행렬의 각 항목을 읽으면서 즉석에서 $k$-선택을 수행하여 메모리 패스를 최소화합니다. 근사 검색에서는 **Product Quantization(PQ)** 기반의 IVFADC 인덱싱을 채택하되, **룩업 테이블** 메커니즘을 통해 거리 계산을 단순한 테이블 조회와 덧셈으로 축소합니다. 특히 거리 표현을 세 개의 항으로 분해하여, 쿼리에 독립적인 항은 사전 계산하고 역색인 리스트에 독립적인 항은 한 번만 계산함으로써 연산 비용을 극적으로 절감합니다. 다중 GPU 환경에서는 **복제(replication)**와 **샤딩(sharding)** 전략을 상호 보완적으로 활용하여 인덱스 크기와 쿼리 수에 따라 최적의 병렬 처리 방식을 선택합니다. + +#### 제안된 방법은 어떻게 구현되었습니까? + +WarpSelect의 구현은 GPU 아키텍처의 세부 사항에 깊이 있게 맞춤화되어 있습니다. 각 워프는 32개의 레인으로 구성되며, 각 레인은 자신의 스레드 큐를 레지스터에 유지합니다. 스레드 큐 길이 $t$는 $k$의 크기에 따라 결정되는데, 실험적으로 $k \leq 32$일 때 $t = 2$, $k \leq 128$일 때 $t = 3$, $k \leq 256$일 때 $t = 4$, $k \leq 1024$일 때 $t = 8$로 설정됩니다. 워프 큐는 레인-스트라이드 레지스터 배열로 구성되어, 레인 $j$가 원소 $\{a_j, a_{32+j}, \ldots, a_{\ell-32+j}\}$를 저장하고, 레지스터 $r$이 원소 $\{a_{32r}, a_{32r+1}, \ldots, a_{32r+31}\}$을 보유합니다. 이러한 배치는 워프 셔플 명령어를 통한 효율적인 데이터 교환을 가능하게 하며, **버터플라이 순열(butterfly permutation)**을 활용하여 병렬 병합을 수행합니다. + +정확 검색의 구현에서는 cuBLAS의 최적화된 GEMM 루틴으로 내적 $\langle x_j, y_i \rangle$을 계산한 후, 별도의 융합 커널에서 $\|y_i\|^2$ 항을 더하면서 동시에 WarpSelect를 수행합니다. 이 융합 전략으로 거리 행렬 $D'$에 대해 GEMM 쓰기와 $k$-선택 읽기의 단 2회 패스만 필요하며, 메모리 대역폭 한계에 근접한 성능을 달성합니다. 근사 검색에서는 각 쿼리에 대해 먼저 정확 검색 방식으로 $\tau$개의 가장 가까운 조대 중심점을 찾은 후, 해당 역색인 리스트들을 순회하면서 룩업 테이블 기반 거리 계산과 $k$-선택을 수행합니다. 룩업 테이블은 공유 메모리에 저장되어 빠른 접근을 보장하며, 각 부분 양자화기별로 256개의 거리값을 사전 계산합니다. 메모리 제약이 있는 경우 2-패스 $k$-선택 전략을 적용하여 병렬성과 메모리 사용량 사이의 균형을 달성합니다. + +다중 GPU 구현에서는 인덱스 크기에 따라 전략을 선택합니다. 인덱스가 단일 GPU 메모리에 들어가는 경우 복제 방식을 사용하여 동일한 인덱스를 여러 GPU에 복제하고 쿼리를 분산 처리하며, 거의 선형적인 속도 향상을 달성합니다. 인덱스가 단일 GPU 메모리에 들어가지 않는 경우 샤딩 방식을 사용하여 인덱스를 여러 GPU에 분산하고, 각 샤드가 전체 쿼리 집합을 처리한 후 부분 결과를 합친 뒤 추가적인 $k$-선택 라운드를 수행합니다. 두 전략은 결합하여 사용할 수 있으며, 이 원리는 다중 머신 분산으로도 자연스럽게 확장됩니다. + +#### 이 연구의 결과가 가지는 의미는 무엇입니까? + +실험 결과는 WarpSelect의 우수한 성능을 명확히 보여줍니다. 단일 Titan X GPU에서 배열 길이 $\ell = 128,000$, $k = 100$일 때 WarpSelect는 기존의 fgknn 라이브러리 대비 1.62배 빠르며, $k = 1000$일 때는 2.01배 빠른 성능을 달성했습니다. 특히 주목할 점은 WarpSelect가 이론적 최대 성능의 55%에서 동작하여, 메모리 대역폭 한계에 매우 근접한 효율을 달성했다는 것입니다. 정확 검색 벤치마크인 SIFT1M 데이터셋에서는 본 논문의 구현이 기존 방법 대비 최소 25% 이상 빠르며, 특히 커널 융합 전략의 중요성이 입증되었습니다. 근사 검색에서는 SIFT1B(10억 개 벡터) 데이터셋에서 동일한 메모리 사용량 조건에서 기존 최신 기법(Wieschollek et al.) 대비 8.5배 빠른 속도로 더 높은 정확도를 달성했습니다. + +가장 인상적인 응용 결과는 $k$-NN 그래프 구축입니다. Yfcc100M 데이터셋의 9,500만 장 이미지에 대한 고정확도 10-NN 그래프를 4개의 Maxwell Titan X GPU에서 35분 만에 구축했으며, 10억 개 벡터에 대한 그래프를 12시간 이내에 완성했습니다. 이는 기존에 128대의 CPU 서버 클러스터에서 3,650만 개 벡터에 대해 108.7시간이 소요된 것과 극명한 대조를 이룹니다. Deep1B 데이터셋에서도 낮은 품질의 그래프를 약 6시간에, 높은 품질의 그래프를 약 반나절에 구축할 수 있었으며, 이는 기존 방법으로는 수백 일이 필요할 것으로 추정되는 규모입니다. 이러한 성능 향상은 단순히 속도 개선을 넘어, 기존에 비실용적이었던 대규모 유사도 검색 응용을 실제로 가능하게 만들었습니다. + +이 연구의 이론적 의의는 GPU 메모리 계층 구조의 극단적인 대역폭 차이(레지스터 파일이 글로벌 메모리의 250배)를 알고리즘 설계에 반영하여, 점유율을 다소 희생하더라도 핵심 데이터를 가장 빠른 메모리에 유지하는 전략의 효과를 입증했다는 점입니다. 또한 홀수 크기 병합 네트워크와 같은 새로운 정렬 프리미티브를 도입하여 임의 크기의 배열을 효율적으로 처리할 수 있음을 보였습니다. 실무적으로는 이 연구의 알고리즘들이 FAISS(Facebook AI Similarity Search) 라이브러리로 오픈소스화되어, 이후 대규모 벡터 검색 분야에서 사실상의 표준 도구로 자리잡게 되었습니다. 이를 통해 기존에 GPU를 보유하고 있지만 유사도 검색에는 활용하지 못했던 연구자와 엔지니어들이 효율적인 유사도 검색을 즉시 수행할 수 있게 되었으며, 정보 검색, 추천 시스템, 이미지 검색, 자연어 처리 등 광범위한 분야에서 실질적인 영향을 미치고 있습니다. +- - - +# GPU를 활용한 십억 규모 유사도 검색 + +## 초록 + +이 논문은 Facebook AI Research에서 발표한 연구로, GPU를 활용하여 십억(billion) 규모의 벡터 데이터셋에서 유사도 검색(similarity search)을 효율적으로 수행하는 방법을 제안합니다. 유사도 검색이란 이미지나 비디오와 같은 복잡한 데이터를 고차원 벡터로 표현한 뒤, 주어진 쿼리 벡터와 가장 가까운 벡터들을 데이터베이스에서 찾아내는 작업입니다. 이 논문의 핵심 기여는 GPU의 빠른 레지스터 메모리에서 동작하는 $k$-선택(k-selection) 알고리즘을 설계하여 이론적 최대 성능의 55%까지 달성한 것이며, 이를 통해 기존 GPU 기반 최신 기법 대비 8.5배 빠른 최근접 이웃 검색을 구현하였습니다. 실제 응용에서는 Yfcc100M 데이터셋의 9,500만 장 이미지에 대한 고정확도 $k$-NN 그래프를 35분 만에 구축하였고, 10억 개 벡터에 대한 그래프를 4개의 Maxwell Titan X GPU에서 12시간 이내에 완성하였습니다. 이 구현체는 FAISS(Facebook AI Similarity Search)라는 이름으로 오픈소스로 공개되었습니다. + +## 서론 + +### 고차원 벡터 표현과 유사도 검색의 필요성 + +현대 AI 시스템에서 이미지, 비디오, 텍스트와 같은 복잡한 데이터는 딥러닝 모델을 통해 고차원 실수 벡터로 변환됩니다. 예를 들어, 텍스트 데이터의 경우 word2vec이라는 기법이 단어를 수십에서 수백 차원의 벡터로 표현하며, 이미지의 경우 합성곱 신경망(CNN)이 50차원에서 1,000차원 이상의 벡터 표현을 생성합니다. 이러한 벡터 표현(임베딩)은 원본 데이터의 의미적 특성을 수치적으로 포착하므로, 두 벡터 간의 거리가 가까울수록 원본 데이터가 유사하다고 해석할 수 있습니다. + +이러한 벡터 표현을 생성하는 과정 자체가 높은 연산 복잡도와 데이터 대역폭을 요구하기 때문에 GPU 시스템에서 효과적으로 수행됩니다. 그런데 벡터가 생성된 이후에도 이를 활용한 유사도 검색 역시 연산 집약적인 작업입니다. 가장 유사한 콘텐츠를 찾거나, 대규모 벡터 컬렉션에 대해 선형 분류기의 최대 응답을 구하는 등의 작업이 이에 해당합니다. 그러나 GPU 자원을 어떻게 효과적으로 활용할 것인지는 자명하지 않은 문제입니다. 이종(heterogeneous) 아키텍처를 활용하는 방법은 데이터베이스 커뮤니티에서도 핵심 연구 주제로 다루어지고 있습니다. + +### $k$-NN 그래프: 핵심 응용 + +대규모 벡터 컬렉션에서 수행할 수 있는 가장 비용이 큰 연산 중 하나는 $k$-NN 그래프의 구축입니다. $k$-NN 그래프란 데이터베이스의 각 벡터가 노드가 되고, 각 노드에서 자신의 $k$개 최근접 이웃으로 향하는 간선이 연결된 방향 그래프입니다. 직관적으로 설명하면, 수백만 장의 사진이 있을 때 각 사진마다 가장 비슷한 $k$장의 사진을 연결해 놓은 네트워크라고 생각할 수 있습니다. 이 논문에서는 이 $k$-NN 그래프 구축을 대표적인 응용 사례로 제시합니다. + +기존의 최신 기법인 NN-Descent는 데이터셋 자체 위에 큰 메모리 오버헤드를 발생시키며, 이 논문에서 다루는 십억 규모의 데이터베이스로 쉽게 확장되지 못한다는 한계가 있습니다. 또한 고차원성의 저주(curse of dimensionality)로 인해 십억 규모 데이터베이스에서는 완전 탐색이나 정확한 인덱싱 모두 비실용적이므로, 근사 검색 및 그래프 구축에 관한 방대한 연구가 존재합니다. + +### 벡터 압축과 Product Quantization + +GPU와 같이 메모리가 제한된 장치에서 거대한 데이터셋을 다루기 위해서는 벡터의 내부 압축 표현을 사용하는 것이 필수적입니다. [Product quantization for nearest neighbor search](https://arxiv.org/abs/1003.4083)에서 제안된 Product Quantization(PQ)은 고차원 벡터 공간을 여러 저차원 부분 공간의 카르테시안 곱으로 분해하고, 각 부분 공간을 독립적으로 양자화하는 기법입니다. 구체적으로, $d$차원 벡터 $\mathbf{x}$를 $m$개의 부분 벡터 $[\mathbf{x}^1, \mathbf{x}^2, \ldots, \mathbf{x}^m]$으로 나눈 뒤, 각 부분 벡터를 해당 부분 공간의 $k$-means 중심점 인덱스로 인코딩합니다. 이를 통해 최소한의 정확도 손실로 수십 배에서 수백 배의 압축을 달성할 수 있으며, 원본 벡터를 복원하지 않고도 압축 도메인에서 직접 거리를 계산할 수 있다는 장점이 있습니다. + +벡터 압축 방법은 크게 이진 코드(binary codes)와 양자화 방법(quantization methods)으로 분류됩니다. 이 논문에서는 이진 코드보다 더 효과적인 것으로 입증된 PQ 코드 기반 방법에 초점을 맞추고 있습니다. 이진 코드는 비완전 탐색 방법에서 상당한 오버헤드를 발생시키는 반면, PQ는 비대칭 거리 계산(Asymmetric Distance Computation)을 통해 쿼리는 압축하지 않은 채로 데이터베이스 벡터만 압축하여 효율적인 검색이 가능합니다. + +### 기존 방법의 한계와 본 연구의 위치 + +원래의 PQ 기반 인덱싱 방법인 IVFADC 이후 여러 개선이 제안되었지만, 대부분 GPU에서 효율적으로 구현하기 어렵습니다. [The inverted multi-index](https://arxiv.org/abs/1201.3655)는 고속/저품질 운영점에서 유용하지만 복잡한 "multi-sequence" 알고리즘에 의존합니다. [Optimized product quantization](https://arxiv.org/abs/1307.1708)(OPQ)은 입력 벡터에 선형 변환을 적용하여 PQ의 정확도를 향상시키는 방법으로, 전처리 단계로 적용할 수 있습니다. 그러나 SIMD 최적화된 IVFADC 구현은 적은 수의 조대 양자화 중심점이라는 차선의 파라미터에서만 동작하며, LOPQ나 [Polysemous codes](https://arxiv.org/abs/1609.01882)와 같은 방법들은 GPU에서 효율적으로 구현하기에 너무 복잡합니다. Polysemous codes의 경우, 이진 코드와 PQ의 이중 해석을 가능하게 하는 혁신적인 접근이지만, 시뮬레이티드 어닐링 기반의 인덱스 할당 최적화가 GPU 병렬화에 적합하지 않습니다. + +기존 GPU 기반 유사도 검색 구현은 대부분 이진 코드, 소규모 데이터셋, 또는 완전 탐색에 한정되어 있었습니다. 양자화 코드를 사용하여 십억 규모 데이터셋에 적합한 것으로 보이는 유일한 기존 연구는 [Wieschollek et al.](https://arxiv.org/abs/1702.05911)의 연구이며, 이것이 본 논문이 비교하는 GPU 기반 기존 최신 기법입니다. + +### 본 논문의 기여 + +이 논문은 세 가지 핵심 기여를 제시합니다. 첫째, 빠른 레지스터 메모리에서 동작하며 다른 커널과 융합(fuse)할 수 있을 만큼 유연한 GPU $k$-선택 알고리즘과 그 복잡도 분석을 제공합니다. 둘째, GPU에서의 정확 및 근사 $k$-최근접 이웃 검색을 위한 거의 최적의 알고리즘 레이아웃을 설계합니다. 셋째, 단일 및 다중 GPU 구성에서 중간 규모부터 대규모까지의 최근접 이웃 검색 작업에서 기존 기법을 큰 차이로 능가함을 보여주는 광범위한 실험을 수행합니다. 이후 섹션에서는 문제 정의와 표기법, GPU 아키텍처와 유사도 검색에서의 문제점, 핵심 기여인 $k$-선택 방법, 알고리즘 계산 레이아웃, 그리고 광범위한 실험 결과가 순차적으로 다루어집니다. +## 문제 정의 + +### 최근접 이웃 검색의 형식적 정의 + +이 논문에서 다루는 핵심 문제는 벡터 컬렉션에서의 유사도 검색(similarity search)입니다. $d$차원 쿼리 벡터 $x \in \mathbb{R}^d$와 데이터베이스 벡터 컬렉션 $[y_i]_{i=0:\ell}$ $(y_i \in \mathbb{R}^d)$가 주어졌을 때, 다음과 같이 $k$개의 최근접 이웃을 찾는 것이 목표입니다. + +$$L = k\text{-}\textrm{argmin}_{i=0:\ell} \|x - y_i\|_2$$ + +여기서 $0:\ell$이라는 표기는 $\{0, \ldots, \ell-1\}$을 의미하는 0-기반 인덱싱 표기법입니다. 이 수식의 의미를 직관적으로 설명하면, 전체 $\ell$개의 데이터베이스 벡터 중에서 쿼리 $x$와의 L2 거리가 가장 작은 $k$개의 벡터 인덱스를 찾는 것입니다. L2 거리가 가장 널리 사용되는 이유는 다양한 임베딩 학습 과정에서 L2 거리가 최적화 대상으로 설계되며, 선형대수적으로 다루기 편리한 성질을 갖기 때문입니다. + +이 검색 과정의 핵심 연산은 $k$-선택($k$-selection)입니다. 배열 $[a_i]_{i=0:\ell}$에 대해 $k$-선택은 가장 작은 $k$개의 값 $[a_{s_i}]_{i=0:k}$를 $a_{s_i} \leq a_{s_{i+1}}$ 순서로 정렬하여 찾아내고, 해당 원소들의 원래 인덱스 $[s_i]_{i=0:k}$ $(0 \leq s_i < \ell)$도 함께 반환합니다. 값 $a_i$는 32비트 부동소수점으로, 인덱스 $s_i$는 32비트 또는 64비트 정수로 표현됩니다. 코사인 유사도처럼 최대값을 찾아야 하는 경우에는 비교 연산자를 변경하여 대응할 수 있으며, 동일한 키 값을 가진 원소들 사이의 순서는 정의되지 않습니다. + +### 배치 처리와 정확 검색 + +실제 응용에서는 단일 쿼리가 아닌 $n_{\text{q}}$개의 쿼리 벡터 $[x_j]_{j=0:n_{\text{q}}}$ $(x_j \in \mathbb{R}^d)$를 동시에 처리하는 배치(batch) 검색이 일반적입니다. 배치 처리는 다중 CPU 스레드나 GPU에서 병렬 실행의 유연성을 높여줍니다. 배치 $k$-선택의 경우, $n_{\text{q}}$개의 독립적인 배열에서 각각 $k$개의 원소와 인덱스를 선택해야 하며, 각 배열의 길이 $\ell_i \geq k$는 서로 다를 수 있습니다. + +정확 검색(exact search)은 전체 쌍별 거리 행렬 $D = [\|x_j - y_i\|_2^2]_{j=0:n_{\text{q}}, i=0:\ell} \in \mathbb{R}^{n_{\text{q}} \times \ell}$을 완전히 계산합니다. 이때 실제 구현에서는 다음과 같은 분해를 활용합니다. + +$$\|x_j - y_i\|_2^2 = \|x_j\|^2 + \|y_i\|^2 - 2\langle x_j, y_i \rangle$$ + +이 분해가 왜 중요한지 직관적으로 이해해 보겠습니다. 좌변의 L2 거리를 직접 계산하려면 매 쿼리-데이터 쌍마다 $d$번의 뺄셈과 제곱 연산이 필요합니다. 그러나 우변으로 분해하면, $\|x_j\|^2$는 쿼리 행렬 $X$에 대해 한 번만, $\|y_i\|^2$는 데이터베이스 행렬 $Y$에 대해 한 번만 사전 계산하면 됩니다. 가장 연산 비용이 큰 부분은 내적 $\langle x_j, y_i \rangle$의 계산인데, 이는 행렬 곱셈 $XY^{\top}$으로 표현되어 고도로 최적화된 BLAS 라이브러리를 활용할 수 있습니다. 최종적으로 각 쿼리에 대해 거리 행렬 $D$의 해당 행에서 $k$-선택을 수행하면 $k$-최근접 이웃을 얻게 됩니다. + +다음 Python 코드는 이 분해를 활용한 정확 검색의 개념을 보여줍니다. + +```python +import numpy as np + +def exact_knn_search(X, Y, k): + """ + X: 쿼리 행렬 (n_q x d) + Y: 데이터베이스 행렬 (ell x d) + k: 반환할 최근접 이웃 수 + """ + # 사전 계산: 각 벡터의 L2 노름 제곱 + x_norms = np.sum(X ** 2, axis=1, keepdims=True) # (n_q, 1) + y_norms = np.sum(Y ** 2, axis=1, keepdims=True).T # (1, ell) + + # 핵심 연산: 행렬 곱셈으로 내적 계산 + inner_products = X @ Y.T # (n_q, ell) - BLAS 최적화 활용 + + # 거리 행렬 조립: ||x_j||^2 + ||y_i||^2 - 2 + D = x_norms + y_norms - 2 * inner_products + + # 각 쿼리에 대해 k-선택 수행 + indices = np.argpartition(D, k, axis=1)[:, :k] # k개의 최소 거리 인덱스 + return indices +``` + +이 코드에서 핵심은 `X @ Y.T` 연산이 전체 계산의 병목이라는 점입니다. 이 행렬 곱셈은 GPU의 cuBLAS와 같은 고성능 라이브러리에서 이론적 최대 성능에 근접하게 실행될 수 있으므로, 거리 분해를 통해 유사도 검색의 핵심 연산을 행렬 곱셈으로 환원하는 것이 매우 효과적입니다. + +### 압축 도메인 검색: IVFADC 인덱싱 + +정확 검색은 데이터베이스 크기가 커질수록 비용이 급격히 증가하므로, 이 논문에서는 근사 최근접 이웃 검색(approximate nearest-neighbor search)에 초점을 맞춥니다. 특히 IVFADC(Inverted File with Asymmetric Distance Computation) 인덱싱 구조를 핵심적으로 다룹니다. + +IVFADC는 두 단계의 양자화(quantization)를 통해 데이터베이스 벡터를 압축합니다. 데이터베이스 벡터 $y$는 다음과 같이 근사됩니다. + +$$y \approx q(y) = q_1(y) + q_2(y - q_1(y))$$ + +여기서 $q_1: \mathbb{R}^d \rightarrow \mathcal{C}_1 \subset \mathbb{R}^d$는 조대 양자화기(coarse quantizer), $q_2: \mathbb{R}^d \rightarrow \mathcal{C}_2 \subset \mathbb{R}^d$는 세밀 양자화기(fine quantizer)입니다. 양자화기란 입력 벡터를 유한 집합의 원소 중 하나로 매핑하는 함수입니다. 이를 도서관에 비유하면, $q_1$은 책을 대략적인 주제별 서가에 배치하는 것이고, $q_2$는 해당 서가 내에서 책의 정확한 위치를 세밀하게 기록하는 것에 해당합니다. $q_2$가 인코딩하는 대상은 원본 벡터가 아니라 잔차 벡터(residual vector) $y - q_1(y)$, 즉 조대 양자화 후 남은 오차입니다. + +유한 집합의 원소로 매핑되므로, $q(y)$는 $q_1(y)$의 인덱스와 $q_2(y - q_1(y))$의 인덱스만으로 인코딩할 수 있어 메모리를 크게 절약합니다. 비대칭 거리 계산(ADC) 방법은 이 압축된 표현을 사용하여 근사 결과를 반환합니다. + +$$L_{\text{ADC}} = k\text{-}\textrm{argmin}_{i=0:\ell} \|x - q(y_i)\|_2$$ + +"비대칭"이라는 이름은 쿼리 $x$는 원본 그대로 사용하고 데이터베이스 벡터 $y_i$만 양자화된 $q(y_i)$로 대체하기 때문에 붙여진 것입니다. + +IVFADC에서는 전체 데이터베이스를 완전 탐색하지 않고, 조대 양자화기 $q_1$을 기반으로 후보 벡터를 사전 선별합니다. 먼저 쿼리와 가장 가까운 $\tau$개의 조대 중심점을 찾습니다. + +$$L_{\text{IVF}} = \tau\text{-}\textrm{argmin}_{c \in \mathcal{C}_1} \|x - c\|_2$$ + +멀티프로브 파라미터 $\tau$는 탐색할 조대 수준 중심점의 수를 결정합니다. 그런 다음 선별된 중심점에 속하는 벡터들에 대해서만 거리를 계산합니다. + +$$L_{\text{IVFADC}} = \underset{i=0:\ell \text{ s.t. } q_1(y_i) \in L_{\text{IVF}}}{k\text{-}\textrm{argmin}} \|x - q(y_i)\|_2$$ + +이 구조에서 데이터베이스 벡터 $y_i$들은 $|\mathcal{C}_1|$개의 역색인 리스트(inverted list) $\mathcal{I}_1, \ldots, \mathcal{I}_{|\mathcal{C}_1|}$로 그룹화됩니다. 같은 리스트에 속한 벡터들은 동일한 $q_1(y_i)$ 값을 공유합니다. 따라서 가장 메모리 집약적인 연산은 $L_{\text{IVFADC}}$를 계산하는 것이며, 이는 $\tau$개의 역색인 리스트를 선형 스캔하는 것으로 귀결됩니다. +## 문제 정의 + +### 양자화기 설계 원리와 Product Quantizer + +앞서 소개한 IVFADC 인덱싱 구조에서 두 양자화기 $q_1$과 $q_2$는 근본적으로 다른 설계 요구사항을 갖습니다. 이 차이를 이해하는 것이 IVFADC의 성능 특성을 파악하는 데 핵심적입니다. + +조대 양자화기 $q_1$은 역색인 리스트의 개수를 직접 결정하므로, 재생산 값(reproduction value)의 수가 지나치게 많아서는 안 됩니다. 역색인 리스트가 너무 많아지면 각 리스트에 포함되는 벡터 수가 극도로 적어져 검색 효율이 떨어지고, 리스트 관리 오버헤드가 증가합니다. 일반적으로 $|\mathcal{C}_1| \approx \sqrt{\ell}$로 설정하며, $k$-means 알고리즘으로 학습됩니다. 예를 들어 데이터베이스에 $10^9$개의 벡터가 있다면 약 $\sqrt{10^9} \approx 31{,}623$개의 조대 중심점을 사용하는 것이 적절합니다. 이렇게 하면 각 역색인 리스트에 평균적으로 약 $31{,}623$개의 벡터가 포함되어, 멀티프로브 파라미터 $\tau$를 조절함으로써 검색 정확도와 속도 사이의 균형을 유연하게 조정할 수 있습니다. + +반면 세밀 양자화기 $q_2$는 더 많은 메모리를 사용하여 보다 정밀한 표현을 제공할 수 있습니다. 역색인 리스트에는 각 벡터의 ID(4바이트 또는 8바이트 정수)도 함께 저장되므로, 양자화 코드가 이보다 짧으면 의미가 없습니다. 즉, $\log_2 |\mathcal{C}_2| > 4 \times 8 = 32$비트 이상의 코드 길이가 필요합니다. 이 조건은 세밀 양자화기가 최소한 벡터 ID만큼의 정보량을 담아야 저장 효율이 정당화된다는 실용적 제약을 반영합니다. + +세밀 양자화기 $q_2$로는 Product Quantizer(PQ)가 사용됩니다. PQ의 핵심 아이디어는 처리 비용을 증가시키지 않으면서도 매우 많은 수의 재생산 값을 제공하는 것입니다. PQ는 $d$차원 벡터 $y$를 $b$개의 부분 벡터(sub-vector)로 분할합니다. 여기서 $b$는 차원 $d$의 짝수 약수여야 합니다. + +$$y = [y^0 \ldots y^{b-1}]$$ + +각 부분 벡터 $y^j$는 독립적인 부분 양자화기 $q^j$에 의해 양자화되어 튜플 $(q^0(y^0), \ldots, q^{b-1}(y^{b-1}))$을 생성합니다. 각 부분 양자화기는 일반적으로 256개의 재생산 값을 가지므로, 하나의 부분 양자화 결과가 정확히 1바이트에 저장됩니다. PQ의 전체 양자화 값은 다음과 같이 표현됩니다. + +$$q_2(y) = q^0(y^0) + 256 \times q^1(y^1) + \ldots + 256^{b-1} \times q^{b-1}(y^{b-1})$$ + +저장 관점에서 이는 각 부분 양자화기가 생성한 바이트들의 단순 연결(concatenation)에 불과합니다. 따라서 PQ는 $b$바이트 코드를 생성하며, 총 $|\mathcal{C}_2| = 256^b$개의 재생산 값을 갖습니다. 이것이 PQ의 가장 강력한 장점입니다. 예를 들어 $b = 8$이면 단 8바이트의 코드로 $256^8 \approx 1.8 \times 10^{19}$개의 서로 다른 재생산 값을 표현할 수 있습니다. 이는 카르테시안 곱(Cartesian product)의 조합적 폭발을 활용한 것으로, 각 부분 공간에서 256개의 중심점만 학습하면 되므로 $k$-means 사전(dictionary)의 크기가 작고 양자화 자체의 계산 비용도 매우 저렴합니다. + +이 구조를 구체적인 수치 예시로 살펴보겠습니다. 128차원 벡터를 $b = 16$개의 부분 벡터로 분할하면, 각 부분 벡터는 $128 / 16 = 8$차원이 됩니다. 각 8차원 부분 공간에서 256개의 $k$-means 중심점을 학습하므로, 총 사전 크기는 $16 \times 256 \times 8 = 32{,}768$개의 부동소수점 값에 불과합니다. 그러나 이 작은 사전으로 $256^{16} \approx 3.4 \times 10^{38}$개의 고유한 벡터 표현이 가능합니다. 원본 128차원 벡터가 512바이트(32비트 $\times$ 128)를 차지하는 반면, PQ 코드는 단 16바이트로 32배의 압축을 달성합니다. + +다음 코드는 PQ 인코딩과 비대칭 거리 계산의 핵심 원리를 보여줍니다. + +```python +import numpy as np + +def pq_encode(y, codebooks, b): + """ + Product Quantizer 인코딩 + y: 원본 벡터 (d,) + codebooks: b개의 부분 양자화기 사전, 각각 (256, d/b) 크기 + b: 부분 벡터 수 + """ + d = len(y) + dsub = d // b # 각 부분 벡터의 차원 + codes = np.zeros(b, dtype=np.uint8) + + for j in range(b): + # j번째 부분 벡터 추출 + sub_vector = y[j * dsub : (j + 1) * dsub] + # 가장 가까운 중심점의 인덱스를 1바이트로 저장 + distances = np.sum((codebooks[j] - sub_vector) ** 2, axis=1) + codes[j] = np.argmin(distances) # 0~255 범위의 인덱스 + + return codes # b바이트 코드 반환 + +def adc_distance(x, codes, codebooks, b): + """ + 비대칭 거리 계산 (ADC) + x: 쿼리 벡터 (원본, 압축하지 않음) + codes: PQ 인코딩된 데이터베이스 벡터 (b,) + """ + d = len(x) + dsub = d // b + dist = 0.0 + + for j in range(b): + # 쿼리의 j번째 부분 벡터와 코드북 중심점 간 거리 + sub_query = x[j * dsub : (j + 1) * dsub] + centroid = codebooks[j][codes[j]] # codes[j]번째 중심점 복원 + dist += np.sum((sub_query - centroid) ** 2) + + return dist # ||x - q_2(y)||^2의 근사값 +``` + +이 코드에서 `pq_encode` 함수는 원본 벡터를 $b$바이트 코드로 압축하고, `adc_distance` 함수는 쿼리 벡터 $x$는 원본 그대로 사용하면서 데이터베이스 벡터는 압축된 코드에서 복원하여 거리를 계산합니다. 실제 GPU 구현에서는 이 거리 계산이 룩업 테이블(lookup table)을 통해 더욱 효율적으로 수행되며, 이에 대한 상세한 계산 레이아웃은 이후 섹션에서 다루어집니다. + +요약하면, 이 논문의 문제 정의는 L2 거리 기반 $k$-최근접 이웃 검색을 정확 검색과 근사 검색 두 가지 관점에서 형식화하고 있습니다. 정확 검색은 앞서 소개한 거리 행렬 분해를 통해 행렬 곱셈으로 환원되며, 근사 검색은 IVFADC 구조를 통해 조대 양자화기의 역색인 기반 후보 선별과 Product Quantizer의 압축 도메인 거리 계산을 결합합니다. 이 두 가지 검색 방식 모두에서 $k$-선택이 핵심 연산으로 등장하며, GPU에서 이를 효율적으로 수행하는 것이 이 논문의 주요 기술적 도전 과제입니다. +## GPU 아키텍처 개요와 $k$-선택 + +이 섹션에서는 Nvidia GPU의 범용 컴퓨팅 아키텍처와 프로그래밍 모델의 핵심 세부사항을 검토하고, 유사도 검색에서 GPU 친화적이지 않은 연산인 $k$-선택의 도전 과제를 심층적으로 분석합니다. GPU 아키텍처에 대한 정확한 이해는 이후 제안되는 WarpSelect 알고리즘의 설계 동기와 최적화 전략을 파악하는 데 필수적입니다. + +### 아키텍처 + +#### GPU 레인과 워프 + +Nvidia GPU는 범용 컴퓨터로서, 32개의 CUDA 스레드로 구성된 벡터 단위인 **워프(warp)**를 통해 명령어 스트림을 실행합니다. 워프 내의 개별 스레드는 **레인(lane)**이라 불리며, 0부터 31까지의 레인 ID를 갖습니다. 여기서 "스레드"라는 용어가 사용되지만, 현대 벡터화된 멀티코어 CPU와의 가장 적절한 비유는 각 워프가 하나의 독립적인 CPU 하드웨어 스레드에 해당한다는 것입니다. 이는 워프가 하나의 명령어 카운터(instruction counter)를 공유하기 때문입니다. 직관적으로 설명하면, 워프는 32명의 군인이 한 줄로 서서 동일한 구령에 맞춰 동시에 같은 동작을 수행하는 것과 유사합니다. 만약 일부 레인이 서로 다른 실행 경로를 취하게 되면 **워프 발산(warp divergence)**이 발생하여 성능이 저하됩니다. 이는 마치 일부 군인만 다른 동작을 해야 할 때 나머지가 기다려야 하는 상황과 같습니다. + +[NVIDIA Tesla 아키텍처](https://ieeexplore.ieee.org/document/4523358)에서 확립된 이 실행 모델에 따르면, 각 레인은 공유 레지스터 파일 내에서 최대 255개의 32비트 레지스터를 사용할 수 있습니다. CPU 관점에서의 비유를 확장하면, 이는 폭 32의 SIMD 벡터 레지스터가 최대 255개 존재하는 것과 동일하며, 워프 레인이 곧 SIMD 벡터 레인의 역할을 수행합니다. + +#### 워프의 집합: 블록과 스트리밍 멀티프로세서 + +사용자가 설정 가능한 1개에서 32개의 워프 모음이 하나의 **블록(block)** 또는 **협력 스레드 배열(CTA, Co-operative Thread Array)**을 구성합니다. 각 블록은 최대 48 KiB 크기의 고속 **공유 메모리(shared memory)**를 보유합니다. 블록 내의 개별 CUDA 스레드는 블록 상대적인 **스레드 ID(thread id)**를 가지며, 이를 통해 작업을 분할하고 할당할 수 있습니다. + +각 블록은 GPU의 단일 코어인 **스트리밍 멀티프로세서(SM, Streaming Multiprocessor)**에서 실행됩니다. 각 SM은 ALU, 메모리 로드/스토어 유닛, 다양한 특수 명령어 유닛을 포함하는 **기능 유닛(functional units)**을 갖추고 있습니다. GPU는 모든 SM에 걸쳐 다수의 워프에서 많은 연산을 동시에 비행 중(in-flight) 상태로 유지함으로써 실행 지연 시간을 숨깁니다. 개별 워프 레인의 명령어 처리량은 낮고 지연 시간은 높지만, 모든 SM의 총합 산술 처리량은 일반적인 CPU보다 5~10배 높습니다. + +#### 그리드와 커널 + +블록들은 **커널(kernel)** 내에서 **블록의 그리드(grid of blocks)**로 조직됩니다. 각 블록에는 그리드 상대적인 ID가 할당됩니다. 커널은 호스트 CPU가 GPU에 실행을 예약하는 작업 단위(인수를 가진 명령어 스트림)입니다. 하나의 블록이 완료되면 새로운 블록이 스케줄링될 수 있으며, 서로 다른 커널의 블록들이 동시에 실행될 수도 있습니다. 커널 간의 순서는 **스트림(streams)**과 **이벤트(events)** 같은 순서 지정 프리미티브를 통해 제어됩니다. + +#### 자원과 점유율 + +동시에 실행되는 블록의 수는 각 블록이 사용하는 공유 메모리와 레지스터 자원에 의존합니다. CUDA 스레드당 레지스터 사용량은 컴파일 시점에 결정되고, 공유 메모리 사용량은 런타임에 선택할 수 있습니다. 이 사용량이 GPU의 **점유율(occupancy)**에 영향을 미칩니다. 예를 들어, 하나의 블록이 48 KiB 공유 메모리 전체를 독점하거나 스레드당 32개가 아닌 128개의 레지스터를 요구하면, 동일 SM에서 1~2개의 다른 블록만 동시 실행 가능하여 낮은 점유율이 됩니다. 높은 점유율에서는 더 많은 블록이 모든 SM에 분포하여 더 많은 작업이 동시에 비행 중 상태가 됩니다. + +[GPU 벤치마킹 및 튜닝 연구](https://dl.acm.org/doi/10.1145/1513895.1513901)에서 체계적으로 분석된 바와 같이, 이러한 점유율과 성능 사이의 트레이드오프는 GPU 알고리즘 설계에서 핵심적인 고려사항입니다. 점유율이 높다고 항상 성능이 좋은 것은 아니며, 때로는 더 많은 레지스터를 사용하여 점유율을 낮추더라도 빠른 메모리에 더 큰 작업 집합을 유지하는 것이 전체 성능에 유리할 수 있습니다. + +#### 메모리 유형과 대역폭 계층 + +서로 다른 블록과 커널은 **글로벌 메모리(global memory)**를 통해 통신하며, 이는 일반적으로 4~32 GB 크기에 CPU 메인 메모리보다 5~10배 높은 대역폭을 제공합니다. 공유 메모리는 속도 면에서 CPU L1 캐시에 비유할 수 있습니다. GPU 레지스터 파일 메모리는 가장 높은 대역폭을 가진 메모리입니다. + +GPU에서 많은 수의 명령어를 동시에 비행 중 상태로 유지하기 위해서는 방대한 레지스터 파일이 필요합니다. 최신 Pascal P100에서는 레지스터 파일이 **14 MB**에 달하는데, 이는 CPU의 수십 KB와 극명한 대조를 이룹니다. 레지스터 대 공유 대 글로벌 메모리의 총합 단면 대역폭 비율은 전형적으로 **250 : 6.25 : 1**이며, 레지스터 파일의 경우 10~100s TB/s에 달합니다. 이 비율의 의미를 구체적으로 이해하면, 글로벌 메모리 대역폭이 약 700 GB/s라면 공유 메모리는 약 4.4 TB/s, 레지스터 파일은 약 175 TB/s에 해당합니다. 이러한 극단적인 대역폭 차이가 바로 이 논문에서 레지스터 기반 $k$-선택 알고리즘을 설계하는 핵심 동기입니다. + +다음 코드는 GPU 메모리 계층의 대역폭 비율을 직관적으로 보여줍니다. + +```python +# GPU 메모리 계층 대역폭 비율 시각화 +class GPUMemoryHierarchy: + """Pascal P100 기준 GPU 메모리 계층 특성""" + + def __init__(self): + # 메모리 유형별 특성 + self.global_memory = { + 'size': '16 GB', + 'bandwidth_ratio': 1, # 기준점 + 'bandwidth_approx': '732 GB/s', + 'analogy': 'CPU 메인 메모리 (단, 5-10x 더 빠름)' + } + self.shared_memory = { + 'size': '48 KiB per block', + 'bandwidth_ratio': 6.25, # 글로벌 대비 6.25배 + 'bandwidth_approx': '~4.6 TB/s', + 'analogy': 'CPU L1 캐시' + } + self.register_file = { + 'size': '14 MB (전체 GPU)', + 'bandwidth_ratio': 250, # 글로벌 대비 250배 + 'bandwidth_approx': '~183 TB/s', + 'analogy': '가장 빠른 메모리 - CPU에는 대응물 없음' + } + + def occupancy_example(self): + """점유율 트레이드오프 예시""" + # 시나리오 1: 높은 점유율, 적은 레지스터 + scenario_high_occ = { + 'registers_per_thread': 32, + 'shared_mem_per_block': '16 KiB', + 'concurrent_blocks_per_SM': 4, # 높은 점유율 + 'working_set_speed': 'shared memory (6.25x)' + } + # 시나리오 2: 낮은 점유율, 많은 레지스터 + scenario_low_occ = { + 'registers_per_thread': 128, + 'shared_mem_per_block': '48 KiB', + 'concurrent_blocks_per_SM': 1, # 낮은 점유율 + 'working_set_speed': 'register file (250x)' # 훨씬 빠름! + } + return scenario_high_occ, scenario_low_occ +``` + +이 코드에서 볼 수 있듯이, 레지스터 파일의 대역폭은 글로벌 메모리의 250배에 달하므로, 점유율을 다소 희생하더라도 핵심 데이터를 레지스터에 유지하는 전략이 매우 효과적일 수 있습니다. + +### GPU 레지스터 파일 활용 + +#### 구조화된 레지스터 데이터 + +공유 메모리와 레지스터 메모리의 사용은 효율성 트레이드오프를 수반합니다. 이들은 점유율을 낮추지만, 더 빠른 메모리에 더 큰 작업 집합을 유지함으로써 전체 성능을 향상시킬 수 있습니다. 점유율을 희생하거나 공유 메모리 대신 레지스터 상주 데이터를 적극 활용하는 것이 종종 이득이 됩니다. + +GPU 레지스터 파일이 매우 크기 때문에, 단순한 임시 피연산자가 아닌 **구조화된 데이터(structured data)**를 저장하는 것이 유용합니다. 단일 레인이 자신의 스칼라 레지스터를 사용하여 로컬 작업을 해결할 수 있지만, 병렬성과 저장 공간이 제한됩니다. 대신, GPU 워프 내의 레인들이 **워프 셔플 명령어(warp shuffle instruction)**를 사용하여 레지스터 데이터를 교환할 수 있으며, 이를 통해 워프 전체 수준의 병렬성과 저장 공간을 확보할 수 있습니다. + +#### 레인-스트라이드 레지스터 배열 + +이를 달성하기 위한 일반적인 패턴이 **레인-스트라이드 레지스터 배열(lane-stride register array)**입니다. 원소 $[a_i]_{i=0:\ell}$이 주어졌을 때, 각 연속적인 값이 이웃 레인의 레지스터에 저장됩니다. 이 배열은 레인당 $\ell/32$개의 레지스터에 저장되며, $\ell$은 32의 배수입니다. + +구체적으로, 레인 $j$는 원소 집합 $\{a_j, a_{32+j}, \ldots, a_{\ell-32+j}\}$를 저장합니다. 반면 레지스터 $r$은 원소 집합 $\{a_{32r}, a_{32r+1}, \ldots, a_{32r+31}\}$을 보유합니다. 이 배치를 직관적으로 이해하기 위해 구체적인 예시를 살펴보겠습니다. $\ell = 64$인 배열 $[a_0, a_1, \ldots, a_{63}]$이 있다고 가정하면, 레인당 $64/32 = 2$개의 레지스터가 필요합니다. + +```python +import numpy as np + +def lane_stride_register_layout(elements, num_lanes=32): + """ + 레인-스트라이드 레지스터 배열의 데이터 배치를 시각화 + elements: 저장할 원소 배열 (길이는 32의 배수) + """ + ell = len(elements) + assert ell % num_lanes == 0, "길이는 32의 배수여야 합니다" + regs_per_lane = ell // num_lanes + + # 각 레인이 저장하는 원소 확인 + for lane_j in [0, 1, 31]: # 대표적인 레인만 표시 + stored = [elements[32 * r + lane_j] for r in range(regs_per_lane)] + print(f"레인 {lane_j:2d}: {stored}") + + print() + # 각 레지스터가 보유하는 원소 확인 + for reg_r in range(regs_per_lane): + held = [elements[32 * reg_r + lane] for lane in range(num_lanes)] + print(f"레지스터 {reg_r}: a[{32*reg_r}] ~ a[{32*reg_r+31}]") + +# 예시: 64개 원소의 레인-스트라이드 배치 +elements = list(range(64)) +lane_stride_register_layout(elements) +# 레인 0: [a_0, a_32] → 레지스터 0과 1에 각각 저장 +# 레인 1: [a_1, a_33] → 레지스터 0과 1에 각각 저장 +# 레인 31: [a_31, a_63] → 레지스터 0과 1에 각각 저장 +# 레지스터 0: a[0] ~ a[31] → 32개 레인에 분산 +# 레지스터 1: a[32] ~ a[63] → 32개 레인에 분산 +``` + +원소 $a_i$를 조작하기 위해서는 해당 원소가 저장된 레지스터 번호(즉, $\lfloor i/32 \rfloor$)와 배열 길이 $\ell$이 **어셈블리 시점**에 알려져야 합니다. 반면 해당 원소가 위치한 레인(즉, $i \bmod 32$)은 **런타임 지식**일 수 있습니다. 이 구분이 중요한 이유는 GPU에서 레지스터 접근은 컴파일 시점에 결정되어야 하지만, 워프 셔플을 통한 레인 간 데이터 교환은 런타임에 동적으로 수행될 수 있기 때문입니다. + +다양한 접근 패턴(시프트, 임의 대 임의)이 제공되며, 이 논문에서는 [병렬 알고리즘 및 아키텍처 입문](https://dl.acm.org/doi/book/10.5555/562413)에서 소개된 **버터플라이 순열(butterfly permutation)**을 광범위하게 활용합니다. 버터플라이 순열은 하이퍼큐브 네트워크에서 유래한 통신 패턴으로, 특정 비트 위치를 기준으로 레인 쌍 간에 데이터를 교환하는 효율적인 방법입니다. + +### CPU 대비 GPU에서의 $k$-선택 + +앞서 문제 정의에서 소개한 $k$-선택은 유사도 검색의 핵심 연산이지만, GPU에서 효율적으로 구현하기가 특히 어려운 부분입니다. 이 하위 섹션에서는 기존 접근법들의 한계를 체계적으로 분석합니다. + +#### 다중 패스 알고리즘의 한계 + +임의로 큰 $\ell$과 $k$에 대한 $k$-선택 알고리즘들은 GPU로 이식될 수 있으며, [GPU에서의 빠른 $k$-선택 알고리즘](https://dl.acm.org/doi/10.1145/2304576.2304623)에서 다루어진 기수 선택(radix selection), 버킷 선택(bucket selection), [확률적 선택(probabilistic selection)](https://dl.acm.org/doi/10.1145/2304576.2304623), 퀵셀렉트(quickselect), 절단 정렬(truncated sorts) 등이 포함됩니다. 그러나 이들의 성능은 글로벌 메모리에서의 **다중 패스(multiple passes)**에 의해 지배됩니다. + +유사도 검색에서는 입력 거리가 즉석에서(on-the-fly) 계산되거나 작은 블록에만 저장되고, 전체가 한꺼번에 존재하지 않는 경우가 많습니다. 전체 명시적 배열이 어떤 메모리에도 들어가지 못할 만큼 클 수 있으며, 처리 시작 시점에 그 크기를 알 수 없을 수도 있어 다중 패스를 요구하는 알고리즘은 비실용적입니다. + +개별 알고리즘의 추가적인 문제점도 존재합니다. 퀵셀렉트는 $\mathcal{O}(\ell)$ 크기의 저장 공간에서 파티셔닝을 수행해야 하는데, 이는 데이터 의존적 메모리 이동을 초래합니다. 이로 인해 과도한 메모리 트랜잭션이 발생하거나, 쓰기 오프셋을 결정하기 위한 병렬 접두사 합(parallel prefix sums)이 필요하여 동기화 오버헤드가 추가됩니다. 기수 선택은 파티셔닝이 없지만 여전히 다중 패스가 필요합니다. + +#### 힙 기반 접근법의 병렬성 문제 + +유사도 검색 응용에서는 일반적으로 $k < 1000$ 정도의 소수 결과만 필요합니다. 이 범위에서 CPU에서의 전형적인 선택은 최대 힙(max-heap)을 통한 선택이지만, 힙은 직렬 트리 갱신(serial tree update)으로 인해 데이터 병렬성을 거의 노출하지 못하며 SIMD 실행 유닛을 포화시킬 수 없습니다. [ad-heap](https://dl.acm.org/doi/10.1145/2588555.2610502)은 이종 시스템에서 가용한 병렬성을 더 잘 활용하지만, 여전히 직렬 작업과 병렬 작업을 적절한 실행 유닛 사이에 분배하려는 시도에 머무릅니다. + +힙의 직렬적 특성에도 불구하고, 작은 $k$에서 CPU는 모든 상태를 L1 캐시에 쉽게 유지할 수 있으며, L1 캐시 지연 시간과 대역폭이 제한 요인으로 남습니다. PQ 코드 조작과 같은 다른 유사도 검색 구성 요소가 CPU 성능에 더 큰 영향을 미치는 경향이 있습니다. + +#### GPU 힙과 병렬 우선순위 큐의 한계 + +힙은 GPU에서도 유사하게 구현될 수 있지만, 직접적인 GPU 힙 구현은 높은 워프 발산과 불규칙한 데이터 의존적 메모리 이동으로 인해 성능이 저하됩니다. 삽입되는 각 원소가 취하는 경로가 힙 내의 다른 값들에 의존하기 때문입니다. + +[GPU 병렬 우선순위 큐](https://dl.acm.org/doi/10.1145/2503210.2503245)는 다수의 동시 갱신을 허용하여 직렬 힙 갱신을 개선하지만, 각 삽입마다 잠재적으로 여러 번의 소규모 정렬과 데이터 의존적 메모리 이동이 필요합니다. 더욱이, 서로 다른 스트림에서의 커널 실행을 통한 다수의 동기화 장벽(synchronization barriers)을 사용하며, 연속적인 커널 실행과 CPU 호스트와의 조율에 따른 추가 지연 시간이 발생합니다. + +#### fgknn 라이브러리로부터의 영감 + +작은 $k$에 대해 더 새로운 GPU 알고리즘도 존재하는데, 특히 [fgknn 라이브러리의 선택 알고리즘](https://dl.acm.org/doi/10.1145/2820783.2820791)이 주목할 만합니다. 이 알고리즘은 복잡하며 너무 많은 동기화 지점, 더 큰 커널 실행 오버헤드, 느린 메모리의 사용, 계층 구조·파티셔닝·버퍼링의 과도한 사용으로 인해 성능 문제가 있을 수 있습니다. 그러나 이 논문의 저자들은 fgknn의 **병합 큐(merge queue) 구조**에서 사용되는 **병렬 병합(parallel merges)** 개념에서 영감을 얻었다고 명시합니다. 이 영감이 이후 섹션에서 소개될 WarpSelect 알고리즘의 핵심 설계 원리로 발전하게 됩니다. + +요약하면, 기존의 GPU $k$-선택 접근법들은 다중 패스 요구(기수/버킷 선택), 데이터 의존적 메모리 이동(퀵셀렉트), 워프 발산과 직렬성(힙 기반), 과도한 동기화 오버헤드(병렬 우선순위 큐, fgknn) 등의 문제를 갖고 있습니다. 이러한 한계들이 바로 레지스터 파일에서 모든 상태를 유지하면서 단일 패스로 동작하는 새로운 $k$-선택 알고리즘의 필요성을 동기 부여하며, 앞서 분석한 GPU 메모리 계층의 250:6.25:1 대역폭 비율이 레지스터 기반 설계의 잠재적 이점을 정량적으로 뒷받침합니다. +## GPU에서의 빠른 $k$-선택 + +GPU 기반 유사도 검색에서 $k$-선택의 성능은 궁극적으로 메모리 대역폭이나 산술 처리량 중 하나에 의해 제한됩니다. 이는 [roofline 성능 모델](https://dl.acm.org/doi/10.1145/1498765.1498785)에서 제시하는 핵심 원리입니다. 글로벌 메모리에서 입력을 읽는 경우, $k$-선택은 최대 메모리 대역폭으로 입력을 한 번 스캔하는 시간보다 빠를 수 없습니다. 이 논문의 목표는 이 이론적 한계에 최대한 근접하는 것이며, 이를 위해 입력 데이터에 대한 **단일 패스(single pass)** 처리를 수행하고, 모든 중간 상태를 가장 빠른 메모리인 **레지스터 파일**에 유지하는 전략을 채택합니다. 레지스터 메모리의 주요 제약은 레지스터 파일에 대한 인덱싱이 어셈블리 시점에 결정되어야 한다는 것이며, 이 제약이 알고리즘 설계에 강한 구조적 조건을 부과합니다. + +### 레지스터 내 정렬 + +WarpSelect 알고리즘의 핵심 빌딩 블록은 레지스터 내에서 동작하는 정렬 프리미티브입니다. SIMD 아키텍처에서는 벡터 병렬성을 활용하는 **정렬 네트워크(sorting networks)**가 널리 사용되며, 이 논문에서는 레인-스트라이드 레지스터 배열 위에 정렬 네트워크를 구축합니다. 기본 구조로는 Batcher의 바이토닉 정렬 네트워크(bitonic sorting network)의 변형이 사용됩니다. 바이토닉 정렬 네트워크는 크기 $2^k$인 배열에 대한 병렬 병합의 집합으로, 각 병합은 길이 $t$인 $s$개의 배열($s$와 $t$는 2의 거듭제곱)을 길이 $2t$인 $s/2$개의 배열로 변환하며, $\log_2(t)$개의 병렬 단계를 사용합니다. 이 병합을 재귀적으로 적용하면, 길이 $\ell$인 배열을 정렬하는 데 총 + +$$\frac{1}{2}(\log_2(\ell)^2 + \log_2(\ell))$$ + +개의 병렬 병합 단계가 필요합니다. 직관적으로 이해하면, 길이 1인 $\ell$개의 배열에서 시작하여 길이 2인 $\ell/2$개의 배열로, 다시 길이 4인 $\ell/4$개의 배열로 점진적으로 병합해 나가는 과정입니다. + +### 홀수 크기 병합 및 정렬 네트워크 + +실제 응용에서는 데이터가 항상 2의 거듭제곱 크기를 갖지 않으며, 일부 입력이 이미 정렬되어 있을 수도 있습니다. 이러한 상황을 효율적으로 처리하기 위해 **홀수 크기 병합 네트워크(odd-size merging network)**가 도입됩니다. Algorithm 1에 제시된 `merge-odd` 함수는 이미 정렬된 임의 길이의 좌측 배열 $[L_i]_{i=0:\ell_L}$과 우측 배열 $[R_i]_{i=0:\ell_R}$을 병합합니다. + +표준 바이토닉 네트워크가 바이토닉 시퀀스(한 방향으로 증가했다가 감소하는 시퀀스)를 병합하는 반면, 여기서는 **단조 시퀀스(monotonic sequences)**, 즉 같은 방향으로 정렬된 시퀀스를 입력으로 받습니다. 바이토닉 병합을 단조 병합으로 변환하는 핵심 기법은 **첫 번째 비교기 단계를 반전(inverting)**시키는 것입니다. 구체적으로, 첫 번째 단계에서 $i \leftarrow 0:\min(\ell_L, \ell_R)$에 대해 $L_{\ell_L - i - 1}$과 $R_i$를 비교-교환(compare-swap)합니다. 좌측 배열의 끝에서부터, 우측 배열의 처음부터 쌍을 이루어 비교하는 이 반전된 패턴이 단조 입력을 올바르게 처리하는 열쇠입니다. + +홀수 크기 처리의 핵심 아이디어는 배열을 다음으로 큰 2의 거듭제곱 크기로 **더미 원소(dummy elements)**로 패딩한 것으로 간주하되, 더미 원소는 절대 교환되지 않고 이미 올바른 위치에 있으므로 관련 비교를 생략(elide)하는 것입니다. 좌측 배열은 시작 부분에, 우측 배열은 끝 부분에 더미 원소가 패딩된 것으로 취급됩니다. `merge-odd-continue` 함수는 이 재귀적 병합을 수행하며, 각 단계에서 $h \leftarrow 2^{\lceil \log_2 \ell \rceil - 1}$ (즉, $\ell$보다 작은 가장 큰 2의 거듭제곱)을 계산하여 비교-교환 스트라이드를 결정합니다. 좌측과 우측에 따라 재귀 분할 방식이 달라지는데, 이는 더미 원소의 위치가 다르기 때문입니다. + +![홀수 크기 네트워크 병합](https://ar5iv.labs.arxiv.org//html/1702.08734/assets/x1.png) + +위 그림은 크기 5와 3인 두 배열의 홀수 크기 네트워크 병합 과정을 보여줍니다. 원형 점(bullet)은 병렬 비교-교환 연산을 나타내고, 점선은 더미 원소로 인해 생략된 비교를 표시합니다. 이 병합은 4개의 병렬 단계로 완료되며, 일반적으로 길이 $\ell_L$과 $\ell_R$인 두 정렬된 배열의 병합에는 + +$$\lceil \log_2(\max(\ell_L, \ell_R)) \rceil + 1$$ + +개의 병렬 단계가 필요합니다. 비교-교환 연산의 실제 구현에서, 스트라이드가 32의 배수인 교환은 동일 레인 내에서 직접 수행되고(해당 레인이 두 원소를 모두 로컬로 보유하므로), 스트라이드가 16 이하이거나 32의 배수가 아닌 교환은 워프 셔플을 통해 수행됩니다. + +Algorithm 2의 `sort-odd`는 이 병합을 완전한 정렬로 확장합니다. 배열을 절반으로 분할하여 각각을 재귀적으로 정렬한 후 `merge-odd`로 병합하며, 입력 데이터에 구조가 없다고 가정할 때 길이 $\ell$의 데이터 정렬에 + +$$\frac{1}{2}(\lceil \log_2(\ell) \rceil^2 + \lceil \log_2(\ell) \rceil)$$ + +개의 병렬 단계가 필요합니다. + +### WarpSelect + +이 논문의 핵심 기여인 WarpSelect는 모든 상태를 레지스터에 유지하면서 단일 패스로 $k$-선택을 수행하며, 워프 간 동기화가 불필요합니다. $k \leq 1024$까지 지원하며, 각 워프는 $n$개의 배열 $[a_i]$ 중 하나에 대한 $k$-선택을 전담합니다. + +![WarpSelect 개요](https://ar5iv.labs.arxiv.org//html/1702.08734/assets/x2.png) + +위 그림은 WarpSelect의 전체 구조를 보여줍니다. 입력 값이 왼쪽에서 스트리밍되어 들어오고, 오른쪽의 워프 큐가 최종 출력 결과를 보유합니다. + +WarpSelect는 두 가지 핵심 데이터 구조를 사용합니다. 첫째, 각 레인 $j$는 $t$개의 원소로 구성된 **스레드 큐(thread queue)** $[T_i^j]_{i=0:t}$를 레지스터에 유지하며, 가장 큰 값에서 가장 작은 값 순서로 정렬됩니다($T_i^j \geq T_{i+1}^j$). 스레드 큐는 새로 들어오는 값에 대한 1차 필터 역할을 합니다. 새로운 $a_{32i+j}$가 큐의 최대값 $T_0^j$보다 크면, 이 값은 최종 $k$개 최소값에 포함될 수 없으므로 즉시 기각됩니다. 둘째, 워프 전체가 공유하는 **워프 큐(warp queue)** $[W_i]_{i=0:k}$는 레인-스트라이드 레지스터 배열로, 지금까지 관찰된 $k$개의 최소 원소를 최소에서 최대 순서로 유지합니다($W_i \leq W_{i+1}$). 두 큐 모두 $+\infty$ 센티넬 값으로 초기화됩니다. + +Algorithm 3에 제시된 갱신 과정은 세 가지 불변 조건을 유지합니다. 모든 레인의 $T_0^j$가 최소-$k$에 포함되지 않아야 하고, 모든 $T_0^j$가 모든 워프 큐 키 $W_i$보다 커야 하며, 지금까지 관찰된 최소-$k$ 값이 어떤 레인의 스레드 큐 또는 워프 큐에 반드시 포함되어야 합니다. + +갱신 시 레인 $j$가 새로운 $a_{32i+j}$를 받으면, $a_{32i+j} > T_0^j$이면 즉시 기각하고, 그렇지 않으면 스레드 큐의 적절한 위치에 삽입하여 기존 $T_0^j$를 밀어냅니다. 모든 레인이 이 작업을 완료한 후, **워프 발로트(warp ballot)** 명령어로 두 번째 불변 조건 위반 여부를 확인합니다. 위반이 없으면 다음 원소 처리를 계속하고, 위반이 있으면 스레드 큐와 워프 큐를 `odd-merge`로 병합하여 불변 조건을 복원합니다. + +불변 조건 복원 과정에서, 32개의 이미 정렬된 스레드 큐를 레인-스트라이드 배열로 재해석한 후 `sort-odd`로 처음부터 정렬합니다. 구조체 배열(struct-of-arrays)에서 배열 구조체(array-of-structs)로의 전치(transposition)를 통해 이미 정렬된 스레드 큐를 `odd-merge`로 병합하는 것도 가능하지만, 이 전치에 비슷한 수의 워프 셔플이 필요하므로 처음부터 정렬하는 방식을 선택합니다. 정렬된 스레드 큐 집합체와 워프 큐의 병합에는 `odd-merge`를 사용하여 상당한 속도 향상을 달성합니다. + +### 복잡도와 파라미터 선택 + +들어오는 32개 원소 그룹마다 WarpSelect는 1, 2, 또는 3가지 상수 시간 연산을 수행합니다. 32개 원소를 읽고 모든 스레드 큐 헤드 $T_0^j$와 비교하는 비용 $C_1$이 $N_1 = \ell/32$번 발생합니다. 어떤 레인에서 $a_{32n+j} < T_0^j$이면 해당 스레드 큐에 삽입 정렬을 수행하는 비용 $C_2 = \mathcal{O}(t)$가 $N_2$번 발생합니다. 어떤 레인에서 $T_0^j < W_{k-1}$이면 큐 정렬 및 병합 비용 $C_3 = \mathcal{O}(t\log(32t)^2 + k\log(\max(k, 32t)))$가 $N_3$번 발생합니다. + +총 비용은 $N_1 C_1 + N_2 C_2 + N_3 C_3$이며, 독립적으로 추출된 랜덤 데이터에서 $N_2 = \mathcal{O}(k\log(\ell))$, $N_3 = \mathcal{O}(k\log(\ell)/t)$입니다. 핵심 트레이드오프는 $N_2 C_2$와 $N_3 C_3$ 사이의 균형입니다. $t$를 키우면 $N_3$가 줄어들어 비용이 큰 정렬-병합 횟수가 감소하지만, $C_2$가 증가하여 삽입 비용이 커집니다. 실험적으로 결정된 최적 파라미터는 $k \leq 32$일 때 $t = 2$, $k \leq 128$일 때 $t = 3$, $k \leq 256$일 때 $t = 4$, $k \leq 1024$일 때 $t = 8$이며, 모두 $\ell$에 무관합니다. +## 계산 레이아웃 + +이 섹션에서는 Product Quantization을 기반으로 구축된 IVFADC 인덱싱 방법이 GPU에서 어떻게 효율적으로 구현되는지를 상세히 다룹니다. 거리 계산의 세부 사항과 $k$-선택과의 연계가 이 방법이 보다 최근의 GPU 호환 근사 최근접 이웃 전략을 능가할 수 있는 이유를 이해하는 핵심입니다. + +### 정확 검색 + +앞서 문제 정의에서 소개한 L2 거리 분해를 실제 GPU에서 구현할 때, cuBLAS 라이브러리의 최적화된 GEMM 루틴을 사용하여 $-2\langle x_j, y_i \rangle$ 항을 계산하고, 이 결과를 부분 거리 행렬 $D'$에 저장합니다. 이 구현의 핵심 최적화는 **커널 융합(kernel fusion)** 전략입니다. $\|y_i\|^2$ 항을 거리 행렬의 각 항목에 더하는 연산과 $k$-선택을 하나의 융합 커널로 결합합니다. 이 융합 커널은 $D'$의 각 항목을 읽으면서 $\|y_i\|^2$를 즉석에서 더한 뒤, 그 값을 곧바로 레지스터 내 $k$-선택에 제출합니다. $\|x_j\|^2$ 항은 동일 쿼리에 대해 모든 데이터베이스 벡터에 공통이므로 $k$-선택 이전에 고려할 필요가 없습니다. 이 융합 전략 덕분에 $D'$에 대해 GEMM 쓰기와 $k$-선택 읽기의 **단 2회 패스**만 필요하며, 다른 구현에서 요구하는 3회 이상의 패스와 대비됩니다. 행 단위 $k$-선택을 잘 튜닝된 GEMM 커널과 융합하는 것은 불가능하거나 전체 효율을 저하시킬 수 있으므로, 별도의 융합 $k$-선택 커널이 사용됩니다. + +현실적인 문제 크기에서 $D'$는 GPU 메모리에 들어가지 않으므로, 쿼리 배치에 대해 **타일링(tiling)**을 적용합니다. $t_q \leq n_q$개의 쿼리를 하나의 타일로 처리하며, $\lceil n_q / t_q \rceil$개의 타일은 각각 독립적인 문제입니다. GPU 활용도를 높이기 위해 두 개의 타일을 서로 다른 스트림에서 병렬 실행하므로, $D$의 실효 메모리 요구량은 $\mathcal{O}(2\ell t_q)$입니다. 데이터베이스 크기 $\ell$에 대해서도 유사하게 타일링할 수 있으며, CPU에서 대량의 입력이 들어오는 경우에는 고정 메모리(pinned memory) 버퍼링을 통해 CPU→GPU 복사와 GPU 연산을 중첩시킵니다. + +### IVFADC 인덱싱: 룩업 테이블 메커니즘과 최적화 + +IVFADC의 핵심은 쿼리 벡터와 Product Quantization 재생산 값 집합 사이의 거리를 효율적으로 계산하는 것입니다. 앞서 정의한 IVFADC의 양자화 구조를 전개하면, 데이터베이스 벡터 $y$에 대한 근사 거리는 다음과 같이 표현됩니다. + +$$\|x - q(y)\|_2^2 = \|x - q_1(y) - q_2(y - q_1(y))\|_2^2$$ + +조대 양자화기 $q_1$ 이후의 잔차 벡터를 $b$개의 부분 벡터로 분해하면, $y - q_1(y) = [\widetilde{y^1} \cdots \widetilde{y^b}]$이고 $x - q_1(y) = [\widetilde{x^1} \cdots \widetilde{x^b}]$로 쓸 수 있으며, 거리는 다음과 같이 부분 거리의 합으로 분해됩니다. + +$$\|x - q(y)\|_2^2 = \|\widetilde{x^1} - q^1(\widetilde{y^1})\|_2^2 + \ldots + \|\widetilde{x^b} - q^b(\widetilde{y^b})\|_2^2$$ + +각 부분 양자화기 $q^1, \ldots, q^b$는 256개의 재생산 값을 가지므로, 쿼리 $x$와 $q_1(y)$가 알려진 상태에서 모든 가능한 부분 거리를 사전 계산하여 크기 256인 테이블 $T_1, \ldots, T_b$에 저장할 수 있습니다. 이후 실제 거리 계산은 $b$번의 테이블 조회와 덧셈으로 귀결됩니다. $n$개의 거리를 계산하는 비용을 비교하면, 명시적 계산은 $n \times d$번의 곱셈-덧셈이 필요한 반면, 룩업 테이블 방식은 $256 \times d$번의 곱셈-덧셈(테이블 구축)과 $n \times b$번의 조회-덧셈(거리 계산)만 필요합니다. $n$이 수만에서 수십만에 달하는 역색인 리스트 크기를 고려하면, 이 차이는 극적입니다. + +역색인 리스트 $\mathcal{I}_L$ 내에서는 $q_1(y)$가 상수이므로 룩업 테이블 방법이 직접 적용됩니다. 테이블 구축을 더욱 최적화하기 위해, 거리 표현을 세 개의 항으로 분해합니다. + +$$\underbrace{\|q_2(\ldots)\|_2^2 + 2\langle q_1(y), q_2(\ldots)\rangle}_{\text{term 1}} + \underbrace{\|x - q_1(y)\|_2^2}_{\text{term 2}} - 2\underbrace{\langle x, q_2(\ldots)\rangle}_{\text{term 3}}$$ + +이 분해의 핵심 통찰은 각 항의 계산 시점이 다르다는 것입니다. Term 1은 쿼리에 독립적이므로 양자화기로부터 사전 계산하여 크기 $|\mathcal{C}_1| \times 256 \times b$인 테이블 $\mathcal{T}$에 저장할 수 있습니다. Term 2는 $q_1$의 재생산 값까지의 거리로, 1단계 양자화기의 부산물입니다. Term 3은 역색인 리스트에 독립적이며 $d \times 256$번의 곱셈-덧셈으로 계산됩니다. + +이 분해의 효율성을 구체적으로 살펴보면, 단일 쿼리가 $\tau$개의 역색인 리스트를 탐색할 때, 테이블을 처음부터 구축하면 $\tau \times d \times 256$번의 곱셈-덧셈이 필요합니다. 반면 이 분해를 사용하면 $256 \times d$번의 곱셈-덧셈과 $\tau \times b \times 256$번의 덧셈만 필요합니다. $d$가 $b$보다 훨씬 크므로($d = 128$, $b = 8$이면 16배 차이), $\tau$가 클수록 절감 효과가 극대화됩니다. 다만 GPU에서는 $\mathcal{T}$의 메모리 사용량이 과도할 수 있으므로, 메모리가 충분한 경우에만 이 분해를 활성화합니다. + +```python +import numpy as np + +def build_ivfadc_lookup_tables(x, q1_centroid, codebooks, b, precomputed_table=None): + """ + IVFADC 룩업 테이블 구축 (3-term 분해 활용) + x: 쿼리 벡터 (d,) + q1_centroid: 해당 역색인 리스트의 조대 중심점 (d,) + codebooks: b개의 PQ 코드북, 각각 (256, d/b) + precomputed_table: Term 1 사전 계산 테이블 (256, b) 또는 None + """ + d = len(x) + dsub = d // b + residual = x - q1_centroid # x - q1(y) + + tables = np.zeros((b, 256), dtype=np.float32) + + for j in range(b): + sub_residual = residual[j * dsub : (j + 1) * dsub] + # Term 3: -2의 부분 벡터 기여분 + sub_x = x[j * dsub : (j + 1) * dsub] + term3_partial = -2.0 * codebooks[j] @ sub_x # (256,) + + if precomputed_table is not None: + # Term 1은 사전 계산됨, Term 2는 ||x - q1(y)||^2 (스칼라, 별도 처리) + tables[j] = precomputed_table[:, j] + term3_partial + else: + # 분해 없이 직접 계산: ||sub_residual - centroid_k||^2 + tables[j] = np.sum((codebooks[j] - sub_residual) ** 2, axis=1) + + return tables # (b, 256) - 각 부분 양자화기별 256개 거리값 + +def scan_inverted_list(tables, pq_codes_in_list): + """ + 역색인 리스트 스캔: 룩업 테이블로 거리 계산 + tables: (b, 256) 룩업 테이블 + pq_codes_in_list: (n_vectors, b) PQ 코드 배열 + """ + b = tables.shape[0] + n = pq_codes_in_list.shape[0] + distances = np.zeros(n, dtype=np.float32) + + for i in range(n): + for j in range(b): + # 핵심: 곱셈 없이 테이블 조회와 덧셈만으로 거리 계산 + distances[i] += tables[j, pq_codes_in_list[i, j]] + + return distances +``` + +이 코드에서 `scan_inverted_list` 함수의 내부 루프가 바로 GPU에서 수조 번 실행되는 핵심 연산이며, 각 반복이 단순한 테이블 조회와 덧셈으로만 구성되어 있다는 점이 PQ의 효율성의 본질입니다. + +### GPU 구현 + +Algorithm 4는 IVFPQ 배치 검색의 전체 과정을 요약합니다. 각 쿼리 $x_i$에 대해 먼저 앞서 설명한 정확 검색 방식으로 $\tau$개의 가장 가까운 조대 중심점 $L_{\text{IVF}}^i$를 찾고, 이후 해당 역색인 리스트들을 순회하면서 룩업 테이블 기반 거리 계산과 $k$-선택을 수행합니다. 역색인 리스트는 PQ 코드와 연관 ID를 별도의 배열로 저장하며, ID는 $k$-선택이 $k$-최근접 멤버십을 결정한 경우에만 참조됩니다. 이 ID 조회는 대규모 배열에서의 희소한 메모리 읽기이므로, ID를 CPU 메모리에 저장해도 성능 비용이 미미합니다. +## 계산 레이아웃 + +이 섹션에서는 Product Quantization을 기반으로 구축된 IVFADC 인덱싱 방법이 GPU에서 어떻게 효율적으로 구현되는지를 상세히 다룹니다. 거리 계산의 세부 사항과 $k$-선택과의 연계가 이 방법이 보다 최근의 GPU 호환 근사 최근접 이웃 전략을 능가할 수 있는 이유를 이해하는 핵심입니다. + +### 정확 검색 + +앞서 문제 정의에서 소개한 L2 거리 분해를 실제 GPU에서 구현할 때, cuBLAS 라이브러리의 최적화된 GEMM 루틴을 사용하여 $-2\langle x_j, y_i \rangle$ 항을 계산하고, 이 결과를 부분 거리 행렬 $D'$에 저장합니다. 이 구현의 핵심 최적화는 **커널 융합(kernel fusion)** 전략입니다. $\|y_i\|^2$ 항을 $D'$의 각 항목에 더하는 연산과 $k$-선택을 하나의 융합 커널로 결합하여, $D'$의 각 항목을 읽으면서 $\|y_i\|^2$를 즉석에서 더한 뒤 그 값을 곧바로 레지스터 내 $k$-선택에 제출합니다. $\|x_j\|^2$ 항은 동일 쿼리에 대해 모든 데이터베이스 벡터에 공통이므로 $k$-선택 이전에 고려할 필요가 없습니다. 이 융합 전략 덕분에 $D'$에 대해 GEMM 쓰기와 $k$-선택 읽기의 **단 2회 패스**만 필요하며, 다른 구현에서 요구하는 3회 이상의 패스와 대비됩니다. 행 단위 $k$-선택을 잘 튜닝된 GEMM 커널과 융합하는 것은 불가능하거나 전체 효율을 저하시킬 수 있으므로, 별도의 융합 $k$-선택 커널이 사용됩니다. + +현실적인 문제 크기에서 $D'$는 GPU 메모리에 들어가지 않으므로, 쿼리 배치에 대해 **타일링(tiling)**을 적용합니다. $t_q \leq n_q$개의 쿼리를 하나의 타일로 처리하며, $\lceil n_q / t_q \rceil$개의 타일은 각각 독립적인 문제입니다. GPU 활용도를 높이기 위해 두 개의 타일을 서로 다른 스트림에서 병렬 실행하므로, $D$의 실효 메모리 요구량은 $\mathcal{O}(2\ell t_q)$입니다. 데이터베이스 크기 $\ell$에 대해서도 유사하게 타일링할 수 있으며, CPU에서 대량의 입력이 들어오는 경우에는 고정 메모리(pinned memory) 버퍼링을 통해 CPU→GPU 복사와 GPU 연산을 중첩시킵니다. + +### IVFADC 인덱싱: 룩업 테이블 메커니즘과 최적화 + +IVFADC의 핵심은 쿼리 벡터와 Product Quantization 재생산 값 집합 사이의 거리를 효율적으로 계산하는 것입니다. 앞서 정의한 양자화 구조를 전개하면, 데이터베이스 벡터 $y$에 대한 근사 거리는 다음과 같이 표현됩니다. + +$$\|x - q(y)\|_2^2 = \|x - q_1(y) - q_2(y - q_1(y))\|_2^2$$ + +조대 양자화기 $q_1$ 이후의 잔차 벡터를 $b$개의 부분 벡터로 분해하면, $y - q_1(y) = [\widetilde{y^1} \cdots \widetilde{y^b}]$이고 $x - q_1(y) = [\widetilde{x^1} \cdots \widetilde{x^b}]$로 쓸 수 있으며, 거리는 부분 거리의 합으로 분해됩니다. + +$$\|x - q(y)\|_2^2 = \|\widetilde{x^1} - q^1(\widetilde{y^1})\|_2^2 + \ldots + \|\widetilde{x^b} - q^b(\widetilde{y^b})\|_2^2$$ + +각 부분 양자화기가 256개의 재생산 값을 가지므로, 쿼리 $x$와 $q_1(y)$가 알려진 상태에서 모든 가능한 부분 거리를 사전 계산하여 크기 256인 테이블 $T_1, \ldots, T_b$에 저장할 수 있습니다. $n$개의 거리를 계산할 때, 명시적 계산은 $n \times d$번의 곱셈-덧셈이 필요한 반면, 룩업 테이블 방식은 $256 \times d$번의 곱셈-덧셈(테이블 구축)과 $n \times b$번의 조회-덧셈만 필요합니다. + +테이블 구축을 더욱 최적화하기 위해, 거리 표현을 세 개의 항으로 분해합니다. + +$$\underbrace{\|q_2(\ldots)\|_2^2 + 2\langle q_1(y), q_2(\ldots)\rangle}_{\text{term 1}} + \underbrace{\|x - q_1(y)\|_2^2}_{\text{term 2}} - 2\underbrace{\langle x, q_2(\ldots)\rangle}_{\text{term 3}}$$ + +Term 1은 쿼리에 독립적이므로 크기 $|\mathcal{C}_1| \times 256 \times b$인 테이블 $\mathcal{T}$에 사전 저장할 수 있고, Term 2는 1단계 양자화기의 부산물이며, Term 3은 역색인 리스트에 독립적으로 $d \times 256$번의 곱셈-덧셈으로 계산됩니다. 단일 쿼리가 $\tau$개의 리스트를 탐색할 때, 처음부터 구축하면 $\tau \times d \times 256$번의 곱셈-덧셈이 필요하지만, 이 분해를 사용하면 $256 \times d$번의 곱셈-덧셈과 $\tau \times b \times 256$번의 덧셈으로 줄어듭니다. 다만 GPU에서는 $\mathcal{T}$의 메모리 사용량이 과도할 수 있으므로, 메모리가 충분한 경우에만 이 분해를 활성화합니다. + +### GPU 구현 + +Algorithm 4는 IVFPQ 배치 검색의 전체 과정을 요약합니다. 각 쿼리에 대해 먼저 정확 검색 방식으로 $\tau$개의 가장 가까운 조대 중심점을 찾고, 이후 해당 역색인 리스트들을 순회하면서 룩업 테이블 기반 거리 계산과 $k$-선택을 수행합니다. 역색인 리스트는 PQ 코드와 연관 ID를 별도의 배열로 저장하며, ID는 $k$-선택이 $k$-최근접 멤버십을 결정한 경우에만 참조됩니다. 이 ID 조회는 대규모 배열에서의 희소한 메모리 읽기이므로, ID를 CPU 메모리에 저장해도 성능 비용이 미미합니다. + +#### 리스트 스캐닝과 공유 메모리 활용 + +리스트 스캐닝 커널은 각 쿼리에 대해 $\tau$개의 가장 가까운 역색인 리스트를 순회하면서, 룩업 테이블 $T_i$를 사용하여 벡터 쌍별 거리를 계산합니다. 이 과정에서 $T_i$는 **공유 메모리(shared memory)**에 저장됩니다. 전체 쿼리 집합에 대해 최대 $n_q \times \tau \times \max_i |\mathcal{I}_i| \times b$번의 테이블 조회가 필요하며, 실제로는 수조(trillions) 회에 달하는 랜덤 접근이 발생합니다. 공유 메모리의 크기 제약으로 인해 $b$의 최대값은 32비트 부동소수점 사용 시 48, 16비트 부동소수점 사용 시 96으로 제한됩니다. 이 제한은 공유 메모리가 블록당 48 KiB라는 하드웨어 제약에서 직접 도출됩니다. 32비트 부동소수점으로 $b$개의 테이블(각 256개 항목)을 저장하면 $b \times 256 \times 4$ 바이트가 필요하므로, $48 \times 256 \times 4 = 49{,}152$ 바이트 $= 48$ KiB가 정확히 한계에 해당합니다. 앞서 소개한 3-term 분해를 사용하지 않는 경우, $T_i$는 스캐닝 전에 별도의 커널에서 계산됩니다. + +#### 멀티패스 $k$-선택 전략 + +$n_q \times \tau$개의 쿼리-역색인 리스트 쌍은 각각 독립적으로 처리될 수 있습니다. 한 극단에서는 각 쌍에 하나의 블록을 할당하여 최대 $n_q \times \tau \times \max_i |\mathcal{I}_i|$개의 부분 결과를 글로벌 메모리에 기록한 후, 이를 $k$-선택하여 $n_q \times k$개의 최종 결과를 얻습니다. 이 방식은 높은 병렬성을 제공하지만 글로벌 메모리를 초과할 수 있으므로, 정확 검색에서와 동일한 타일링 원리를 적용하여 메모리 복잡도를 $\mathcal{O}(2t_q \tau \max_i |\mathcal{I}_i|)$로 제한합니다. + +단일 워프가 각 $t_q$개 쿼리의 전체 리스트 집합에 대한 $k$-선택을 전담하면 병렬성이 낮아질 수 있습니다. 이를 해결하기 위해 **2-패스 $k$-선택**이 도입됩니다. 첫 번째 패스에서 $t_q \times \tau \times \max_i |\mathcal{I}_i|$개의 부분 결과를 분할 인자 $f$를 사용하여 $t_q \times f \times k$개로 축소하고, 두 번째 패스에서 이를 다시 $k$-선택하여 최종 $t_q \times k$개의 결과를 얻습니다. 이 2-패스 전략은 병렬성과 메모리 사용량 사이의 균형을 효과적으로 달성합니다. + +#### 융합 커널의 트레이드오프 + +정확 검색에서 효과적이었던 커널 융합 전략을 IVFADC에도 적용하는 실험이 수행되었습니다. 단일 블록이 하나의 쿼리에 대한 모든 $\tau$개 리스트를 스캔하면서 거리 계산과 $k$-선택을 융합하는 방식입니다. 앞서 소개한 WarpSelect가 공유 메모리 자원을 경쟁하지 않으므로 이러한 융합이 가능합니다. 이 방식은 거의 모든 중간 결과를 제거하여 글로벌 메모리 기록을 줄여줍니다. 그러나 정확 검색과 달리, IVFADC에서는 런타임의 상당 부분이 공유 메모리의 $T_i$ 테이블에서의 수집(gather)과 글로벌 메모리의 $\mathcal{I}_i$ 선형 스캔에 소비되며, 글로벌 메모리 기록이 지배적 비용이 아닙니다. 따라서 융합 커널의 성능 향상은 최대 15%에 그치며, 일부 문제 크기에서는 낮은 병렬성으로 인해 오히려 성능이 저하될 수 있습니다. 이러한 이유와 구현 단순성을 위해 융합 레이아웃은 채택되지 않았습니다. + +이 설계 결정은 중요한 통찰을 제공합니다. 정확 검색에서는 GEMM 이후의 글로벌 메모리 기록이 병목이므로 커널 융합이 큰 효과를 발휘하지만, IVFADC에서는 공유 메모리 랜덤 접근과 역색인 리스트의 순차 읽기가 병목이므로 융합의 이점이 제한적입니다. 즉, 동일한 최적화 전략이라도 연산의 병목 지점에 따라 효과가 크게 달라진다는 것입니다. + +### 다중 GPU 병렬 처리 + +현대 서버는 여러 GPU를 지원하며, 이를 연산 능력과 메모리 양쪽 모두에 활용할 수 있습니다. 이 논문에서는 **복제(replication)**와 **샤딩(sharding)**이라는 두 가지 상호 보완적인 전략을 제시합니다. + +복제는 인덱스가 단일 GPU 메모리에 들어가는 경우에 적용됩니다. 동일한 인덱스를 $\mathcal{R}$개의 GPU에 복제하고, $n_q$개의 쿼리를 각 복제본이 $n_q / \mathcal{R}$개씩 분담하여 처리합니다. 결과는 단일 GPU 또는 CPU 메모리에서 합쳐집니다. 이 방식은 거의 선형적인 속도 향상을 제공하지만, $n_q$가 작을 때는 효율이 다소 떨어질 수 있습니다. + +샤딩은 인덱스가 단일 GPU 메모리에 들어가지 않는 경우에 적용됩니다. 인덱스를 $\mathcal{S}$개의 GPU에 분산하여 각 샤드가 $\ell / \mathcal{S}$개의 벡터를 보유합니다. 쿼리 시에는 각 샤드가 전체 쿼리 집합 $n_q$를 처리하고, 부분 결과를 합친 후 추가적인 $k$-선택 라운드를 수행합니다. 동일한 인덱스 크기 $\ell$에서 샤딩은 각 샤드가 $\ell / \mathcal{S}$에 대해 $n_q$를 처리하는 반면, 복제는 $\ell$에 대해 $n_q / \mathcal{R}$을 처리하므로 샤딩이 속도 향상을 제공하지만, 고정 오버헤드와 후속 $k$-선택 비용으로 인해 순수 복제보다는 효율이 낮습니다. + +두 전략은 결합하여 사용할 수 있습니다. $\mathcal{S}$개의 샤드 각각을 $\mathcal{R}$개로 복제하면 총 $\mathcal{S} \times \mathcal{R}$개의 GPU를 활용합니다. 이 원리는 다중 머신 분산으로도 자연스럽게 확장됩니다. + +```python +def multi_gpu_search_strategy(n_q, ell, k, num_gpus, single_gpu_memory): + """ + 다중 GPU 검색 전략 선택 및 실행 시뮬레이션 + """ + index_size_bytes = ell * 16 # 예: 벡터당 16바이트 PQ 코드 + + if index_size_bytes <= single_gpu_memory: + # 복제 전략: 인덱스가 단일 GPU에 적재 가능 + R = num_gpus # 모든 GPU에 복제 + S = 1 + queries_per_replica = n_q // R + # 각 복제본: queries_per_replica개 쿼리 × ell개 벡터 + # 결과 합치기: 단순 연결 (추가 k-선택 불필요) + speedup = R * 0.95 # 거의 선형, 소규모 n_q에서 약간 손실 + else: + # 샤딩 필요. 인덱스가 단일 GPU에 적재 불가 + S = max(2, index_size_bytes // single_gpu_memory + 1) + R = num_gpus // S + vectors_per_shard = ell // S + # 각 샤드: n_q개 쿼리 × (ell/S)개 벡터 + # 결과 합치기: S개 샤드의 부분 결과에 대해 추가 k-선택 필요 + # 각 샤드에서 k개 후보 → S*k개 후보에서 최종 k개 선택 + merge_cost = S * k # 추가 k-선택 비용 + speedup = S * R * 0.85 # 고정 오버헤드로 인해 복제보다 낮은 효율 + + return { + 'shards': S, 'replicas': R, + 'total_gpus': S * R, + 'estimated_speedup': speedup + } +``` + +이 코드는 인덱스 크기와 GPU 메모리 용량에 따라 복제와 샤딩 전략을 자동으로 선택하는 로직을 보여줍니다. 핵심은 샤딩 시 각 샤드의 부분 결과를 합친 후 추가적인 $k$-선택이 필요하다는 점이며, 이 추가 비용이 샤딩의 효율을 순수 복제보다 낮추는 원인입니다. +## 실험 및 응용 + +이 섹션에서는 앞서 설계한 GPU $k$-선택 알고리즘과 최근접 이웃 검색 방법의 실제 성능을 기존 라이브러리와 체계적으로 비교합니다. 모든 실험은 별도 명시가 없는 한 2개의 2.8GHz Intel Xeon E5-2680v2 CPU와 4개의 Maxwell Titan X GPU, CUDA 8.0 환경에서 수행되었습니다. + +### $k$-선택 성능 + +WarpSelect의 성능을 평가하기 위해 두 가지 기존 GPU 소규모 $k$-선택 구현과 비교가 수행되었습니다. 첫 번째는 [Tang et al.](https://arxiv.org/abs/1404.1831)의 fgknn 라이브러리에서 추출한 행 기반 Merge Queue with Buffered Search and Hierarchical Partition이며, 두 번째는 [Sismanis et al.](https://dl.acm.org/doi/10.1145/2304576.2304623)의 Truncated Bitonic Sort(TBiS)입니다. 두 방법 모두 각각의 정확 검색 라이브러리에서 추출된 것입니다. + +실험 설정은 단일 Titan X GPU에서 $n_q \times \ell$ 크기의 행 우선(row-major) 행렬로부터 각 행의 $k = 100$ 및 $k = 1000$개 최소값을 선택하는 것입니다. 배치 크기 $n_q$는 10,000으로 고정되고, 배열 길이 $\ell$은 1,000에서 128,000까지 변화합니다. 입력과 출력 모두 GPU 메모리에 상주하며, 출력 크기는 $n_q \times k$입니다. 따라서 입력 문제 크기는 $\ell = 1{,}000$일 때 40 MB에서 $\ell = 128{,}000$일 때 5.12 GB까지 범위를 갖습니다. TBiS는 대규모 보조 저장 공간을 요구하기 때문에 테스트에서 $\ell \leq 48{,}000$으로 제한되었습니다. + +![k-선택 방법별 런타임 비교](https://ar5iv.labs.arxiv.org//html/1702.08734/assets/x3.png) + +위 그림은 배열 길이 $\ell$에 따른 서로 다른 $k$-선택 방법의 런타임을 보여줍니다. 실선은 $k = 100$, 점선은 $k = 1000$에 해당하며, 메모리 대역폭 한계선도 함께 표시되어 있습니다. 가장 큰 배열 길이 $\ell = 128{,}000$에서 WarpSelect는 fgknn 대비 $k = 100$일 때 $1.62\times$, $k = 1000$일 때 $2.01\times$ 빠른 성능을 달성하였습니다. 특히 주목할 점은 $k$가 커질수록 WarpSelect의 상대적 우위가 더 커진다는 것이며, TBiS조차 $k = 1000$에서 큰 $\ell$에 대해 fgknn을 능가하기 시작합니다. + +이론적 최대 성능 대비 효율을 살펴보면, WarpSelect는 $k = 100$에서 Titan X 메모리 대역폭 피크의 55%에서 동작하지만, $k = 1000$에서는 16%로 떨어집니다. 이러한 성능 저하는 큰 $k$에서 스레드 큐 크기 증가와 병합/정렬 네트워크의 추가 오버헤드에 기인합니다. 앞서 소개한 복잡도 분석에서 $C_3 = \mathcal{O}(t\log(32t)^2 + k\log(\max(k, 32t)))$ 항이 $k$에 따라 증가하는 것과 일치하는 결과입니다. + +fgknn과의 핵심적인 차이점으로는, WarpSelect가 모든 상태를 레지스터에 유지하여 공유 메모리를 사용하지 않는다는 점, 워프 간 동기화나 버퍼링이 불필요하다는 점, 계층적 파티셔닝이 없다는 점, 다른 커널과 융합이 가능하다는 점, 그리고 효율적인 병합과 정렬을 위해 홀수 크기 네트워크를 사용한다는 점이 있습니다. + +### $k$-means 클러스터링 + +$k = 1$인 정확 검색은 $k$-means 클러스터링의 할당 단계에서 직접 활용될 수 있습니다. $n_q$개의 학습 벡터를 $|\mathcal{C}_1|$개의 중심점에 할당하는 이 작업은 IVFADC를 사용하지 않으며, $k = 1$ 선택은 WarpSelect 대신 병렬 리덕션(parallel reduction)으로 처리됩니다. 그럼에도 $k$-means는 양자화기 $q_1$을 학습하는 데 사용되는 클러스터링의 좋은 벤치마크입니다. + +MNIST8m 데이터셋(810만 장의 28×28 그레이스케일 숫자 이미지, 784차원 벡터로 선형화)에서 실험이 수행되었으며, GPU $k$-means 구현인 [BIDMach](https://github.com/BIDData/BIDMach)와 비교되었습니다. BIDMach는 수십 대의 머신을 요구하는 여러 분산 $k$-means 구현보다 효율적인 것으로 알려진 방법입니다. 두 알고리즘 모두 20회 반복으로 실행되었습니다. + +| 중심점 수 | 방법 | GPU 수 | 실행 시간 | +|-----------|------|--------|----------| +| 256 | BIDMach | 1 | 320 s | +| 256 | 본 논문 | 1 | 140 s | +| 256 | 본 논문 | 4 | 84 s | +| 4096 | BIDMach | 1 | 735 s | +| 4096 | 본 논문 | 1 | 316 s | +| 4096 | 본 논문 | 4 | 100 s | + +위 표에서 확인할 수 있듯이, 본 논문의 구현은 BIDMach 대비 $2\times$ 이상 빠르며, 두 방법 모두 cuBLAS를 기반으로 합니다. 성능 차이의 일부는 $k$-선택을 L2 거리 계산에 융합하는 전략에서 비롯됩니다. 다중 GPU 실행 시 복제 방식을 통해 충분히 큰 문제에서 거의 선형적인 속도 향상이 관찰되었으며, 4096개 중심점에서 4개 GPU 사용 시 $3.16\times$의 속도 향상을 달성하였습니다. + +대규모 실험에서는 $10^8$개의 128차원 벡터를 85,000개 중심점으로 클러스터링하는 작업도 수행되었습니다. 근사 CPU 방법이 클러스터링에 46분, 벡터 인코딩 전처리에 최소 56분(총 102분 이상)이 소요되는 반면, 본 논문의 방법은 4개 GPU에서 전처리 없이 정확 $k$-means를 52분 만에 완료하였습니다. + +### 정확 최근접 이웃 검색 + +최근접 이웃 검색 평가의 고전적 데이터셋인 SIFT1M($\ell = 10^6$, $d = 128$, $n_q = 10^4$)에서 정확 검색 실험이 수행되었습니다. 부분 거리 행렬 $D'$ 계산 비용은 $n_q \times \ell \times d = 1.28$ Tflop이며, 현재 GPU에서 1초 미만에 실행됩니다. + +![SIFT1M 정확 검색 k-NN 시간](https://ar5iv.labs.arxiv.org//html/1702.08734/assets/x4.png) + +위 그림은 SIFT1M 데이터셋에서 $k$를 변화시키며 측정한 정확 검색 시간을 보여줍니다. 앞서 설명한 GEMM 타일링을 통한 $-2\langle x_j, y_i \rangle$ 항 계산 비용과, 거리 행렬 크기 $n_q \times \ell$에 대한 피크 $k$-선택 성능이 함께 표시되어 있습니다. 핵심 관찰 결과로, `thrust::sort_by_key`를 사용하여 전체 결과 배열을 정렬하는 나이브 알고리즘은 비교 대상 방법들보다 $10\times$ 이상 느립니다. 또한 본 논문의 방법을 제외한 모든 구현에서 L2 거리와 $k$-선택 비용이 지배적이며, 본 논문의 방법은 GEMM 사용과 $D'$ 타일링이 최적에 가깝다고 가정할 때 피크 가능 성능의 85%를 달성합니다. 특히 융합 L2/$k$-선택 커널의 중요성이 두드러지는데, 융합 없이 $D'$에 대한 추가 패스가 필요한 동일 알고리즘은 최소 25% 느립니다. + +효율적인 $k$-선택은 근사 방법으로 거리를 계산하는 상황에서 더욱 중요해집니다. 근사 거리 계산의 비용이 정확 계산보다 낮아지면, 전체 파이프라인에서 $k$-선택이 차지하는 상대적 비중이 증가하기 때문입니다. + +### 십억 규모 근사 검색 + +GPU 기반 근사 최근접 이웃 검색에 대한 대규모($\ell \gg 10^6$) 연구는 드문 편입니다. 표준 데이터셋과 평가 프로토콜을 사용하여 인덱스 검색 성능이 비교되었습니다. + +SIFT1M에서 [Wieschollek et al.](https://arxiv.org/pdf/1702.05911v1)의 GPU 검색 속도와 비교한 결과, 동일한 시간 예산(쿼리당 0.02 ms)에서 Wieschollek et al.이 R@1 = 0.51, R@100 = 0.86을 달성한 반면, 본 논문의 구현은 R@1 = 0.80, R@100 = 0.95를 달성하여 동일 시간 내에 훨씬 높은 정확도를 보여줍니다. + +SIFT1B(10억 개 SIFT 이미지 특징, $n_q = 10^4$)에서는 동일한 메모리 사용량 조건에서 비교가 수행되었습니다. 단일 GPU에서 벡터당 $m = 8$ 바이트 인코딩을 사용하여 R@10 = 0.376을 쿼리당 17.7 $\mu$s에 달성한 반면, Wieschollek et al.은 R@10 = 0.35를 쿼리당 150 $\mu$s에 달성하였습니다. 즉, 본 논문의 구현이 더 높은 정확도를 $8.5\times$ 빠른 속도로 달성한 것입니다. + +DEEP1B(10억 개 CNN 이미지 표현, $n_q = 10^4$)에서는 [OPQ](https://arxiv.org/abs/1307.1708)를 통해 $d = 80$으로 차원 축소한 후 $m = 20$의 PQ 인코딩과 $|\mathcal{C}_1| = 2^{18}$을 사용하였습니다. 단일 GPU 메모리에 적재할 수 없으므로 4개 GPU를 $\mathcal{S} = 2$, $\mathcal{R} = 2$ 구성으로 사용하였으며, 원본 논문과 비교 가능한 20 GB의 데이터셋 저장 공간을 사용합니다. R@1 = 0.4517을 쿼리당 0.0133 ms에 달성하였는데, 원본 데이터셋 논문의 CPU 결과(1 스레드)가 R@1 = 0.45를 쿼리당 20 ms에 달성한 것과 비교하면, 하드웨어 플랫폼이 다르지만 GPU 검색이 단일 머신에서 달성 가능한 속도 측면에서 게임 체인저임을 보여줍니다. +### $k$-NN 그래프 구축 + +$k$-NN 그래프는 유사도 검색 방법의 대표적인 응용 사례로, 데이터셋의 모든 벡터를 쿼리로 사용하여 전체 인덱스에 대해 브루트 포스 방식으로 구축됩니다. 이 실험에서는 속도, 정밀도, 메모리 사이의 트레이드오프를 두 가지 대규모 데이터셋에서 평가합니다. + +첫 번째 데이터셋은 Yfcc100M에서 추출한 9,500만 장의 이미지이며, [ResNet](https://arxiv.org/abs/1512.03385) 아키텍처의 마지막에서 두 번째 레이어에서 CNN 디스크립터를 추출한 후 PCA로 $d = 128$차원으로 축소하였습니다. 두 번째 데이터셋은 앞서 근사 검색 실험에서도 사용된 Deep1B입니다. 평가 지표는 세 가지 축으로 구성됩니다. 속도는 IVFADC 인덱스를 처음부터 구축하고 전체 $k$-NN 그래프($k = 10$)를 완성하는 데 걸리는 종단 간(end-to-end) 시간을 측정하며, 품질은 10,000개의 샘플 이미지에 대해 정확한 최근접 이웃을 계산한 뒤, 발견된 10개의 최근접 이웃 중 실제 10개 최근접 이웃에 포함되는 비율(10-intersection at 10)로 측정됩니다. + +Yfcc100M에서는 $2^{16}$개의 조대 양자화 중심점과 $m = 16, 32, 64$ 바이트 PQ 인코딩을 사용하였고, Deep1B에서는 OPQ를 통해 $d = 120$으로 전처리한 후 $|\mathcal{C}_1| = 2^{18}$과 $m = 20, 40$을 사용하였습니다. 주어진 인코딩에서 멀티프로브 파라미터 $\tau$를 1에서 256까지 변화시켜 효율성과 품질 사이의 트레이드오프를 얻습니다. + +![k-NN 그래프 구축의 속도/정확도 트레이드오프](https://ar5iv.labs.arxiv.org//html/1702.08734/assets/x5.png) + +![Deep1B k-NN 그래프 구축 결과](https://ar5iv.labs.arxiv.org//html/1702.08734/assets/x6.png) + +위 두 그림은 Yfcc100M과 Deep1B 데이터셋에 대한 브루트 포스 10-NN 그래프 구축의 속도/정확도 트레이드오프를 보여줍니다. Yfcc100M에서는 $\mathcal{S} = 1$, $\mathcal{R} = 4$ 구성(4개 Titan X GPU 전체를 복제본으로 사용)으로 0.8 이상의 정확도를 35분 만에 달성하였습니다. PQ 인코딩 크기가 클수록($m = 64$) 동일 시간 내에 더 높은 정확도를 달성하지만, 메모리 사용량이 증가하는 트레이드오프가 존재합니다. + +Deep1B에서는 낮은 품질의 그래프를 약 6시간에, 높은 품질의 그래프를 약 반나절에 구축할 수 있었습니다. 8개의 Maxwell M40 GPU(Titan X와 대략 동등한 성능)로 복제 세트를 두 배로 늘린 실험에서는 $m = 20$일 때 약 $1.6\times$, $m = 40$일 때 약 $1.7\times$의 준선형(sub-linear) 성능 향상이 관찰되었습니다. 이 준선형 스케일링은 앞서 다중 GPU 병렬 처리에서 설명한 샤딩의 고정 오버헤드와 후속 $k$-선택 비용에 기인합니다. + +기존 연구와의 비교는 이 결과의 의미를 더욱 명확히 합니다. 기존에 알려진 가장 큰 $k$-NN 그래프 구축은 3,650만 개의 384차원 벡터에 대해 128대의 CPU 서버 클러스터에서 [NN-Descent](https://dl.acm.org/doi/10.1145/1963405.1963487) 알고리즘으로 108.7시간이 소요되었습니다. NN-Descent는 본 논문에서 다루는 데이터셋에 대해서도 $k$-NN 그래프를 구축하거나 개선할 수 있지만, 그래프 저장 자체에 대한 큰 메모리 오버헤드가 문제입니다. Deep1B의 경우 그래프 저장만 80 GB이며, 모든 벡터에 대한 랜덤 접근(384 GB)이 추가로 필요합니다. GPU 기반 기존 최대 규모 $k$-NN 그래프 구축은 2,000만 개의 15,000차원 벡터에 대해 32개의 Tesla C2050 GPU 클러스터에서 GEMM 기반 브루트 포스로 10일이 소요되었습니다. GEMM 비용 기준으로 이 접근법을 Deep1B에 적용하면 해당 클러스터에서 약 200일이라는 비실용적인 시간이 필요할 것으로 추정됩니다. + +### $k$-NN 그래프의 활용 + +이미지 데이터셋에 대해 $k$-NN 그래프가 구축되면, 임의의 두 이미지 사이에서 그래프 내 경로를 탐색할 수 있습니다. 이는 그래프가 단일 연결 성분(connected component)을 형성하는 경우에 가능하며, 실험에서 이 조건이 충족되었습니다. 출발 이미지 $S$와 도착 이미지 $D$가 주어졌을 때, 경로 $P = \{p_1, \ldots, p_n\}$ ($p_1 = S$, $p_n = D$)에 대해 다음 목적 함수를 최소화합니다. + +$$\min_P \max_{i=1..n} d_{p_i p_{i+1}}$$ + +이 수식은 경로 상의 연속된 노드 쌍 사이의 최대 거리를 최소화하는 것으로, 직관적으로 말하면 경로를 따라 이동할 때 가장 큰 "점프"를 최소화하여 부드러운 전환(smooth transition)을 선호하는 것입니다. 이는 최단 경로(총 거리 합 최소화)와는 다른 미니맥스(minimax) 기준으로, 시각적으로 자연스러운 이미지 전환을 생성하는 데 적합합니다. + +![k-NN 그래프에서의 이미지 경로](https://ar5iv.labs.arxiv.org//html/1702.08734/assets/figs/flowers.jpg) + +위 그림은 Yfcc100M의 9,500만 장 이미지에 대한 $k$-NN 그래프에서 두 꽃 이미지 사이의 경로를 보여줍니다. 첫 번째와 마지막 이미지가 주어지고, 알고리즘이 그 사이의 가장 부드러운 경로를 계산합니다. $k = 15$ 이웃을 가진 $k$-NN 그래프에서 20초의 전파(propagation) 후에 이 결과가 얻어졌습니다. 데이터셋에 다양한 꽃 이미지가 풍부하게 존재하기 때문에 전환이 매우 자연스럽게 이루어지는 것을 확인할 수 있습니다. 이 응용은 대규모 $k$-NN 그래프가 단순한 검색을 넘어 데이터셋의 구조적 탐색과 시각적 내비게이션에 활용될 수 있음을 보여주는 인상적인 사례입니다. +## 결론 + +이 논문은 GPU의 산술 처리량과 메모리 대역폭이 각각 테라플롭스(teraflops)와 수백 기가바이트/초(hundreds of gigabytes per second)에 달함에도 불구하고, 이러한 성능 수준에 근접하는 알고리즘을 구현하는 것이 복잡하고 직관에 반하는 작업임을 지적하면서 연구의 의의를 정리합니다. + +앞서 상세히 다루어진 WarpSelect 알고리즘, 레인-스트라이드 레지스터 배열 기반의 정렬 네트워크, 커널 융합 전략, 그리고 IVFADC의 룩업 테이블 최적화 등은 모두 GPU에서 유사도 검색의 거의 최적(near-optimal) 성능을 달성하기 위한 알고리즘적 구조를 구성합니다. 이 논문의 핵심 메시지는 이러한 최적화된 GPU 구현이 기존에 복잡한 근사 알고리즘을 필요로 했던 응용들을 근본적으로 변화시킨다는 것입니다. + +구체적으로, 앞서 실험 섹션에서 확인한 바와 같이 정확(exact) $k$-means 클러스터링이나 $k$-NN 그래프 구축과 같은 작업을 단순한 브루트 포스 방식으로 수행하면서도, CPU 단일 머신이나 CPU 클러스터가 근사적으로 동일한 작업을 수행하는 데 걸리는 시간보다 더 짧은 시간 안에 완료할 수 있음이 입증되었습니다. 이는 GPU의 하드웨어 특성을 알고리즘 설계에 깊이 반영한 결과로, 단순히 CPU 알고리즘을 GPU로 이식하는 것이 아니라 GPU의 메모리 계층 구조와 병렬 실행 모델에 맞춤화된 새로운 알고리즘 패러다임을 제시한 것입니다. + +또한 이 논문은 GPU 하드웨어가 머신러닝 알고리즘의 인기 덕분에 과학 워크스테이션에서 이미 매우 보편화되어 있다는 현실적 관찰을 바탕으로, 이러한 GPU가 데이터베이스 응용에도 큰 가치를 지닌다는 점을 강조합니다. 유사도 검색은 정보 검색, 추천 시스템, 이미지 검색, 자연어 처리 등 광범위한 분야에서 핵심 연산으로 활용되므로, GPU 기반의 효율적인 유사도 검색 구현은 이러한 응용 전반에 걸쳐 실질적인 영향을 미칠 수 있습니다. + +이 연구와 함께 논문에서 제안된 알고리즘들의 정교하게 엔지니어링된 구현체가 오픈소스로 공개되었으며, 이것이 바로 FAISS(Facebook AI Similarity Search) 라이브러리입니다. FAISS의 공개는 학술적 기여를 넘어 실무적 파급력을 극대화하는 결정으로, 이후 대규모 벡터 검색 분야에서 사실상의 표준 도구로 자리잡게 되었습니다. 이를 통해 기존에 GPU를 보유하고 있지만 유사도 검색에는 활용하지 못했던 연구자와 엔지니어들이 효율적인 유사도 검색을 즉시 수행할 수 있게 되었습니다. + +종합적으로, 이 논문의 가장 중요한 기여는 GPU 아키텍처의 세밀한 이해를 바탕으로 유사도 검색의 각 구성 요소($k$-선택, 거리 계산, 인덱싱)를 최적화하고, 이들을 유기적으로 결합하여 십억 규모 데이터셋에서도 실용적인 성능을 달성한 것입니다. 레지스터 파일의 극단적인 대역폭 우위를 활용한 WarpSelect, 행렬 곱셈으로의 거리 계산 환원, 룩업 테이블 기반의 압축 도메인 검색, 그리고 복제와 샤딩을 결합한 다중 GPU 전략이 하나의 통합된 시스템으로 구현됨으로써, 정확 검색과 근사 검색 모두에서 기존 방법을 큰 폭으로 능가하는 결과를 달성하였습니다. +## 부록: WarpSelect의 복잡도 분석 + +이 부록에서는 앞서 소개한 WarpSelect 알고리즘에서 갱신(update)이 트리거되는 평균 횟수를 수학적으로 유도합니다. 이 분석은 WarpSelect의 복잡도와 파라미터 선택에 대한 이론적 근거를 제공합니다. + +### 기본 설정과 연속 최소-$k$ 확률 + +$k$-선택의 입력은 서로 다른 원소들의 집합을 무작위로 배열한 순열 $\{a_1, a_2, \ldots, a_\ell\}$ (1-기반 인덱싱)입니다. 원소들은 크기 $w$인 그룹(워프) 단위로 순차적으로 읽히며, $w = 32$이고 $\ell$이 $w$의 배수라고 가정하면 $c = \ell / w$개의 그룹이 존재합니다. $t$는 스레드 큐 길이를 나타냅니다. 위치 $n$까지의 원소들 중 최소-$k$에 해당하는 원소들을 **연속 최소-$k$ (successive min-$k$)**라 정의하며, $a_n$이 위치 $n$에서의 연속 최소-$k$에 포함될 확률은 다음과 같습니다. + +$$\alpha(n, k) := \begin{cases} 1 & \text{if } n \leq k \\ k/n & \text{if } n > k \end{cases}$$ + +이는 모든 순열이 동일한 확률을 가지므로, $n > k$일 때 $a_n$이 지금까지 본 $n$개 원소 중 상위 $k$개에 포함될 확률이 정확히 $k/n$이 되기 때문입니다. + +### 삽입 정렬 횟수 ($N_2$) + +주어진 레인에서 삽입 정렬은 들어오는 값이 연속 최소-$(k+t)$ 값에 포함될 때 트리거됩니다. 해당 레인이 "본" 값의 수는 $wc_0 + (c - c_0)$이며, 여기서 $c_0$는 이전에 승리한 워프 발로트 횟수입니다. 이 사건의 확률은 다음과 같이 근사됩니다. + +$$\alpha(wc_0 + (c - c_0), k + t) \approx \frac{k + t}{wc} \quad \text{for } c > k$$ + +이 근사는 스레드 큐가 자신의 레인에 할당된 값만이 아니라 전체 $wc$개의 값을 모두 관찰한 것으로 간주합니다. 임의의 레인에서 삽입 정렬이 트리거될 확률은 $w$개 레인 중 적어도 하나가 트리거될 확률이므로 다음과 같습니다. + +$$1 - \left(1 - \frac{k+t}{wc}\right)^w \approx \frac{k+t}{c}$$ + +여기서 근사는 1차 테일러 전개에 의한 것입니다. 이 확률을 $c$에 대해 합산하면 삽입의 기대 횟수는 다음과 같습니다. + +$$N_2 \approx (k+t)\log(c) = \mathcal{O}(k\log(\ell/w))$$ + +### 전체 정렬 횟수: 단일 레인 ($w = 1$) + +전체 정렬의 기대 횟수 $N_3 = \pi(\ell, k, t, w)$를 구합니다. 먼저 $w = 1$인 경우를 분석하므로 $c = \ell$입니다. $\gamma(\ell, m, k)$를 길이 $\ell$인 순열에서 순차 스캐너가 정확히 $m$개의 연속 최소-$k$ 원소를 만날 확률이라 정의하면, 이는 다음 점화식으로 주어집니다. + +$$\gamma(\ell, m, k) := \begin{cases} 1 & \ell = 0 \text{ and } m = 0 \\ 0 & \ell = 0 \text{ and } m > 0 \\ 0 & \ell > 0 \text{ and } m = 0 \\ \gamma(\ell-1, m-1, k) \cdot \alpha(\ell, k) + \gamma(\ell-1, m, k) \cdot (1 - \alpha(\ell, k)) & \text{otherwise} \end{cases}$$ + +마지막 경우는 두 가지 상호 배타적 사건의 합입니다. 길이 $\ell - 1$인 접두 순열에 $m - 1$개의 연속 최소-$k$ 원소가 있고 현재 원소가 연속 최소-$k$에 포함되는 경우와, 접두 순열에 이미 $m$개가 있고 현재 원소는 포함되지 않는 경우입니다. + +이제 $\delta(\ell, b, k, t)$를 정확히 $b$번의 발로트 승리를 유발하는 모든 길이 $\ell$ 순열의 비율로 정의합니다. + +$$\delta(\ell, b, k, t) := \sum_{m = bt}^{\min((bt + \max(0, t-1)), \ell)} \gamma(\ell, m, k)$$ + +스레드 큐가 $t$번 채워져야 한 번의 발로트 승리가 발생하므로, $b$번의 정렬이 일어나려면 연속 최소-$k$ 원소의 수 $m$이 $bt$에서 $bt + \max(0, t-1)$ 사이여야 합니다. 최대 $\lfloor \ell / t \rfloor$번의 발로트 승리가 가능하며, $\pi(\ell, k, t, 1)$은 이에 대한 기댓값입니다. + +$$\pi(\ell, k, t, 1) = \sum_{b=1}^{\lfloor \ell / t \rfloor} b \cdot \delta(\ell, b, k, t)$$ + +이 값은 동적 프로그래밍으로 계산할 수 있습니다. + +### 특수 경우의 해석적 결과 + +$t = 1$, $k = 1$인 경우, $\pi(\ell, 1, 1, 1)$은 조화수 $H_\ell = 1 + \frac{1}{2} + \frac{1}{3} + \cdots + \frac{1}{\ell}$이 되며, $\ell \to \infty$에서 $\ln(\ell) + \gamma$ (오일러-마스케로니 상수 $\gamma$)로 수렴합니다. + +$t = 1$, $k > 1$, $\ell > k$인 경우에는 다음이 성립합니다. + +$$\pi(\ell, k, 1, 1) = k + k(H_\ell - H_k)$$ + +처음 $k$개 원소는 모두 연속 최소-$k$에 포함되고, 나머지에 대한 기댓값이 $\frac{k}{k+1} + \frac{k}{k+2} + \cdots + \frac{k}{\ell}$이므로 $\mathcal{O}(k\log(\ell))$입니다. + +$t > 1$, $k > 1$, $\ell > k$인 일반적 경우, 각 순열에 대해 연속 최소-$k$ 판정 횟수 $D$ ($k \leq D \leq \ell$)가 존재하고, 발로트 승리 횟수는 정의상 $\lfloor D / t \rfloor$입니다. 따라서 다음이 성립합니다. + +$$\pi(\ell, k, t, 1) = \mathcal{O}(k\log(\ell) / t)$$ + +### 다중 레인 ($w > 1$) + +$w > 1$인 경우, $w$개의 워커 중 둘 이상이 동시에 정렬을 트리거하더라도 실제로는 한 번의 정렬만 수행되므로 결합 확률을 고려해야 합니다. $\pi'(\ell, k, t, w)$를 $w$개 워커 간 발로트 승리의 상호 간섭이 없다고 가정한 기대 발로트 승리 횟수로 정의합니다. $k \geq w$를 가정하면 다음 상한이 성립합니다. + +$$\pi'(\ell, k, 1, w) \leq w\left(\left\lceil \frac{k}{w} \right\rceil + \sum_{i=1}^{\lceil \ell/w \rceil - \lceil k/w \rceil} \frac{k}{w(\lceil k/w \rceil + i)}\right) \leq w \cdot \pi(\lceil \ell/w \rceil, k, 1, 1) = \mathcal{O}(wk\log(\ell/w))$$ + +이 부등식은 $w$개 워커 각각이 연속 최소-$k$ 원소를 관찰할 확률의 상한을 첫 번째 워커의 확률로 대체한 것입니다. 스레드 큐 길이 $t$에 의한 스케일링을 적용하면 다음을 얻습니다. + +$$\pi'(\ell, k, t, w) = \mathcal{O}(wk\log(\ell/w) / t)$$ + +상호 간섭은 발로트 횟수를 줄이는 방향으로만 작용하므로, $\pi(\ell, k, t, w)$에 대해서도 동일한 상한 $\mathcal{O}(wk\log(\ell/w) / t)$가 성립합니다. 이 결과는 WarpSelect의 전체 정렬 횟수가 스레드 큐 길이 $t$에 반비례하여 감소하며, 입력 길이 $\ell$에 대해 로그적으로만 증가한다는 것을 보장합니다. 이는 앞서 실험에서 관찰된 WarpSelect의 효율적인 성능을 이론적으로 뒷받침하는 핵심 결과입니다. +- - - +### References +* [Billion-scale similarity search with GPUs](https://arxiv.org/pdf/1702.08734v1) \ No newline at end of file