잉여의 생각저장소

참고

Efficient and robust approximate nearest neighbor search using HNSW 논문

C++ HNSW implementation with python bindings

Lucene95HnswVectorsFormat code

배경

 ANN(Approximate Nearest Neighbor)이란 공간에서 query 벡터와 유사한 벡터(문서, 이미지 등이 학습된 데이터)들을 빠르게 검색하는 알고리즘 이다. Approximate라 표현하는 이유는 ANN이 반환한 벡터들이 실제 정답셋이 아니다. 벡터 크기가 보통 1000차원 이상 크고 학습하는 데이터도 규모있는 서비스 경우 수 억건이 넘어가기에 실제 정답을 구하는 비용이 크다. 효율적으로 계산하기 위해 recall(재현율) 및 precision(정밀도)를 기준치 이상 유지하며 빠르게 검색하는 알고리즘들이 개발되고 있다.
(http://ann-benchmarks.com/ 참고)
 

이번 글은 ANN 알고리즘 중 HNSW에 대해 알아보자. HNSW는 본래 c++ 기반으로 위 참고 github에서 먼저 공개되었고 Lucene9x, ES8x 부터 적용되었다. Lucene에서는 native java로 변환하여 적용되었으며 SIMD 등 칩에 최적화하는 과정이 빠져있어 c++보다 성능이 떨어질 수 있다. 참고로 현재 recall 대비 가장 빠른 검색 성능을 내는 알고리즘은 2020년 구글이 공개한 ScaNN 이다.

분석

1. NSW

NSW는 공간상의 벡터들을 배치하고 이들을 적절하게 연결하여 벡터 간 탐색 가능하도록 만드는 알고리즘이다. 크게 벡터를 저장하는 방식(Dense, Sparse ...), 거리 함수, 그리고 벡터를 연결하는 알고리즘(방향성 그래프 기반)으로 구성한다.

그림1 은 편의상 2차원 공간에 4개의 벡터들로 NSW를 구성한 예시이다. M은 이웃 수를 나타내며 보통 10 ~ 20개로 구성한다. 제한된 이웃 수로 효율적인 NSW를 구성하기 위해 아래 두 가지 목표를 설정한다.

1. 최소 방문으로 원하는 목적지 탐색
2. 탐색 불가능한 노드가 없도록

1번은 NSW 내 임의의 두 벡터를 선택했을 때 최단경로 길이의 평균을 최소화 하는 것이다. 2번은 edge(이웃)가 방향성을 가지기 때문에 발생하는 탐색불가능 노드를 최소화 하는 것이다. 예를들어 아래와 같은 상황을 만들면 ANN 성능이 크게 떨어진다.

그림2 는 가장 가까운 벡터들로만 이웃을 구성했을 때 NSW 예시이다. 왼쪽 3개의 벡터들은 서로 연결되어 탐색 가능하지만 오른쪽위 벡터로 가는 길은 없다. HNSW 논문의 Algorithm 3(SELECT-NEIGHBORS-SIMPLE)은 단순 거리계산만 적용된 이웃탐색 알고리즘이며 사용하지 말라고 알려주는 것이다. 이를 대신하여 Algorithm 4(SELECT-NEIGHBORS-HEURISTIC)를 제안한다. 

Algorithm 4는 처음 이웃을 구할 땐 가장 가까운 벡터를 선택하지만 다음 벡터부터 이전에 선택한 방향과 다른 벡터들을 탐색하는 것이 특징이다. 아래 예시를 통해 자세히 알아보자.

그림3 처럼 벡터(노란색, q)가 새로 추가될 때 첫 번째 이웃은 가장 가까운 벡터를 선택한다.

이후 첫 번째 이웃과 직선의 중심에서 수직선을 그어 아래 파란영역에 포함된 벡터들을 모두 후보군에서 제외한다. 만약에 벡터가 3차원 공간이라면 2차원 평면이 기준이 되어 q와 먼 거리의 벡터들을 제외한다. N차원 공간에 대해 N-1차원으로 경계를 만드는 것이다.
(실제 알고리즘과 절차는 약간 다르지만 이 프로세스가 결과를 직관적으로 이해하기 편하다.)

이어서 두 번째 이웃으로 오른쪽위 벡터를 선택하고 마찬가지로 초록색 영역의 벡터들은 후보군에서 제외된다. 이처럼 SELECT-NEIGHBORS-HEURISTIC은 이웃들을 다방면으로 연결하여 그래프 내 강한연결요소(Strongly Connected Component)를 줄인다. (탐색 속도 증가) 또한 군집에서 멀리 떨어진 벡터를 연결하여 탐색 불가능 영역을 최소화 한다. (recall 증가, 논문 Fig. 2. 참고)

2. HNSW

HNSW는 Hierarchical Navigable Small World의 약자로 NSW를 수직 계층적으로 구성하는 것을 뜻한다. 앞서 설명한 NSW를 여러 계층으로 구성하여 탐색 시간을 줄이기 위함이다.

HNSW의 가장 하단 $Lv_0$ 에 실제 찾고자 하는 벡터들을 배치한다. 그리고 $Len(Lv_{n}) = log(Len(Lv_{n-1}))$ 이 되도록 레벨이 높을수록 벡터 수를 log-scale로 줄인다. 이 때 벡터를 선택하는 기준은 랜덤이며 각 $Lv$의 NSW는 독립적으로 동작한다. 최상위 레벨에는 하나의 벡터만 있으며 이 벡터가 시작 $ep$(entry point) 이다.

이 계층을 두는 이유는 검색 시 NSW에서 이웃들을 탐색할 때 최적의 시작위치($ep$)를 구하기 위함이다. HNSW 논문에 Algorithm 1(색인)과 Algorithm 2(검색) 모두 이 계층을 통해 최적화되어 있으며 두 알고리즘은 유사해 검색 예시 하나로 설명한다.

Algorithm 1 (색인)

 

HNSW가 그림5 처럼 색인되어 있고 query 벡터가 들어오면 $Lv_2$의 $ep$ q부터 탐색을 시작한다. $Lv_2$의 $ep$는 탐색할 이웃이 없기 때문에 동일한 $Lv_1$의 $ep$에서 탐색을 진행한다. $Lv_1 ep$의 이웃들을 탐색하면서 query와 떨어진 거리를 기준으로 priorityQueue에 후보 벡터들을 저장한다. 이는 $ef$ 또는 $efConstruction$ 이라 명명한 후보군 제한 수 설정 때문이다. 그림5 예시와 달리 벡터가 많을 경우 검색 시 모든 벡터들을 탐색할 수 있다. 때문에 $ef$ 수를 넘은 만큼 priorityQueue에서 query와 가장 멀리 벡터들을 제거하여 검색 성능을 높인다.

 

$Lv_1 ep$는 그림5-2 처럼 query와 가까운 쪽으로 이동하고 이를 최종 $Lv_0 ep$로 한다. 마지막 $Lv_0$ 검색은 $ep$를 결과 priorityQueue에 저장하여 원하는 결과 벡터 수가 나올 때까지 탐색한다.

다음

Lucene9.x에 적용된 HNSW와 개선점

선행

이번 글은 Lucene 역색인 전체구조에서 FST 구조만 설명합니다.

Lucene의 FST는 Direct Construction of Minimal Acyclic Subsequential Transducers 논문을 기반으로 하며 TestFSTs 로 테스트할 수 있다.

목표

전체 역색인 과정에서 FST의 역할은 Terms의 prefix와 이와 매칭되는 tim의 FP(파일포인터)를 함께 저장할 수 있어야 한다. 또한 term 검색 시 log(BYTE x len(term)) 이내로 빠르게 찾을 수 있어야 하며, Segment 간 병합을 위해 사전순으로 Terms를 iteration 할 수 있어야 한다. 

Lucene에서 FST의 input은 Terms의 공통 prefix를 BYTE 단위로 쪼개 node로 생성하며, output은 .tim의 FP로 long 타입이다. 논문에서 output은 string으로 전제하였다. FST는 output 또한 공통부분을 추출하여 저장하기에 타입에 따라 효율이 크게 달라진다. 예를들어 "100"과 "101"는 공통부분이 "10" 이지만 100과 101은 공통부분이 100이다. 

논문분석

Definition 1.

Subsequential Transducer는 inputs, outputs와 상태 및 함수들로 구성된 tuple로 이 글에선 $\mathbb{T}$ 로 표현한다. Subsequential Transducer는 모든 상태가 결정되었음을 전제로 하며 이 뜻은 inputs 집단과 outputs 집단이 변하지 않음을 뜻한다. Lucene도 주기에 맞춰 Segment 내에서 색인할 때 추출된 Terms로만 FST를 구성할 뿐 기존 FST에 추가할 수 없다. FST에 input을 추가하거나 output을 변경할 때는 처음부터 새로 생성해야 하며 이는 Segment 병합 시 FST를 새로 생성하는 이유기도 하다.

$\mathbb{T}  =  <\Sigma, \Delta, S, s, F, \mu, \lambda, \Psi>$

$\Sigma \ \ is \ a \ finite \ input \ alphabet;$
$\Delta \ \ is \ a \ finite \ output \ alphabet;$
$S \ \ is \ a \ finite \ set \ of \ states;$
$s \in S \ \ is \ the \ starting \ state;$
$F \subseteq S \ \ is \ the \ set \ of \ final \ states;$
$\mu \ : \ S \ \times \ \Sigma \rightarrow S \ \ is \ a \ partial \ function \ called \ the \ transition \ function;$
$\lambda \ : \ S \ \times \ \Sigma \rightarrow \Delta^{*} \ \ is \ a \ partial \ function \ called \ the \ output \ function;$
$\Psi \ : \ F \rightarrow 2^{\Delta^{*}} \ \ is \ the \ final \ function;$

Subsequential Transducer의 구성은 위와 같다. $\Sigma, \Delta, S, F$ 는 집합이며 나머지 함수들은 파라미터와 결과 타입을 의미한다. $\mu, \lambda$ 의 첫 번째 파라미터는 node이며 두 번째 파라미터는 온전한 Term이 아닌 부분으로 쪼갠 Term의 일부다. $\lambda$의 결과 또한 공통 부분만 추출한 output의 일부를 표현한 것이다.

이해를 돕기위해 월별 데이터로 FST를 구성하면 위 과정을 거친다. 입력은 월별 영어 약자($\Sigma$) 이며 출력은 월별 말일($\Delta$) 이다. ○(compiled, 확정된), □(compiling, 계산중) 들은 node를 의미하며 $S$에 속한다. 각 함수의 결과 값은 node 및 arc에 저장되며 아래 예시처럼 동작한다. (Lucene FST 코드에서 정점을 node, 간선을 arc로 명명하여 이 글에서도 똑같이 표현한다.)

$\Sigma$ = {['a','p','r'], ['a','u','g'], ['d','e','c'], ['f','e','b'], ... }
$\Delta$ = {"30", "31", "31", ["28", "29"], ...}
$\mu$(s5, 'p') $\rightarrow$ s2
$\lambda$(s5, 'p') $\rightarrow$ "0"
$\Psi$(s1) $\rightarrow$ ""
$\Psi$(s8) $\rightarrow$ ["8", "9"]

$\mu, \lambda$ 의 확장함수($^*$) 들은 두 번째 파라미터로 하나의 문자가 아닌 prefix 문자열($\sigma$)을 받으며 아래 규칙을 따른다.

${\forall}r \in S, \forall\sigma \in \Sigma^*, {\forall}a \in \Sigma$ ($\Sigma^*$ 은 모든 prefix inputs 집합)
$\mu^*(r, {\sigma}a) = \mu(\mu^*(r, \sigma), a)$
$\lambda^*(r, {\sigma}a) = \lambda^*(r, \sigma)\lambda(\mu^*(r, \sigma), a)$

e.g.
$\mu^*$(t0, "apr") = $\mu$($\mu^*$(t0, "ap"), 'r') = $\mu$(s2, 'r') = s1
$\lambda^*$(t0, "apr") = $\lambda^*$(t0, "ap")$\lambda$($\mu^*$(t0, "ap"), 'r') = "30"$\lambda$(s2, 'r') = "30"

최종으로 input language 함수 $L$과 output 함수 $O_{\mathbb{T}}$ 정의는 아래와 같다.

$L(\mathbb{T}) \ = \{ \sigma \in \Sigma^* \ | \ \mu^*(s, \sigma) \in F \}$
$O_\mathbb{T}(\sigma) \ = \lambda^*(s, \sigma) \cdot \Psi(\mu^*(s, \sigma))$

e.g.
$L(\mathbb{T})$ = {"apr", "aug", "dec", "feb", ...}
$O_\mathbb{T}$("apr") = "30"

임의의 두 Transducers가 $L(\mathbb{T}), O_{\mathbb{T}}$ 의 입출력 값들이 모두 같다면 두 Transducers는 같다고(equivalent) 한다. 

Definition 2, 3.

정의 2, 3은 outputs 측면에서 Minimal Subsequential Transducer를 구성하기 위한 정리들을 설명한다. Minimal Subsequential Transducer란 Equivalent Transducers 중에 가장 node, arc 수가 작으면서 outputs 저장공간이 가장 작은 Transducer를 뜻한다.

정의2 는 $L(\mathbb{T})$ 집합의 prefix 집합 $D(\mathbb{T})$에 대하여 공통의 output을 가장 앞선 node에 배치하기 위한 $g_{\mathbb{T}}(u)$ 함수를 선언한다.

$D(\mathbb{T}) \ = \{ u \in \Sigma^* \ | \ {\exists}w \in \Sigma^* \ (uw \in L(\mathbb{T}) \ \}$
$g_{\mathbb{T}}(u) \ = \ \wedge_{w \in \Sigma^* \ \& \ uw \in L(\mathbb{T})} {\wedge}O_{\mathbb{T}}(uw) $ 
($\exists$는 존재함을 $\wedge$는 and 연산을 의미)

$D(\mathbb{T})$ 는 $L(\mathbb{T})$ 로 만들 수 있는 모든 prefix 집합.
$g_{\mathbb{T}}(u)$ 는 u로 시작하는 모든 input들의 output들을 모아 공통 부분을 추출한 것. 

e.g.
$u$ = "j" 일 때 $uw$ 집합은 {"jan", "jul", "jun"} 이며 각 output은 아래와 같음.

$O_{\mathbb{T}}$("jan") = "31"
$O_{\mathbb{T}}$("jul") = "31"
$O_{\mathbb{T}}$("jun") = "30"

i.g.
$g_{\mathbb{T}}$("j") = $O_{\mathbb{T}}$("jan") $\wedge$ $O_{\mathbb{T}}$("jul") $\wedge$ $O_{\mathbb{T}}$("jun") = "31" $\wedge$ "31" $\wedge$ "30" = "3"

정의3 은 $g_{\mathbb{T}}$ 함수를 활용하여 transducer의 총 $\lambda$ 합을 최소화 하기 위한 함수를 선언한다. 그리고 아래 조건을 만족하는 transducer를 Canonical Subsequential Transducer 라고 명명한다.

${\forall}r \in S, \forall\sigma \in \Sigma^*, {\forall}a \in \Sigma$
$(\mu^*(s, \sigma) = r \ \& \ !\mu(r, a)) \ \rightarrow \ \lambda(r, a) = [g_{\mathbb{T}}(\sigma)]^{-1}g_{\mathbb{T}}({\sigma}a)$
($!\mu$ 는 node가 존재하지 않을 때)

e.g.
$g_{\mathbb{T}}$("a") = "3"
$g_{\mathbb{T}}$("ap") = "30"
$\lambda$(s5, 'p') =  $g_{\mathbb{T}}$("a")$^{-1}g_{\mathbb{T}}$("au") = "0"

Canonical Subsequential Transducer는 논문에 언급한 것처럼 가장 앞선 node에 output을 최대한 많이 저장한다. 그림1 에서 ("jan", "31"), ("jul", "31") 을 저장하는 과정을 다시보자. $\lambda$(t0, 'j')가 빈 값이라면 $\lambda$(t1, 'a'), $\lambda$(t1, 'u') 두 arcs에 "31"을 저장하므로 중복이 발생한다. 그러므로 가장 앞선 node인 $\lambda$(t0, 'j')에 "31"을 저장하는 것이 공간을 최소화하는 방법이다. 

이후 ("jun", "30")을 저장할 때는 $g_{\mathbb{T}}$ 에 따라 $\lambda$(t0, 'j') = "3", $\lambda$(t1, 'a') = "1", $\lambda$(t2, 'l') = "1", $\lambda$(t2, 'n') = "0" 으로 바뀐다. 즉, 사전순으로 Term을 입력하면서 가장 앞선 node에 공통 output을 저장하여 공간을 최소화 한다. 자세한 증명은 인용 논문인 Minimization algotithms for sequential transducers에서 다루며 node, arc 빌드 과정은 다음 정리 이후 설명한다.

여기까지 Canonical Subsequential Transducer를 구성하기 위한 정의이다. 그리고 논문에서는 Canonical Subsequential Transducer 내에 equivalent node가 없다면 minimal Transducer라고 한다. 쉽게 설명하면 outputs까지 최적화 하였으니 중복된 node, arc를 제거하여 최소화하라는 것이다. 그림1 s1 node가 하나의 예시이며 특정시점 이후 node, arc가 모두 같을 때 해당 node를 재사용하는 과정을 다음 정리에서 설명한다.

Definition 4.

정의 4는 $\mathbb{T}  =  <\Sigma, \Delta, S, s, F, \mu, \lambda, \Psi>$ 가 아래 조건들을 만족할 시 마지막 Term의 prefix인 $w$ 를 제외하고 minimal이라고 한다. 다시 그림1 예시로 조건들을 설명하면 다음과 같다.

  1.  모든 node가 시작 node로 부터 접근 가능해야 한다.

  2. 사전순으로 마지막 단어인 "jul"의 prefix인 $w$의 node들은 아직 $\mathbb{T}$에 포함되지 않은 빌드 중인 상태(□) 이다.

      그리고 아래 정의와 조건을 만족한다.

$w = w_{1}^{\mathbb{T}}w_{2}^{\mathbb{T}}...w_{k}^{\mathbb{T}}, \ w_{i}^{\mathbb{T}} \in \Sigma, \ i \in 1...k$
$t_{0}^{\mathbb{T}} = s; \ t_{1}^{\mathbb{T}} = \mu(t_{0}^{\mathbb{T}}, w_{1}^{\mathbb{T}}); \ ... \ ; \  t_{k}^{\mathbb{T}} = \mu(t_{k-1}^{\mathbb{T}}, w_{k}^{\mathbb{T}})$ 
$T = \{ t_{0}^{\mathbb{T}}, t_{1}^{\mathbb{T}}, ..., t_{k}^{\mathbb{T}}\}$

$\forall{r} \ \in \ S, \forall{i} \ \in \ \{ 1 ... k\}, \forall{a} \ \in \ \Sigma$
$\mu(r,a) = t_i \leftrightarrow (i > 0 \ \& \ r = t_{i-1} \ \& \  a = w_i^{\mathbb{T}})$

   3. $S \setminus T$에 equivalent states 들이 없다. (즉, 모든 node, arc가 unique 함)

   4. $\mathbb{T}$는 Canonical Subsequential Transducer다.

정의4는 부분으로 구성된 Minimal Subsequential Transducer(이하 MST)에 새로운 node, arc를 추가하면서 minimal 상태를 유지하기 위함이다. 이를 위해 MST에 포함된 node(해당 node에서 출발하는 arcs 포함)를 ○로 표시하고 아직 계산 중인 $T$의 node들을 □로 표시한다. $T$ 에서 MST에 node를 추가하는 과정은 가장 끝의 node부터 시작한다. 그림1 예시로 "jul"에서 'l'에 해당하는 node 부터 minimal 상태를 유지한채 transducer에 입력하는 것이다. 이는 아래 보조정리들로 자세히 설명한다.

Lemma1은 $t_{k}$와 동일한 기존 node가 없을 때 이를 $\mathbb{T}$에 추가해도 minimal 이다.

Lemma2는 $t_{k}$와 동일한 기존 node인 p가 있을 때 $t_{k-1}$에서 $w_{k}$ arc의 target을 p로 가리켜도 minimal 이다.

Lemma3는 $t_{k}$와 동일한 기존 node의 조건을 말한다. 요약하면 동일한 노드는 final일 때 출력이 같아야하며 final이 아니면 가리키는 모든 arc 정보가 같아야한다.

다시 그림1의 마지막 사전순 단어 "jul" 에서 마지막 단어부터 $\mathbb{T}$에 추가한다. 이때 다음 단어인 "jun"과 common prefix인 "ju" 만 $T$로 남기고 'l' arc가 추가된다. $g_{\mathbb{T}}$("ju") = "3" 이므로 'l' arc의 output은 "1"이며 $w$ = "jun" 으로 바뀐다. 결과는 아래 그림2와 같다.

Lemma2에 의해 <$\mu$(t2,'l') = s1, $\lambda$(t2,'l') = "1">로 추가해도 $\mathbb{T}$는 minimal이다. 이후 ("MAR", "31)을 색인할 때 common prefix가 하나도 없으므로 t1, t2, t3를 $\mathbb{T}$에 추가해야 한다. Lemma3에 의해 $\mu$(t2,'n')는 s1을 가리키고 t1과 t2는 새로운 node로 추가된다. t1과 t2는 동일한 node가 없어서 추가된 것이므로 Lemma1에 의해 $\mathbb{T}$는 minimal이다.

종합하여 Minimal Subsequential Transducer를 생성하는 과정은 뜨개질과 유사하다. 구성할 모든 단어들을 사전 순으로 탐색하며 마지막 단어의 끝부터 공통의 input, output만 남기고 나머지 부분을 확정짓는 방식이다. 확정된 node들은 이후 수정할 수 없기에 뜨개질 중간에 튀어나온 올을 없애기 힘든 것과 같다. 한 번의 뜨개질로 minimal 상태를 만들어야 하므로 정렬되고 결정된(finite) inputs를 전제로 한다.

다음

이번 글에서 Lucene FST를 이해하기 위한 정의들을 다뤘지만 Theorem3, 4 및 Algorithm이 빠졌습니다. Theorem3, 4는 $T$에서 common prefix만 남기고 나머지 node, arc를 추가하는 상세한 과정과 증명인데 너무 길어서 위에 간략하게만 설명했습니다. 사실 여기까지도 다 읽으실 분들이 많지 않을 것 같아서 요청하시면 추가할게요!

Algorithm은 Lucene의 FST 기반으로 설명드리려 합니다. Lucene의 FST는 파일로 en/decoding 하기에 arc가 많을 시 binary search 할 수 있도록 codec을 구성합니다.