안녕하세요!
오늘은 혁신적인 검색 기법을 제안하고 있는 MUVERA 라는 논문을 깊게 공부해보는 시간을 갖겠습니다.
오늘 살펴볼 MUVERA는 구글 리서치에서 작년 5월에 발표한 논문으로, 다중 벡터 유사도 검색을 단일 벡터 검색으로 변환하는 혁신적인 다중 벡터 검색 알고리즘을 제안합니다.
Introduction
먼저 임베딩 모델 기반 검색 시스템의 발전에 대해 살펴보도록 하겠습니다.
임베딩 모델과 단일 벡터 모델
딥러닝 기반 임베딩 모델은 방대한 데이터셋에서 사용자 쿼리에 대한 연관 정보를 빠르게 찾기 위한 핵심적인 도구로 사용되고 있는데요.
이에 따라 단일 벡터 모델이 등장하게 됩니다.
단일 벡터 모델이란 문서나 질의당 하나의 임베딩 벡터를 생성하고, 내적 기반 유사도 계산(MIPS, Maximum Inner Product Search을 적용하는, 우리가 흔히 아는 검색 기법입니다. 이 모델은 내적 기반 연산을 사용하기에 빠른 검색 성능을 제공하지만, 하나의 임베딩 벡터로는 문서의 복잡한 의미를 모두 표현하기에는 한계가 있었습니다.
다중 벡터 모델의 도입
이러한 문제점을 해결하기 위해 등장한 것이 다중 벡터 모델입니다.
다중 벡터 모델은 문서나 질의를 토큰별로 분할하여 각각에 대해 임베딩 벡터를 생성하고, 내적이 아닌 Chamfer 유사도와 같은 복잡한 유사도 함수를 적용합니다. 즉, 문서나 질의당 다수의 임베딩 벡터들을 생성하게 되는 것입니다.
토큰별 임베딩을 통해 단일 임베딩보다 문서의 세밀한 의미와 복잡한 관계 정보를 더 잘 포착할 수 있으며, 기존 단일 벡터 모델보다 높은 검색 정확도를 달성합니다.
다중 벡터 모델의 한계
하지만 이러한 다중 벡터 모델도 한계를 가지고 있긴 했는데요.
- 일단 첫번째로는 토큰별로 임베딩 벡터를 생성해야 하기에 임베딩 벡터 수가 증가하면서 연산량과 메모리, 연산 시간이 대폭 증가하게 됩니다. 쿼리 단일 벡터와 문서 단일 벡터끼리 1회 시행했던 내적 연산을 각 쿼리 벡터마다 모든 문서 벡터와 내적을 하게 되면서, O(1)이 O(|Q| x |P|, Q: 쿼리의 토큰 개수, P: 문서의 토큰 개수)가 되게 됩니다.
- 두번째로는 Chamfer 유사도 계산의 문제입니다.
Chamfer 유사도 수식을 단계별로 살펴보겠습니다.
아래의 수식을 보시면 Q는 쿼리를, P는 문서를 의미합니다.
1. 쿼리를 각 임베딩 벡터에 대해 문서의 임베딩 벡터들과의 내적 계산을 하고,
2. 그 내적 값들 중 최댓값을 선택한 뒤
3. 모든 쿼리 임베딩 벡터들에 대해 이를 반복하여 합산하여 쿼리와 문서 간의 전체 유사도를 계산하게 됩니다.
여기서 문제는 max 연산의 비선형성입니다.
max 연산이 들어가면, 쿼리의 모든 임베딩 벡터마다 문서의 모든 임베딩 벡터와의 내적값을 다 비교해야 하며, O(|Q|x|P|)의 복잡도의 높은 연산 비용을 발생시킵니다.
max 연산으로 인해 Chamfer 유사도 연산을 단일 벡터 간의 내적처럼 선형화된 계산으로 변환이 불가능하게 됩니다.
MUVERA의 핵심 아이디어
따라서 MUVERA는 이러한 다중 벡터 모델의 비용 문제를 해결하고자 쿼리/문서당 생성되는 다수의 임베딩 벡터들을 고정 차원의 단일 벡터, 즉 Fixed Dimensional Encoding, FDE로 변환하는 기법을 제안합니다.
만약 Chamfer 유사도에서 q라는 쿼리 임베딩 벡터에 대해 어떤 문서 임베딩 벡터가 최대 내적 값을 가져오는지, 즉 최적일 지를 모든 임베딩 벡터들을 탐색하지 않고 바로 알 수 있다면, 아래의 수식과 같이 Chamfer 유사도 검색을 max 연산없이 기존의 MIPS(내적 기반 유사도 검색)으로 축소할 수 있습니다.
MUVERA의 전체 절차는 다음과 같습니다.
1. 먼저 문서의 다중 벡터 집합을 고정 길이의 단일 벡터 FDE로 변환하고,
2. 변환된 문서 FDE를 MIPS(내적 기반 유사도 검색) 인덱스에 저장합니다.
쿼리가 입력되고 쿼리 FDE가 계산되면 MIPS solver는 유사한 문서 FDE 후보들을 초고속으로 검색합니다.
3. 검색된 후보군들에 대해 Chamfer유사도를 사용해서 정밀하게 재정렬합니다.
Methods
이제 MUVERA에서 다중 벡터 집합을 고정 길이의 단일 벡터로 변환하는 과정에 대해 단계별로 살펴보겠습니다.
이 부분이 논문의 핵심 부분이라 "내적 연산 반복없이 최적의 문서 임베딩 벡터 구하는 법"에 집중해서 읽으시면 좋으실 것 같습니다.
FDE 변환 과정 (1): 공간 분할
먼저 첫번째 단계는 공간 분할입니다.
임베딩 벡터 차원의 임베딩 공간을 특정 개수의 클러스터로 분할하고 문서, 쿼리 임베딩 벡터들을 클러스터에 매핑합니다.
이 때 논문에서는 서로 가까운 임베딩 벡터들이 동일한 클러스터에 매핑되도록 분할하는 데에 집중하였는데 주로 SimHash 기반의 공간 분할 방식을 사용하였습니다.
SimHash 공간 분할 방식에 대해 간략하게 설명해보자면,
1)먼저 k_sim개의 임베딩 벡터와 같은 차원을 갖는 랜덤 가우시안 벡터들(g1, ..., gksim)을 샘플링합니다.
2) 어떤 벡터 x에 대해 가우시안 벡터들과의 내적을 계산합니다.
3) 그 내적 값이 0보다 큰지 작은지 여부를 나타내는 비트 문자열로 벡터 x의 클러스터를 정의합니다. (클러스터 개수 B=2^k_sim)
FDE 변환 과정 (2): 클러스터 단위의 FDE 생성
문서, 쿼리 임베딩 벡터들을 클러스터에 매핑하면 이제 쿼리 임베딩과 문서 임베딩이 같은 클러스터에 위치하게 되는 경우가 발생합니다. 이 때 어떤 쿼리 임베딩에 대해 같은 클러스터에 속한 문서 임베딩 벡터들의 평균으로 가장 유사한 문서 임베딩 벡터를 근사할 수 있다는 것이 FDE 생성의 핵심 아이디어입니다.
아래의 예시를 보시면 클러스터 C에 쿼리 임베딩 w가 문서 임베딩 y와 z와 같은 클러스터로 묶여 있는데, 이 상황에서 쿼리 임베딩 w에게 가장 유사한 문서 임베딩은 y와 z의 평균이라고 근사할 수 있게 되는 것입니다.
이러한 아이디어를 기반으로 쿼리 Q와 문서 P의 Chamfer 유사도를 단일 벡터간의 내적으로 근사할 수 있게 됩니다.
원래는 n개의 쿼리 임베딩 벡터들과 m개의 문서 임베딩 벡터 간의 내적을 n*m회만큼 실행한 후, 각 쿼리 임베딩 벡터들에 대해 최대 내적값을 합해서 구합니다.
하지만 n*m회만큼의 내적을 실행하지 않고도 최대 내적값을 구할 수 있는 문서 임베딩 벡터를 클러스터링으로 바로 알 수 있다면 concatenate한 쿼리 임베딩 벡터와 concatenate한 문서 임베딩 벡터의 내적으로 Chamfer 유사도를 근사할 수 있다는 것입니다.
어려운 개념일수도 있으니 후에 수식과 함께 설명하겠습니다.
문서 FDE 만들기
지금까지의 아이디어를 가지고 문서의 멀티 임베딩 벡터를 FDE로 변환하는 과정을 아래 영상과 함께 살펴봅시다.
먼저 문서 임베딩 벡터들을 클러스터에 매핑하고 각 클러스터 k에 속하는 문서 임베딩 벡터들을 평균하여 p(k) 블록을 생성합니다.
더 자세하게 설명해보자면,
1. 각 클러스터 k에 속하는 문서 임베딩 p들을 평균하여 p(k) 블록을 생성합니다.
영상으로 예시를 들어보자면 클러스터 3에 속하는 문서 임베딩 a와 b를 평균하여 p(3) = (a+b)/2를 생성합니다.
2. 모든 클러스터에 대해 블록을 만든 후, 이를 모두 연결(concatenate)하여 문서 FDE인 FDE(P)를 생성합니다.
쿼리 FDE 만들기
쿼리도 문서 FDE와 비슷합니다.
아래 영상을 보시면 쿼리 임베딩 벡터들을 클러스터에 매핑하고 각 클러스터 k에 속하는 쿼리 임베딩 벡터들을 합산하여 q(k) 블록을 생성합니다.
더 자세하게 설명해보자면,
1. 각 클러스터 k에 속하는 쿼리 임베딩 q들을 합하여 q(k) 블록을 생성합니다.
2. 모든 클러스터에 대해 블록을 만든 후, 모두 연결(concatenate)하여 문서 FDE(Q)를 생성합니다.
여기서 의문이 생길 수 있습니다. 왜 쿼리는 합산인데 문서는 평균일까?
Chamfer 유사도 수식으로 다시 돌아가보면 해답을 찾을 수 있습니다.
전에 말했다시피 Chamfer 유사도는 쿼리 임베딩 벡터를 기준으로 문서를 스코어링하는 형태입니다.
하지만 앞선 내용과 같이 MUVERA는 공간 분할을 도입하여 Chamfer 유사도 계산 시 "가장 유사한 벡터"를 찾는 연산을 "가장 유사한 클러스터의 대표 벡터(평균 벡터)"를 사용하는 근사 연산으로 대체합니다.
이 때 연산은 쿼리 임베딩 벡터 단위의 합산에서 클러스터 단위의 합산으로 변환됩니다.
따라서 내적의 선형성(<q1, p(k)> + <q2, p(k)> = <q1+q2, p(k)>) 에 따라 클러스터 k에 속한 각 쿼리 벡터 q1, q2가 동일한 p(k)와 내적되는 경우, 각 쿼리 벡터와 p(k)와의 내적을 쿼리 벡터의 합산과 p(k)의 내적으로 단순화할 수 있습니다.
이렇게 Chamfer 유사도 계산이 어떤 클러스터 k에 대해 쿼리 블록 q(k)와 문서 p(k)와의 내적을 하고 이것을 모든 클러스터에 대해 반복하는 연산으로 변환된다면 max 연산이 사라져 모든 클러스터에 대해 concatenate한 q(k)와 p(k)의 내적, 즉 단일 벡터간의 내적으로 추가적으로 변환할 수 있게 됩니다.
따라서 O(|Q|x|P|)의 복잡도를 O(1)로 바꿀 수 있게 됩니다.
FDE 변환 과정(3): 비어있는 클러스터 채우기
사실 concatenate하기 전에 좀 더 해야 할 단계들이 있습니다.
어떤 쿼리 임베딩 이 속해있는 클러스터에 문서 임베딩이 없는 경우, 아래 이미지처럼 q(k)가 0이 되고 해당 쿼리 토큰이 유사도 측정에서 무시되는 문제가 발생합니다.
따라서 만약 어떤 클러스터 k에 문서 임베딩 벡터가 없다면, 클러스터 k에서 가장 가까운 클러스터 k'의 문서 임베딩 벡터들의 평균으로 대체합니다.
여기서 "가장 가깝다"는 SimHash를 사용했을 때 클러스터 k의 비트 문자열에서 가장 적은 수의 불일치 비트를 가짐을 의미합니다.
FDE 변환 과정(4): 랜덤 선형 투영
임베딩 차원 dim에 대한 의존성을 줄이기 위해 각 클러스터 k에 대해 생성된 q(k), p(k) 블록에 개별적으로 random linear projection을 수행합니다.
FDE 변환 과정(5): 1~4단계 반복
공간 분할부터 선형 투영까지 전체 과정을 R_reps회 반복하여 여러 개의 FDE들을 얻고 이들을 concatenate하여 최종 FDE를 구성합니다.
따라서 최종 FDE의 차원은 클러스터의 개수(B) * linear projection 차원(d_proj) * 위의 것들을 반복한 횟수(R_reps)가 됩니다.
반복을 통해 다양한 랜덤 분할 및 투영의 장점을 결함하여 근사 정확도를 향상시킵니다.
이론적 보장
본 논문에서는 FDE 기법이 Chamfer 유사도를 정확하게 근사함을 이론적으로 증명하였습니다.
- 먼저 다중 벡터 집합을 압축시킨 FDE간의 내적만으로 복잡한 Chamfer 유사도를 앱실론 오차 이내로 근사함을 증명하였습니다.
쿼리 Q와 문서 P의 정규화된 Chamfer 유사도(NCHAMFER(Q, P))와 쿼리 및 문서의 FDE 내적을 쿼리 수로 나눈 값(1/|Q|<Fq(Q), Fdoc(P)>) 사이의 차이가 ε보다 작다는 것을 의미합니다.
- 또한 MUVERA로 찾은 최적 문서와 쿼리 간의 Chamfer 유사도와 실제 최적 문서와 쿼리 간의 Chamfer 유사도가 앱실론 오차 이내로 근사함을 증명하였습니다.
이는 MUVERA가 실제 가장 유사한 문서를 정확히 찾아내지 못하더라도, 그 유사도 값이 실제 최댓값과 거의 차이가 없는 문서를 찾아낼 수 있음을 의미합니다.
Evaluation
논문에서는 MUVERA에 대한 평가를 두 가지 방식으로 진행합니다.
- 첫번째 방식은 Offline 평가로 FDE의 품질을 평가하는 방식이고,
- 두번째 방식은 Online 평가로 실제 검색 시스템에서의 end-to-end 성능을 측정하는 방식입니다.
평가 지표의 경우,
- Recall@N: FDE 내적 기반 top-N 후보군 안에 정답 문서가 포함된 쿼리 수 / 전체 쿼리 수, 검색 정확도를 의미
- Latency: 평균 응답 시간
- QPS: 초당 처리 가능한 쿼리 수
데이터셋은 BEIR 벤치마크의 6개의 지식 검색 데이터셋을 사용합니다.
임베딩 모델은 토큰 단위로 임베딩을 생성하는 ColBERTv2(임베딩 차원:128)을 사용하며, 쿼리당 임베딩 개수는 32개로 문서당 임베딩 개수는 가변적으로 처리하도록 하였습니다.
Offline Evaluation: FDE Quality vs Dimensionality
본 논문에서는 FDE를 구성하는 다양한 파라미터들이 성능에 미치는 영향에 대해 실험하였습니다.
당연하게도 아래 Figure를 보면 FDE의 차원 (= 클러스터의 개수(B) * linear projection 차원(d_proj) * 위의 것들을 반복한 횟수(R_reps))가 증가할수록 더 많은 정보를 담을 수 있기에 FDE의 검색 품질(Recall@N)이 향상되는 것을 확인할 수 있습니다.
FDE 차원을 구성하는 파라미터들 중에서도 반복 횟수인 R_reps가 가장 큰 영향을 미치는 파라미터이며 나머지(클러스터 개수, 선형 투영 차원)은 그에 비해 영향이 미미합니다.
Offline Evaluation: Single Vector Heuristic vs FDE retrieval
다음으로는 기존 다중 멀티 검색 시스템의 주류 방식인 Single Vector Heuristic과 MUVERA의 FDE 검색을 비교합니다.
그러기 위해서 Single Vector Heuristic에 대해 대략 알 필요가 있습니다. Single Vector Heuristic은 Chamfer 유사도의 전체 문서에 대한 직접적인 계산을 피하고 효율적으로 후보군을 추려내어 이 후보군에 대해서만 Chamfer 유사도를 계산하는 방식입니다.
먼저, 스캔되는 부동 소수점을 기준으로 비용을 비교합니다.
- Single Vector Heuristic의 경우, 쿼리 토큰 개수 * 문서 토큰 개수 만큼의 내적 연산이 필요하기에 어마어마한 수의 부동 소수점을 스캔해야 합니다.
실험 환경(임베딩 벡터 차원:128, 쿼리당 임베딩 개수:32)과 MS MARCO 데이터셋(문서당 평균 임베딩 벡터 개수mavg=78.8)를 예시로 들자면, 128차원의 임베딩 차원에 대해 32개의 쿼리 임베딩 벡터의 브루트 포스 스캔을 n(문서의 총 개수) * mavg(문서당 평균 임베딩 벡터 개수)개의 문서 임베딩 벡터에 대해 실행해야 하므로, 스캔된 부동 소수점 수는 128 * 32 * n * 78.8, 즉 322,764n이 됩니다. - FDE 검색의 경우, 다수의 멀티 임베딩 벡터를 단일 고정 차원 벡터로 변환했기에 단일 쿼리 FDE 벡터와 단일 문서 FDE 벡터간의 내적만을 수행하여 dFD를 10,240이라고 가정하면 싱글 벡터 휴리스틱에 대해 31배 적은 부동 소수점 수를 스캔하게 됩니다.
실험 환경이나 데이터셋 상관없이 FDE 검색 적용 시 n(문서 개수) * dFDE(FDE의 차원) 만큼의 부동 소수점 수를 스캔합니다.
이는 FDE 검색이 계산 복잡도의 근복적인 개선을 이뤘음을 의미합니다.
다음으로는 검색하는 후보군 수를 변경해가면서 검색 정확도(Recall@N)를 비교합니다.
아래 테이블을 보시면
- 효율성 측면에서는 FDE 검색이 중복 제거된 Single Vector Heuristic 대비 2~5배 더 적은 후보 문서를 검색해도 동일한 Recall@N을 달성하였음을 확인할 수 있습니다.
- 정확도 측면에서는 동일한 수의 후보 문서를 사용할 경우에도 최대 56%, 평균 10% 더 높은 Recall@N 성능을 달성하였음을 확인할 수 있습니다.
이는 FDE 검색이 기존 방법 대비 정확도도 개선하였음을 의미합니다.
Online Evaluation: PLAID vs MUVERA
다음으로는 MUVERA를 활용한 실제 검색 시스템의 성능을 평가합니다.
단일 벡터 MIPS 검색을 위해 DiskANN이라는 최첨단 그래프 기반 이웃 검색(ANNS) 알고리즘을 사용합니다.
압축되지 않은 문서 FDE를 사용해서 DiskANN 인덱스를 구축하고, 쿼리 입력 시에 DiskANN 인덱스에 Beam Search를 사용하여 후보들을 검색하고, 검색된 후보군을 Chamfer 유사도로 reranking하는 방식으로 검색을 수행하게 됩니다.
논문에서는 검색 성능을 향상시키기 위해 Ball-Carving과 Product Quanization 방식을 적용하여 추가적인 최적화를 적용하였습니다.
- Ball-Carving: 다수의 쿼리 임베딩을 클러스터링하여 여러 벡터를 하나의 대표 벡터(합산)로 압축하는 방식으로 쿼리 임베딩 집합의 중복성을 제거하여 reranking 속도를 향상시킵니다.
- Product Quantization: 고차원 벡터를 G개의 서브 벡터로 나누고, 각 서브 벡터마다 k-means로 C개의 클러스터 중심을 만들어서 각 서브 벡터를 가장 가까운 클러스터 중심 인덱스로 대체하는 방식으로 문서 FDE를 더 작은 코드북으로 압축하여 메모리 사용량 감소 및 검색 속도를 향상시킵니다.
이렇게 구축한 검색 시스템을 바탕으로 기존 SOTA인 PLAID와 종합적인 성능 비교를 수행하였습니다.
여기서 PLAID는 기존 다중 벡터 검색 시스템으로, ColBERT의 높은 정확도를 유지하면서 latency를 줄이기 위해 제안된 토큰 수준의 multi-vector 검색 방식을 말합니다.
MUVERA는 평가된 모든 데이터셋에 대해 기존 SOTA 대비 평균적으로 90% 더 낮은 latency와 10% 더 높은 recall을 달성하였습니다.
Conclusion
MUVERA가 활용되었던 사례로 morphik에서의 활용이 있는데요.
morphik은 Colpali 모델 기반 시각 문서 검색 플랫폼입니다. PDF를 이미지로 인식하고 패치 단위로 분할하여 시각+텍스트 정보 기반 다중 벡터 임베딩을 생성하여 검색에 활용했었다고 합니다. 이 때 Colpali가 생성한 고차원 다중 벡터를 MUVERA 기반 FDE로 변환하여 빠르고 정확한 단일 벡터 검색을 가능케 하였습니다.
MUVERA는 다중 벡터 검색의 성능과 단일 벡터 검색의 효율성을 동시에 달성하는, 다중 벡터 검색 문제를 근본적으로 해결하는 새로운 패러다임을 제시하였습니다. 다만, 실제 검색 시스템에 적용하기 위해서는 추가적인 실험이 필요하다고 생각합니다. 따라서 본 논문 리뷰 이후 MUVERA에 대한 추가 실험을 진행하여 블로그에 포스팅할 예정이니 많은 관심 부탁드립니다. ㅎㅎ