ANN(Approximate Nearest Neighbor) 알고리즘의 이해 #3 DiskANN

Approximate Nearest Neighbor (ANN) 알고리즘은 고차원 데이터에서 가장 가까운 이웃을 빠르게 찾기 위한 중요한 기법입니다. ANN은 추천 시스템, 이미지 검색, 자연어 처리(NLP) 등 다양한 분야에서 사용됩니다. 이번 글에서는 DiskANN 알고리즘에 대해 살펴보고, 이 알고리즘이 대규모 데이터셋에서 효율적인 성능을 제공하는 이유를 설명하겠습니다.

 

왜 DiskANN인가?

대부분의 주류 ANNS 알고리즘은 인덱스 구축 성능, 검색 성능 및 리콜 간에 일부 균형을 이룹니다. HNSW 및 NSG와 같은 그래프 기반 알고리즘은 현재 검색 성능 및 리콜 측면에서 최첨단 방법입니다. 메모리-resident 그래프 기반 인덱싱 방법은 메모리를 너무 많이 차지하므로, 제한된 메모리 리소스를 가진 단일 머신을 사용하여 대규모 데이터 세트를 인덱싱하고 검색하는 것은 상대적으로 어렵습니다.

많은 애플리케이션은 10억 규모의 데이터 세트에서 유클리드 거리 기반 ANNS의 빠른 응답을 요구합니다. 기존에 주요 사용되는 솔루션은 아래 두가지였습니다:

• Inverted index + quantization: 데이터 세트를 M개의 파티션으로 클러스터링하고 PQ(제품 양자화)와 같은 양자화 방식을 사용하여 데이터 세트를 압축합니다. 이 솔루션은 데이터 압축으로 인한 정밀도 손실로 인해 낮은 리콜을 생성합니다. topk를 늘리면 리콜이 향상되지만 QPS는 그에 따라 떨어집니다.
• Divide and index: 데이터 세트를 여러 개의 분리된 샤드로 분할하고 각 샤드에 대한 메모리 내 인덱스를 빌드합니다. 쿼리 요청이 오면 각 샤드의 인덱스에서 검색이 수행되고 병합 후 결과가 반환됩니다. 이러한 솔루션은 데이터 세트 규모를 지나치게 확장하게 되고, 단일 머신의 메모리 리소스가 제한되어 더 많은 머신이 필요하게 되어 QPS가 낮아집니다.

위에서 언급한 두 솔루션은 모두 단일 머신의 메모리 제한으로 제한됩니다. DiskANN은 이 문제를 해결하기 위해 SSD-resident 인덱싱 메커니즘의 설계를 제안합니다. SSD 상주 인덱싱의 챌린지는 무작위 디스크 액세스 수와 디스크 액세스 요청 수를 줄이는 것입니다.

 

DiskANN은 SSD를 활용한 ANNS(Approximate Nearest Neighbor Search) 시스템으로, 대규모 데이터셋에서도 효과적으로 검색을 지원합니다. 본 알고리즘의 주요 기여는 다음과 같습니다:

• 대규모 데이터 처리:DiskANN은 100차원을 초과하는 수십억 개의 데이터셋을 단일 64GB RAM 머신에서 처리할 수 있습니다. 이때 검색 정확도(Recall@1)가 95% 이상이며, 지연 시간은 5밀리초 이내로 유지됩니다.
• Vamana 알고리즘 도입:DiskANN은 NSG나 HNSW보다 작은 검색 반경을 가지는 Vamana라는 새로운 그래프 기반 알고리즘을 제안했습니다. 이는 디스크 접근 횟수를 최소화하여 효율성을 극대화합니다.
• 메모리 내 성능 유지:Vamana는 메모리 내에서도 작동하며, NSG나 HNSW와 비교해도 성능 저하가 없습니다.
• 지속 가능한 그래프 병합:Vamana로 생성된 작은 인덱스 그래프들을 겹치는 데이터셋 파티션으로부터 병합하여 하나의 연결된 그래프로 만들 수 있습니다.
• 양자화와의 결합:Vamana는 PQ(Product Quantization)와 같은 양자화 기법과 결합할 수 있습니다. 이 경우, 그래프 구조와 원본 데이터는 디스크에 저장되고, 압축된 데이터는 메모리에 유지됩니다.

 

---

 

Vamana 알고리즘

Vamana 알고리즘은 NSG(Navigable Small World)와 유사한 아이디어를 기반으로 하지만, 트리밍 전략(trimming strategy)에 차이를 둡니다. NSG는 대상 점(target point)의 이웃 선택이 최대한 다양하도록 설계되었습니다. 예를 들어, 대상 점의 이웃 중 하나가 또 다른 이웃보다 더 가까이 있다면 해당 점은 추가하지 않는 방식으로, 그래프의 out-degree를 효과적으로 제어합니다. 이 방법은 메모리 사용량을 줄이고 검색 속도를 높이는 데 기여하지만, 검색 정확도를 낮추는 단점도 있습니다.

Vamana는 이 트리밍 과정을 보다 유연하게 제어하기 위해 알파(alpha)라는 파라미터를 도입했습니다. 트리밍 조건에서 거리(dist)를 알파(α ≥ 1)로 곱하여 기준 거리를 확장시킵니다. 이로 인해 대상 점 이웃 간 상호 배제 조건이 완화되고, 더 많은 이웃 점을 허용할 수 있게 됩니다.

Vamana의 인덱싱 과정

1. 랜덤 그래프 초기화:

초기에는 무작위로 생성된 그래프를 사용합니다.

2. 시작점 계산:

NSG의 내비게이션 포인트와 유사한 시작점을 계산합니다. 먼저 글로벌 중심을 찾은 다음 글로벌 중심과 가장 가까운 지점을 내비게이션 포인트로 찾습니다. Vamana와 NSG의 차이점은 NSG의 입력이 이미 최근접 이웃 그래프라는 점입니다. 즉, 사용자는 초기 이웃 그래프에서 중심점에 대한 근사적 최근접 이웃 검색을 직접 수행할 수 있습니다. 그러나 Vamana는 임의의 최근접 이웃 그래프를 초기화하므로 사용자는 임의의 그래프에서 직접 근사 검색을 수행할 수 없습니다.

3. 알파 값에 따른 검색 및 트리밍:

2단계에서 결정된 검색 시작점과 초기화된 랜덤 이웃 그래프를 기반으로 각 지점에 대한 근사 최근접 이웃 검색을 수행합니다. 검색 경로의 모든 점을 후보 이웃으로 설정하고 알파(α)를 1로 설정하여 NSG와 유사한 트리밍 과정을 수행합니다. 

4. 반복 및 품질 향상:

첫 번째 반복(iteration) 이후 그래프 품질이 낮기 때문에, 알파 > 1로 조정하여(논문에서는 1.2를 권장합니다) 3단계를 반복합니다. 추가 반복을 통해 그래프 품질을 지속적으로 향상시킵니다. 이는 검색 정확도(recall rate)를 높이는 데 중요한 단계입니다.

 

---

 

DiskANN

메모리가 64GB에 불과한 개인용 컴퓨터는 원시 데이터를 10억 ​​개도 보관할 수 없고, 이를 기반으로 구축된 인덱스는 더더욱 보관할 수 없습니다. DiskANN에서는 이러한 문제를 해결하기 위해 다음과 같은 해결책을 제시합니다.

제한된 메모리 리소스로 이렇게 대규모 데이터 세트를 인덱싱하는 방법은?

먼저 k-means를 사용하여 데이터를 k개의 클러스터로 나눈 다음, 각 포인트를 가장 가까운 i개의 클러스터에 할당합니다. 일반적으로 i개의 클러스터에는 2개면 충분합니다. 각 클러스터에 대한 메모리 기반 Vamana 인덱스를 빌드하고 마지막으로 k개의 Vamana 인덱스를 하나로 병합합니다.

원본 데이터를 메모리에 로드할 수 없는 경우 검색할 때 거리를 계산하는 방법은?

original 벡터에 인덱스를 빌드하고 compressed 벡터를 쿼리합니다. original 벡터에 인덱스를 빌드하면 그래프의 품질이 보장되고, compressed 벡터는 coarse-grained 검색을 위해 메모리에 로드될 수 있습니다. compressed 벡터로 검색하면 정확도가 떨어질 수 있지만, 그래프의 품질이 충분히 높으면 일반적인 방향은 정확합니다. 최종 거리 결과는 original 벡터를 사용하여 계산됩니다.

 

DiskANN의 인덱스 레이아웃은 일반 그래프 인덱스와 유사합니다. 각 포인트의 이웃 집합과 original 벡터 데이터는 함께 저장됩니다. 이렇게 하면 데이터의 지역성을 더 잘 활용할 수 있습니다.

앞서 언급했듯이, 인덱스 데이터가 SSD에 저장되는 경우, 낮은 검색 지연을 보장하기 위해 디스크 액세스 수와 디스크 읽기 및 쓰기 요청을 최대한 줄여야 합니다. 따라서 DiskANN은 두 가지 최적화 전략을 제안합니다.

• Cache hotspot: 메모리의 시작 지점에서 C 점프 내의 모든 지점을 캐시합니다. C의 값은 3~4 이내로 설정하는 것이 좋습니다.
• Beam search: 간단히 말해서, 이웃 정보를 미리 로드하는 것입니다. p 지점을 검색할 때, p의 이웃 지점이 메모리에 없으면 디스크에서 로드해야 합니다. 소량의 SSD 랜덤 액세스 동작은 SSD 단일 섹터 액세스 동작과 거의 같은 시간이 걸리므로, 액세스되지 않은 W 지점의 이웃 정보를 한 번에 로드할 수 있습니다. W는 너무 크거나 작게 설정할 수 없습니다. 큰 W는 컴퓨팅 리소스와 SSD 대역폭을 낭비하는 반면, 작은 W는 검색 지연을 증가시킵니다.

 

성능

Comparison among memory-based indexes: Vamana VS. NSG VS. HNSW

 

Comparison between a one-time built index and a large merged index

 

---

 

 

Milvus

Search parameters

• search_list: 후보 목록의 크기. 크기가 클수록 성능은 떨어지지만 Recall은 높아집니다.

 

DiskANN-related Milvus configurations

...
DiskIndex:
  MaxDegree: 56
  SearchListSize: 100
  PQCodeBugetGBRatio: 0.125
  SearchCacheBudgetGBRatio: 0.125
  BeamWidthRatio: 4.0
...

• MaxDegree: Vamana 그래프의 최대 차수. 값이 클수록 Recall이 높아지지만 인덱스를 구축하는 데 필요한 크기와 시간이 늘어납니다.
• SearchListSize: 후보 목록의 크기. 값이 클수록 인덱스를 구축하는 데 소요되는 시간이 늘어나지만 Recall이 더 높아집니다. 인덱스 구축 시간을 줄여야 하는 경우가 아니면 MaxDegree보다 작은 값으로 설정하세요.
• PQCodeBugetGBRatio: PQ 코드의 크기 제한. 값이 클수록 Recall이 높아지지만 메모리 사용량이 증가합니다.
• SearchCacheBudgetGBRatio: 캐시된 노드 수 대 raw 데이터의 비율. 값이 클수록 메모리 사용량이 증가하고, 인덱스 구축 성능이 향상됩니다.
• BeamWidthRatio: 검색 반복당 최대 IO 요청 수와 CPU 수 간의 비율입니다.