KyungHyun Lim

KyungHyun Lim

ML/AI/SW Developer

© 2023

Dark Mode

Scaling up with FAISS
Oct 15, 2021
  • NLP
  • ML_AI

1. Passage Retrieval and Similarity Search

  • MIPS
    • 주어진 question 벡터 q에 대해 Passage 벡터 v 들중 가장 질문과 관련된 벡터를, 내적 값을 기준으로 찾는것
  • 이전 포스팅(sparse, dense embedding)에서는 모든 벡터를 비교하는 브루트포스 방식 사용
  • Challenges
    • 실제로 검색해야할 데이터는 훨씬 방대함
      • 5,000,000 개 위키피디아
      • 수십 억, 조 단위 까지 커질 수 있음
      • 즉, 모든 문서 임베딩을 일일히 보면서 검색할 수 없음

1.2 Tradeoffs

  • Search speed
    • 쿼리 당 유사한 벡터를 k개 찾는데 얼마나 걸리는지?
    • 가지고 있는 벡터가 많을 수록 더 오래 걸림
  • Memory usage
    • 벡터를 사용할 때, 어디에서 가져올 것인지?
    • RAM에는 너무 많은 것을 올리기 어려움
    • 디스크에서 불러오는 것은 병목현상 등의 시간적 비용이 존재
  • Accuracy
    • 브루트포스 검색결과와 얼마나 비슷한가?
    • 속도가 높아지면 정확도가 떨어지는 경우가 많음

2. Approximating Similarity Search

2.1 Compression

  • 벡터를 압축하여, 하나의 vector가 적은 용량을 차지하도록!
    • 압축량이 늘어나면 메모리는 비용은 줄어들지만, 정보 손실이 발생할 수 있음

2.2 Pruning

  • Inverted File
    • Search space를 줄여서 속도를 개선, dataset의 subset만 방문
    • Clustering과 inverted file을 활용
      • Clustering: 전체 vector space를 k개의 cluster로 나눔(E.g. K-means)
      • Inverted file: vector의 index = inverted list structure
        • 각 cluster의 centroid id와 해당 cluster의 벡터들이 연결되어 있는 형태
    • 과정
      1. 주어진 query vector에 대해 유사 centroid id 탐색
      2. 찾은 cluster의 list 내 vector들에 대해서만 search 수행

3. Introduction to FAISS

3.1 FAISS 란?

  • 유사도 탐색을 효율적으로 도와주는 라이브러리

3.2 Passage Retrieval with FAISS

  1. Train index and map vectors
    • cluster를 확보하기 위해 학습이 필요
    • 퀀타이제이션 과정에서도 얼마나 scale up 할지 팍악할 필요가 있음
    • 학습할 데이터와 더할 데이터를 어떻게 나누는지는 사용자의 몫
  2. Search based on FAISS index
    • nprobe: 몇 개의 가장 가까운 cluster를 방문하여 search 할 것인지?

4. Scaling up with FAISS

4.1 브루트포스로 벡터와 쿼리 비교해서 찾아보기

  1. 더미 데이터 만들기
d = 64          # 벡터의 차원
nb = 100000     # 데이터베이스 크기
nq = 1000       # 쿼리 개수

# 임의 데이터 베이스 및 쿼리 벡터 만들기
xb = np.random.random((nb, d).astype('float32'))
xq = np.random.random((nq, d).astype('float32'))
  1. 인덱스 만들기
index = faiss.IndexFlatL2(d)    # 인덱스 빌드하기
index.add(xb)                   # 인덱스에 벡터 추가하기
  1. 검색하기
k = 4                       # 찾고 싶은 top-k
D, I = index.search(xq, k)  # 검색하기
# D: 쿼리와의 거리, I: 검색된 벡터의 인덱스

4.2 IVF로 찾아보기

  1. IVF 인덱스 생성 후 탐색
    • 클러스터 내에서는 여전히 전체 벡터와 거리 비교 (Flat)
     nlist = 100                         # 클러스터 개수
 quantizer = faiss.INdexFlatL2(d)    
 index = faiss.IndexIVFFlat(quantizer, d, nlist) # Inverted File 만들기
 index.train(xb)                     # 클러스터 학습하기
     # K-means 알고리즘으로 nlist개의 클러스터 생성
 index.add(xb)                       # 클러스터에 벡터 추가하기
 D, I = index.search(xq, k)          # 검색하기
  2. 다른 압축기법 사용하기
    • 전체 벡터를 저장하지 않고 압축된 벡터만을 저장
    • 메모리 사용량을 줄일 수 있음
     nlist = 100                         
 m = 8                               # subquantizer 개수
 quantizer = faiss.INdexFlatL2(d)    
 index = faiss.IndexIVFPQ(quantizer, d, nlist, m, 8)
                                     # 각 sub-vector가 8 bit로 인코딩
 index.train(xb)                     
 index.add(xb)                       
 D, I = index.search(xq, k)          # 검색하기

  3. GPU도 활용해보기

    • 빠른 연산을 위해 활용 가능 - 벡터 연산
    • GPU 메모리 제한에 따라 시간이 약간 느려질 수도 있음
     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 = gpu_index_flat.search(xq, k)
    • 멀티 GPU 활용
     index_flat = faiss.IndexFlatL2(d)

 # 멀티 GPU로 인덱스 옮기기
 gpu_index_flat = faiss.index_cpu_to_all_gpus(index_flat)

 gpu_index_flat.add(xb)
 D, I = gpu_index_flat.search(xq, k)

5. 참조