[MRC] Passage Retrieval – Scaling up

작성일 | In |

Passage Retrieval and Similarity Search

passage의 개수가 많아질수록 쿼리와 가장 가까운 문서(쿼리에 대한 답)를 찾기 힘들 것이다. 질문과 가장 가까운 답을 찾아야만 한다 => similarity search.

참고로 sparse/dense embedding은 인코딩에 관한 것이었고, 지금은 검색 방법에 대해 말한다.


주어진 질문 벡터 q에 대해 패시지 벡터 v들 중 가장 질문과 관련성이 높은 벡터를 찾아야 한다. 내적(inner product)이 가장 클수록 관련성이 높다.

데이터가 방대하니 브루트포스 검색(exhaust search)은 좋지 않다. 다만 다른 검색 방법들은 trade-off가 있다.

  • 가지고 있는 벡터가 많을수록 쿼리당 특정 개수의 벡터를 찾는 데에 오래 걸린다.
  • 마찬가지로 RAM에 벡터들을 모두 저장하기엔 메모리 요구량이 많을 수밖에 없다. 디스크에서 불러 온다면 속도가 느려진다.
  • 다른 방법을 사용한다면 브루트포스 방법과 정확도면에서 얼마나 차이나는지도 고려해야 한다.

더 정확한 검색을 하려면 오랜 시간과 많은 메모리가 필요하기 마련이다.

corpus 크기가 커질수록 탐색 공간이 커지고 검색에 어려움이 많은데, 메모리도 많이 잡아 먹는다. 특히, sparse embedding은 이러한 문제가 더 심하다.


  • compression - scalar quantization
    • 벡터를 압축해서, 하나의 벡터가 적은 용량을 차지하도록 한다. 압축량이 높아질수록 메모리는 적게 잡아먹지만 정보손실이 커진다.
  • pruning - inverted file : pruning : search space를 줄여서 검색 속도 개선 (clustering + inverted file을 활용한 search)

    • 점들을 정해진 클러스터에 소속시켜서 군집을 이루도록 한다. 그리고 해당 쿼리에 근접한 특정 개수의 클러스터들만 보는데, 이 클러스터들에 대해서는 전수검색한다.

    • clustering : 전체 벡터 공간을 k개의 cluster로 나눔(e.g. k-means clustering)
    • iverted file : 벡터의 인덱스 = inverted list structure (각 클러스터에 속해 있는 점들을 역으로 인덱스로 가지고 있음…각 클러스터의 centroid id와 해당 클러스터의 벡터들이 연결되어 있음.)

정리하자면 주어진 쿼리벡터에 대해 근접한 centroid 벡터를 찾고, 찾은 클러스터의 inverted list 내 벡터들에 대해 탐색 수행하는 과정.


Faiss

일단 위의 탐색 방법을 쓰기 위해서 cluster를 정해줘야 한다. 이때, 데이터 점들의 분포를 보고 클러스터를 지정하게 된다. 그래서 벡터를 학습해야 함.

먼저 Faiss index를 결정하고, 그걸 바탕으로 탐색한다.


Scaling up with FAISS

  • 브루트포스 방식으로 모든 쿼리와 벡터를 탐색하는 단순한 인덱스 만들기
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
# 세팅
d = 64 # 벡터의 차원
nb = 100000 # 데이터 베이스 크기
nq = 10000 # 쿼리 개수
xb = np.random.random((nb, d)).astype('float32') # 데이터 베이스 예시
xq = np.random.random((nq, d)).astype('float32') # 쿼리 예시

# 인덱스 생성
index = faiss.IndexFlatL2(d) # 인덱스 빌드하기
index.add(xb) # 인덱스에 벡터 추가하기
# Clustering + Quantize 하지 않으므로 Train 안 함.

# 탐색
k = 4 # 가장 가까운 벡터 4개를 찾고 싶음
D, I = index.search(xq, k) # 검색하기
# D : 쿼리와의 거리, I : 검색된 벡터의 인덱스


  • IVF(Inverted File) with Faiss
  • 클러스터를 사용해서 가까운 클러스터 내 벡터들만 비교
  • 빠른 검색 가능
  • 클러스터 내에서는 전체 벡터들와 거리 비교
1
2
3
4
5
6
7
8
nlist = 100 # cluster 개수
quantizer = faiss.IndexFlatL2(d)
index = faiss.IndexIVFFlat(quantizer, d, nlist) # Inverted File 만들기
index.train(xb) # clustering 학습

index.add(xb) # cluster에 벡터 추가하기
D, I = index.search(xq, k) # 검색하기


  • IVF-PQ with Faiss
  • 벡터 압축 기법(PQ) 사용
  • 전체 벡터가 아닌 압축된 벡터만 저장해서 메모리 사용량을 줄임
1
2
3
4
5
6
7
8
9
10
nlist = 100 # cluster 개수
m = 8 # subquantizer 개수
quantizer = faiss.IndexFlatL2(d)
index = faiss.IndexIVFPQ(quantizer, d, nlist, m, 8) # 각각의 sub-vector가 8 bits로 인코딩

index.train(xb) # clustering 학습
index.add(xb) # cluster에 벡터 추가하기

D, I = index.search(xq, k) # 검색하기


  • Using GPU with Faiss
  • GPU를 사용해서 거리 계산을 위한 행렬곱 연산 과정에서 속도가 빠르다.
  • GPU 메모리 제한과 메모리 random access 시간이 느린 것에 유의해야 한다.
1
2
3
4
5
6
7
8
9
res = faiss.StandardGpuResources() # 단일 GPU 사용하기
index_flat = faiss.IndexFlatL2(d) # 인덱스 (CPU) 빌드하기

# GPU 인덱스로 옮기기
gpu_index_flat = faiss.index_cpu_to_gpu(res, 0, index_flat)
gpu_index_flat.add(xb) 

D, I = index.search(xq, k) # 검색하기


  • Using Multiple GPUs with Faiss
  • 여러 개의 GPU로 연산 속도를 높인다.
1
2
3
4
5
6
7
8
cpu_index = faiss.IndexFlatL2(d) 

# GPU 인덱스로 옮기기
gpu_index = faiss.index_cpu_to_all_gpus(cpu_index)
gpu_index.add(xb) 

D, I = index.search(xq, k) # 검색하기

참고로 CPU를 사용한 Approximating Similarity Search 방법이 GPU를 사용한 브루트포스 방식보다 훨씬 빠르다.