Vector Search - IVF

Vector Similarity Search

Vector Similarity Search란 오브젝트(검색어, 사진, 비디오, 상품 등) 간의 “유사성”을 정량적으로 측정하기 위해 각 오브젝트를 다차원 벡터 공간으로 Embedding한 후, 벡터 공간에서 정의된 “거리”를 이용하여 가장 유사한 오브젝트를 찾는 것이다.

이를 위해서는 먼저 숫자가 아닌 오브젝트를 Vector Representative로 변환하는 기술이 필요하다. 이전까지 Locality Sensitive Hashing(LSH)와 같이 직접 유사도을 정의하는 Hash 함수를 사용했다면 최근 들어서는 Word2Vec이나 Transformer처럼 유사도 정보를 최대한 잃지 않도록 학습된 머신러닝 모델을 사용하는 방식이 많아졌다.

Vector Similarity Search를 이용하면 다양한 분야에서 머신러닝을 서비스에 녹여낼 수 있다. 사용자에게 검색 결과를 찾아주거나, 판매자가 올린 상품 설명으로 카테고리를 추천해주는 서비스를 만들 수도 있다. 우리가 이미 잘 아는 GPT 또한 Context의 Vector Representative와 가장 유사한 단어들을 찾아나가는 과정이라 볼 수도 있다.

이러한 검색 모델은 비교할 후보군의 크기가 작으면 쉬워 보이지만 그 크기가 수백만에 달하면 훨씬 까다로워 진다. 만약 후보가 되는 모든 오브젝트를 탐색(Exhaustive Search)한다면 가장 관련성이 높은 후보를 찾을 순 있으나, 후보군의 크기(N)가 크면 시간이 굉장히 많이 걸릴 것이다. 또 머신러닝 모델이 표현하는 벡터의 차원(D)이 크다면 유사도를 한번 계산하는데도 훨씬 오래 걸릴 것이다.

이를 해결할 수 있는 방법 중 하나로 Approximate Nearest Neighbors (ANN) Search가 있다. 모든 오브젝트와의 거리를 계산하는 대신에 일종의 구역을 만들어 그 구역의 대표값과 먼저 거리를 계산하게 하는 것이다. 그런 다음에는 가까운 구역에 속한 오브젝트들부터 다시 거리를 계산하여 순위를 매긴다. 즉, 벡터 공간에서 일종의 Index를 만드는 것이다.

K Nearest Neighbors (KNN) Search

질의(Query)가 하나 주어질 때마다 가장 유사한 K개의 오브젝트를 찾는다고 가정하자. 일반적으로 유사도를 측정할 때 사용하는 metric은 크게 3가지이다.

1. L2 Distance (L2)
   • 작을수록 유사함
   • 0 보다 큰 값
2. Inner Product (IP)
   • 클수록 유사함
3. Cosine Similarity
   • 클수록 유사함
   • -1 에서 1 사이의 값

참고로 임베딩 벡터가 Normalization되어 있다면 셋 모두 거의 동일한 의미를 지니게 된다.

Flat Index

Flat Index란, Index를 만들 때 후보 벡터들을 수정하지 않고 그대로 사용하는 것을 말한다. 근사나 클러스터링을 하지 않기 때문에 검색이 정확하지만 후보 벡터의 개수가 많아질수록 검색 시간이 크게 증가한다. 따라서 검색의 품질이 매우 중요하거나 데이터의 크기가 작을 때만 사용 가능하다.

“검색 속도”와 “검색 품질” 사이에는 trade-off 가 존재한다. Flat Index의 경우 검색 속도를 높이기 위해서는

1. Vector Dimension을 줄이거나
2. Vector Representative의 bit를 줄이면 된다.

하지만 위와 같이 벡터의 크기를 줄이는 방법은 모델 및 학습 방법에 크게 의존한다. 후보 데이터가 커질 때 마다 SLA를 맞추기 위해서 모델 스펙을 바꾸고 새로 학습시키는 것은 시간도 오래걸리고 언젠가 한계에 부딪힐 것이다. 따라서 벡터의 크기를 줄이는 대신 검색 범위를 줄이는 방법을 찾게 되었다.

Inverted File Index (IVF)

ANN Search에 가장 직관적으로 써먹어볼 수 있는 알고리즘은 k-mean Clustering일 것이다. 벡터 공간 위에서 k개의 Cluster를 미리 만들어둔 뒤, 질의(Query)가 들어오면 k개의 Centroid와 먼저 유사도 계산을 진행한다. 그 다음 가장 가까운 Centroid의 Cluster에 속한 오브젝트들과 거리를 계산한다. Cluster가 곧 Index인 셈이다.

Cluster의 개수가 적을 수록 검색이 정확해지는 대신 검색 시간이 많이 증가한다. 따라서 검색의 품질이 매우 중요하거나 데이터의 크기가 작을 때는 학습시킬 Cluster 개수를 줄이거나, 학습 후 찾을 Cluster의 개수를 늘리면 된다.

1. k-mean clustering로 데이터를 파티셔닝
   • nlist = (생성할 파티션 개수, 즉 k)
2. Query vector가 주어지면 각 클러스터의 centroid와 유사도를 비교
   • Query vector가 클러스터 사이의 경계에 걸치면 성능이 떨어질 수 있으므로, nprobe를 조절해 탐색할 클러스터의 개수를 늘린다.
   • nprobe = (비교할 centroid 개수)
3. 가까운 centroid와 대응되는 클러스터의 벡터들과 유사도 계산
4. 가장 유사한 순으로 top K 리턴

“검색 속도”와 “검색 품질” 사이의 trade-off

• 클러스터 개수(nlist)가 많아지면 더 많은 centroid와 비교하지만, 그 속 데이터들과는 덜 비교하게되어 검색 속도가 빨라진다.
• 탐색할 클러스터 개수(nprobe)가 많아지면 유사도를 비교할 데이터의 개수가 커지므로 검색 속도는 느리지만 검색 품질이 좋아진다.

Product Quantization (PQ)

탐색할 벡터의 개수를 줄여주는 IVF와 달리, PQ는 탐색할 벡터의 크기를 줄여주는 데이터 압축 기법으로, 정확도(Recall)를 조금 더 희생해서 검색 속도를 높일 수 있다.

1. D 차원의 공간을 m 개의 subspace로 나눈다.
   • 각 vector는 m 개의 D/m 차원 sub-vector로 나뉜다.
2. 각 subspace 마다 $k^* := k/m$ 개의 클러스터를 생성한다. ($k^*$-mean clustering)
   • 총 클러스터 개수는 여전히 $k$ 개
   • 각 subspace 마다 $1$ ~ $k^*$의 클러스터 ID 생성
3. 각 vector는 sub-vector와 가장 가까운 클러스터의 ID로 조합된 quantized vector가 된다.
   • 각 entry가 $1$ ~ $k^*$의 값을 가진 m 차원의 벡터가 된다.
   • nbits = (quantized vector의 element의 bit 수)
   • $2^{nbits}$ = (각 subspace 마다 생성할 클러스터 개수, k/m)
4. 후보 벡터들의 quantized vector를 계산하여 저장한다.
   • $c = (c_1, \ldots , c_m)$.
5. 쿼리가 들어오면 각 subspace 별로 모든 centroid와의 거리를 미리 계산한 lookup 테이블을 만든다.
   • $d_{i}(j)$ = “i번째 subspace에서 클러스터 ID가 j인 클러스터의 centroid” 와 “쿼리의 i번째 sub-vector” 사이의 거리
6. 이후 각 후보의 quantized vector로 lookup 테이블을 조회 및 집계하여 partial distance를 계산한다.
   • L2 공간인 경우, $d(c) = \sqrt{d_1(c_1)^2 + d_2(c_2)^2 + \cdots + d_m(c_m)^2}$
   • $d(c)$ 값이 작은 순으로 top K 리턴

IVF-PQ

IVF와 PQ를 동시에 사용할 수도 있다. PQ를 이용해 후보 벡터의 quantized vector를 계산하고, IVF로 후보 벡터의 인덱스를 만든다.

PQ와 마찬가지로 쿼리가 들어오면, 각 subspace 별로 쿼리의 sub-vector와 centroid들 간에 거리가 계산된 lookup table을 만든다. 이후 쿼리와 가까운 인덱스 속 quantized vector들 사이에 partial distance를 계산하여 가장 가까운 순으로 top K를 뽑는다.

Framework

다음 프레임워크들이 IVF 또는 IVF-PQ를 지원하고 있다.

• faiss
• lance
• milvus
• pgvector

사실 Vector DB를 사용해보면 어떤 알고리즘을 지원하느냐 또는 성능이 얼마나 좋은가는 크게 중요치 않은 것 같다. 실무에 적용할 때는 벡터 외의 속성 추가가 가능하여 필터링을 지원하는지, 또 백엔드 로직을 얼만큼 수정하지 않아도 되는지가 중요했다.

이 중 pgvector는 postgres DB의 extension 중 하나로, 테이블의 인덱스를 생성하는 것과 비슷하게 IVF 인덱스를 생성할 수 있다. 따라서 벡터 외의 속성을 다루기가 편리하고 러닝 커브가 크지 않아 가장 생산성이 높다고 느꼈다.

faiss는 인덱스 파일을 받아와 프로세스 메모리에 올려야 하기 때문에 번거롭긴 하다. 대신 Disk IO를 탈 일이 없기 때문에 가장 빠른 성능을 자랑한다.

Let’s Practice

다음 Github 링크에 상세한 설정을 정리해 두었습니다.

• https://github.com/ivoryRabbit/mini-search

Sentence BERT와 pgvector를 사용해 간단한 영화 검색 엔진을 개발해볼 수 있었다. Docker에서 pgvector가 설치된 postgres image를 이용해 컨테이너를 띄우고 CREATE EXTENSION IF NOT EXISTS vector 쿼리를 실행시키면 pgvector를 사용해 볼 수 있다.

version: "3.7"

services:
  database:
    image: ankane/pgvector:v0.5.0
    container_name: pgvector-dev
    hostname: pgvector-dev
    ports:
      - "5432:5432"
    environment:
      - POSTGRES_USER=postgres
      - POSTGRES_PASSWORD=postgres
    volumes:
      - ./init.sql:/docker-entrypoint-initdb.d/init.sql

Python에서는 SQLAlchemy가 pgvector를 위한 ORM을 제공하고 있다. Vector 타입의 컬럼이 포함된 Entity를 정의하고 Session을 열면 데이터를 insert 하거나 select할 수 있다.

Model

from pgvector.sqlalchemy import Vector
from sqlalchemy import Integer, String, SmallInteger, Float
from sqlalchemy.orm import mapped_column

from api.entity.base import Base


class Movie(Base):
    __tablename__ = "movie"

    id = mapped_column(Integer, primary_key=True)
    title = mapped_column(String, nullable=False)
    genres = mapped_column(String)
    year = mapped_column(SmallInteger)
    embedding = mapped_column(Vector(768))

Insert

with get_session().begin() as session:
    session.execute(insert(Movie), values)

Make Index

pgvector는 초기에 IVF-Flat 만 지원하다 최근 HNSW가 추가되었다. Contributor 피셜로는 머지않아 IVF-PQ와 ScaNN 알고리즘이 추가될 예정이라 한다.

쿼리를 통해 다음과 같이 IVF Index를 추가할 수 있다. 여기서 lists와 probes는 각각 “생성할 클러스터의 개수”와 “탐색할 클러스터의 개수”를 의미한다. 참고로 테이블 하나에 여러 Index를 추가할 수도 있고, segment 별로 서로 다른 Index를 운영할 수도 있다.

NUM_LISTS = 128

table_name = Movie.__tablename__

with get_session().begin() as session:
    session.execute(
        text(
            f"""
                CREATE INDEX ON :table_name
                USING ivfflat (embedding vector_cosine_ops) WITH (lists = :lists)
            """
        ).bindparams(table_name=table_name, lists=NUM_LISTS)
    )

Select

NUM_PROBES = 8

with get_session().begin() as session:
    session.execute(
        text("SET ivfflat.probes = :probes").bindparams(probes=NUM_PROBES)
    )

    row = session.scalars(
        select(Movie)
        .order_by(
            Movie.embedding.cosine_distance(query_vector),
            Movie.year.desc(),
        )
        .limit(size)
    ).all()