PStage MRC 6강 - Scaling up with FAISS

Passage Retrieval and Similarity Search

저번 강의에서 임베딩을 통해 리트리벌을 진행하려고 할 때 질문과 지문 쪽 각각 인코더가 존재했다는 것을 배웠다. 질문은 그때그때 임베딩을 하며, 지문은 미리 임베딩을 해놓지만 새로운 것이 들어올 때 추가하여 적절한 임베딩을 반환하게 된다.

Nearest Neighbor Search에서는 Passage 개수가 늘어날 수록 Top-k 때 Dot product 연산이 부담스러워질 수 있는데, 이러한 과정에서 Similarity Search를 아는 것이 중요하다.

MIPS(Maximum Inner Product Search)

기본적으로 Nearest Neighbor Search보다, Inner product Search가 더 많이 사용되고 있다. Nearest Neighbor와 같은 L2 유클리디언 거리를 계산하는 것보다 Dot Product의 Maximum을 찾는 문제로 돌리는 것이 더 수월하기 때문이다. 가장 가까운 벡터를 찾겠다는 것은 Maximum Inner Product를 찾겠다는 것으로 이해하면 된다. 하지만 개념 정리나 상상을 할 때는 Nearest Neighbor Search를 떠올리는 것이 더 쉽긴 한다.

MIPS는 주어진 질문(query) 벡터 q에 대해 Passage 벡터 v들 중 가장 질문과 관련된 벡터를 찾는 문제이며, 관련성은 내적이 가장 큰 값으로 찾는다.

하지만 실제 검색해야할 데이터는 위키피디아에만 5백만개이며, 수십억, 수십조까지 커질 수 있기 때문에 사실상 모든 문서 임베딩을 Bruteforce로 찾는 것은 불가능하다.

Tradeoffs of similarity search

1) Search Speed - 쿼리당 유사한 벡터 k개를 찾는 시간. 많은 벡터를 가지고 있을수록 당연히 시간이 오래걸린다.

Pruning 을 사용하면 속도 개선 가능

2) Memory Usage - 벡터를 어디에서 가져올 것인지. RAM에 올려놓으면 가장 빠르지만 RAM은 비싸다.

Compression 으로 메모리 압축 가능

3) Accuracy - Bruteforce 검색과 얼마나 유사한지. 속도를 증가시키려면 정확도를 어느정도 희생해야 한다.

Exhaustive Search로 정확도 개선 가능

Approximating Similarity Search

Compression - Scalar Quantization (SQ)

Vector를 압축하여 어느정도의 정보 손실을 감안하면서 메모리 사용량을 낮추는 방법이다. 보통 32bit floating point를 사용하는데, 8bit의 unsigned integer로 압축하여 사용한다.

Pruning - Inverted File (IVF)

클러스터링 기법으로, 일정 벡터들을 군집화하여 Search Space를 줄인 후 검색속도를 개선하는 방법이다. 군집화 이후 쿼리가 들어왔을 때 쿼리와 가장 비슷한 군집만을 방문하도록 하는 것이다. 군집화 방법으로는 k-means clustering을 가장 많이 사용한다.

Inverted file (IVF) 라고 불리는 이유는 각 클러스터에 있는 포인트들을 인덱스로 가지고 있기 때문이다. 각 클러스터의 아이디와 클러스터 벡터들이 연결되어 있는 형태로 데이터 구조가 형성된다.

Introduction to FAISS

FAISS는 large scale 데이터에 대해서 효율적으로 유사도 검색과 군집화를 도와주는 페이스북의 오픈소스 라이브러리이다.

FAISS를 사용하기 위해서는 클러스터링된 벡터들을 확보해야 한다. 클러스터를 랜덤하게 지정하면 비효율적이기 때문에 적절한 군집화를 위해 데이터셋을 활용해야 하며, Quantization을 할 때도 scale할 비율을 계산해야 하기 때문에 학습데이터를 활용해야 한다.

따라서 Train과 Adding 단계가 필요하다.

Scaling up with FAISS

IVF 인덱스 만들기 (SQ 방식)

nlist = 100 # 클러스터 개수
quantizer = faiss.IndexFlatL2(d)
index = faiss.IndexIVFFlat(quantizer, d, nlist) # Inverted File 만들기
index.train(xb) # 클러스터 학습

index.add(xb) # 클러스터에 벡터 추가
D, I = index.search(xq, k) # 검색

IVF 인덱스 만들기 (PQ 방식)

• 벡터 압축 기법으로, 전체 벡터를 저장하지 않고 압축된 벡터만을 저장하여 메모리 사용량을 줄일 수 있음.

nlist = 100 # 클러스터 개수
m = 8 # subquantizer 개수
quantizer = faiss.IndexFlatL2(d)
index = faiss.IndexIVFPQ(quantizer, d, nlist, m, 8) # 각각의 sub-vector가 8 bits로 인코딩됨

index.train(xb) # 클러스터 학습
index.add(xb) # 클러스터에 벡터 추가
D, I = index.search(xq, k) # 검색