벡터 DB 구현하기

RAG가 화두가 되면서 벡터 DB를 사용하는 방법은 많았다. 하지만 벡터DB가 실제로 데이터를 어떻게 저장하고 관리하는지에 대해서는 잘 와닿지 않았다. 그래서 부스트캠프 미션으로 벡터 DB를 라이브러리에 기대지 않고 직접 구현해보기로 했다. 짝 설계로 큰 구조를 잡은 뒤 구현과 성능 개선을 이어갔고, 그 과정에서 계속 마주한 질문은 하나였다. 제한된 자원 안에서 검색의 정확도와 속도를 어떻게 함께 잡을 것인가.

벡터 DB는 무엇이며, 어떻게 구성되는가?

벡터 DB는 RDB나 문서형 NoSQL과 목적이 다르다. RDB는 테이블과 관계로 무결성을 지키고 NoSQL은 유연한 스키마로 대용량을 다루지만, 벡터 DB의 본질은 벡터 데이터를 저장하고 벡터 사이의 유사도로 검색하는 것이다. AI 모델이 만든 임베딩을 담고 유사성 검색을 수행하는 데 특화돼 있다.

벡터 DB는 두가지 핵심 요소로 구성된다.

• 벡터를 영속적으로 담는 데이터 저장소
• 쿼리 벡터에 가장 가까운 벡터를 찾는 검색 엔진

데이터 저장소와 검색엔진을 움직이는 기술은 다음과 같다.

• 텍스트와 같은 비정형 데이터를 의미가 담긴 고차원 벡터로 바꾸는 벡터 임베딩
• 방대한 벡터에서 유사한 것을 빠르게 찾는 인덱싱.
• 두 벡터가 얼마나 비슷한지 재는 유사도와 거리

무엇을 저장하고, 무엇으로 검색할 것인가?

벡터 DB에는 성격이 다른 두 종류의 데이터가 있다. 계산에 쓰는 임베딩 벡터와, 검색 결과로 사람에게 보여줄 원본 정보다. 형태도 다르고 접근 패턴도 다르다. 그래서 저장을 처음부터 둘로 나눴다. 접근 패턴이 다른 데이터를 한곳에 섞으면 결국 한쪽 기준에 맞추느라 다른 쪽이 손해를 보기 때문이다.

• data_meta.json: id, 생성/수정 시각, 원본 필드처럼 사람이 읽고 관리할 정보. JSON으로 둬서 조회와 수정이 직관적이다.
• vectors.bin: 임베딩 벡터를 id 순서대로 담은 바이너리 파일.

벡터는 바이너리로 둔다. 임베딩 벡터는 고정 길이의 실수 배열인데, JSON 텍스트로 저장하면 부동소수점 하나하나를 문자열로 풀어 쓰느라 용량이 커지고 파싱 비용도 든다. 바이너리로 순차 저장하면 압축적이고, 차원 크기를 알면 오프셋 계산만으로 특정 벡터에 바로 접근할 수 있다. 차원이 384면 벡터 하나가 384 * 4바이트로 고정되므로, id의 위치만 알면 파일 전체를 읽지 않고 그 지점만 읽으면 된다.

constructor(options = {}) {
  this.#vectorsPath  = options.vectorsPath  || path.join(__dirname, 'vectors.bin');
  this.#metadataPath = options.metadataPath || path.join(__dirname, 'data_meta.json');
  this.#indexPath    = options.indexPath    || path.join(__dirname, 'vectors_index.json');
  this.#vectorByteSize = (options.dimension || 384) * 4; // float32 * 차원
  // ...
}

임베딩 대상도 문서 전체가 아니라 검색에 쓰이는 필드(제목과 줄거리)만으로 좁혔다. 나머지 필드는 검색 대상이 아니라 결과를 보여줄 정보이므로 메타에 둔다.

처음에는 위와 같이 설계했으나 인덱스 파일vectors_index.json은 GC를 구현하며 추가되었다. 삭제를 soft delete로 처리하면 빈 공간이 쌓이고, 이를 주기적으로 정리하는 순간 id와 파일 내 물리적 위치가 더 이상 1:1로 맞지 않는다. 그래서 id와 바이너리 오프셋을 매핑하는 인덱스가 필요해졌다.

graph TD
  User[사용자] -->|텍스트| Embed[임베딩 모듈]
  Embed -->|벡터| DB[VectorDB]
  DB -->|벡터 저장| Bin[(vectors.bin)]
  DB -->|메타 저장| Meta[(data_meta.json)]
  DB -->|위치 매핑| Idx[(vectors_index.json)]
  User -->|쿼리| DB
  DB -->|KNN 검색| Bin
  DB -->|id로 메타 조회| Meta

임베딩 모듈은 특정 모델에 종속되지 않게 분리했다. 이렇게 두면 모델을 바꿀 때 벡터 DB 코드를 건드리지 않아도 되는데, 뒤에서 이 선택이 두 번 값을 한다. 임베딩 파이프라인은 프로그램당 하나면 충분하므로 싱글톤으로 뒀다.

class EmbeddingModule {
  static #pipe = null;
  async init(modelName = 'koxy-ai/gte-small') {
    if (EmbeddingModule.#pipe === null) {
      const { pipeline } = await import('@xenova/transformers');
      EmbeddingModule.#pipe = await pipeline('feature-extraction', modelName);
    }
  }
  async embedText(text) {
    const output = await EmbeddingModule.#pipe(text, { pooling: 'mean', normalize: true });
    return Array.from(output.data);
  }
}

임베딩에서 벡터의 방향은 곧 의미다. word2vec은 단어의 문맥적 관계를 벡터 공간에 반영해서, 의미가 비슷한 단어의 벡터가 서로 가깝게 놓인다. king에서 man을 빼고 woman을 더하면 queen에 가까워진다. 벡터의 위치와 방향 자체가 의미를 담는다는 이 성질이 다음 선택을 결정한다.

거리를 무엇으로 측정할 것인가?

KNN 검색의 핵심은 "가장 가까운 K개"를 찾는 것이고, 그 "가깝다"를 무엇으로 잴지가 검색의 성격을 정한다. 척도마다 재는 것이 다르다.

텍스트 검색에는 코사인 유사도가 맞다. 텍스트 임베딩에서 중요한 건 벡터의 크기가 아니라 방향, 즉 의미가 같은 쪽을 가리키는지다. 유클리드 거리는 문서 길이 같은 요인으로 크기가 달라지면 영향을 받지만, 코사인은 방향만 보기 때문에 문서 길이가 달라도 단어 분포가 비슷하면 높은 유사도를 준다. 앞서 임베딩할 때 normalize: true로 벡터를 정규화한 것도 이 선택과 이어진다. 방향만 볼 거라면 크기를 1로 맞춰두는 게 자연스럽다.

#cosineSimilarity(vecA, vecB) {
  let dotProduct = 0, magnitudeA = 0, magnitudeB = 0;
  for (let i = 0; i < vecA.length; i++) {
    dotProduct += vecA[i] * vecB[i];
    magnitudeA += vecA[i] * vecA[i];
    magnitudeB += vecB[i] * vecB[i];
  }
  magnitudeA = Math.sqrt(magnitudeA);
  magnitudeB = Math.sqrt(magnitudeB);
  if (magnitudeA === 0 || magnitudeB === 0) return 0;
  return dotProduct / (magnitudeA * magnitudeB);
}

정규화된 벡터라면 크기가 1이라 사실상 내적이 유사도를 지배하지만, 구현은 크기까지 나누는 정식 공식으로 두고 0으로 나누는 경우까지 막았다.

검색 구현: Brute-force KNN

검색은 최적화된 인덱스 없이 Brute-force KNN으로 시작했다. 게으른 선택이 아니라 의도한 순서다.

인덱싱은 근사(approximate) 검색이다. 모든 벡터를 다 비교하지 않고 범위를 좁혀 속도를 얻는 대신, 정답을 놓칠 가능성을 감수한다. 그렇다면 근사가 얼마나 정확한지 판단하려면 먼저 정확한 정답이 있어야 한다. Brute-force는 쿼리 벡터와 모든 벡터의 거리를 전부 계산하므로 느리지만 항상 정답을 준다. 이게 이후 인덱싱 개선의 기준선이 된다.

대신 비용이 크다. 벡터가 N개이고 차원이 D면 매 쿼리마다 $O(N \cdot D)$가 든다. 데이터가 늘수록 검색 시간이 선형으로 늘어난다.

임베딩이 검색 품질이다

기본 구현을 마치고 도서 검색을 돌리자 문제가 드러났다. "어린 왕자"로 검색했는데 정작 "어린 왕자" 책이 1위가 아니라 2위로 밀렸다. 검색 로직의 버그가 아니라, 임베딩 모델의 표현력이 결과 품질을 그대로 결정한다는 신호였다.

kobert는 문장을 하나의 벡터로 압축하는 과정에서 의미 손실이 컸다. 그래서 koxy-ai/gte-small로 교체했다.

성능 테스트

Brute-force KNN으로 검색을 구현하고, 임베딩 모듈로 koxy-ai/gte-small을 선택한 뒤 성능 테스트를 진행했다. 데이터를 100개, 1000개, 5000개로 늘려가며 시간을 재고, Node.js 내장 V8 프로파일러로 CPU와 메모리를 함께 확인했다.

--prof 플래그로 파일을 실행하면 v8.log 파일이 생성되고 해당 파일을 --prof-process로 변환하면 어떤 함수가 CPU를 가장 많이 사용했는지 확인된다.

데이터가 많아질 수록 병목이 확인되었다. 유사도 계산과 KNN 정렬 구간이 리소스를 가장 많이 사용하는 것을 확인했고, 쿼리 한 번에 메모리 1GB를 사용하기도 하였다.

병목이 전수 비교 자체에 있다는 것으로 판단하고, 단순 Brute-force KNN으로 검색을 하는 것이 아니라 비교 대상을 줄이는 방법인 인덱싱을 추가하도록 하였다.

개선 단계

다른 데이터에는 다른 임베딩

초기 주어진 데이터는 도서 정보 데이터로 정형 데이터에서 도서 제목, 도서 내용으로 임베딩을 하기 때문에 트랜스포머 임베딩을 사용하였다. 그러나 개선 단계에서 주어진 데이터는 약 234만건의 대용량 질문 쌍이었다. 이 데이터에서 단어들 사이의 관련도, 노출 빈도, 거리를 벡터로 표현하기 위해서는 TF-IDF가 구조적으로 적절하다고 고려하였다.

TF-IDF는 단어의 빈도와 차별성을 동시에 재는 지표이다. 한 문서에서 자주 등장할수록 커지는 TF와 전체 문서에서 드물게 등장할 수록 커지는 IDF를 종합한 지표이다. 노출빈도, 단어의 변별력은 측정할 수 있으나 문맥 의도는 측정하기 어렵다.

$$TF(t, d) = (\text{문서 d에서 단어 t의 등장 횟수}) / (\text{문서 d의 전체 단어 수})$$

TF-IDF 벡터도 코사인 유사도로 비교할 수 있게 마지막에 정규화했다. 검색 로직(코사인 기반 KNN)을 그대로 두고 임베딩 방식만 교체할 수 있었던 것은, 벡터 DB가 주입과 검색만 책임지고 임베딩 모듈을 분리해둔 설계 덕이다.

embedText(text) {
  const tokens = this.#tokenize(text);
  const vocaSize = Object.keys(this.vocabulary.word_to_idx).length;
  const vector = new Array(vocaSize).fill(0);

  const termFrequency = {};
  tokens.forEach(t => termFrequency[t] = (termFrequency[t] || 0) + 1);

  for (const token in termFrequency) {
    const idx = this.vocabulary.word_to_idx[token];
    if (idx !== undefined) {
      const tf = termFrequency[token];
      const df = this.vocabulary.df[token];
      const idf = Math.log(this.vocabulary.total_documents / (df + 1)); // 0 division 방지
      vector[idx] = tf * idf;
    }
  }

  // 코사인 유사도 계산을 위한 정규화
  const norm = math.norm(vector, 2);
  const normalized = norm === 0 ? vector : vector.map(v => v / norm);
  return new Float32Array(normalized);
}

LSH 도입

기존 병목이 전수 비교였기 때문에 벡터 인덱싱 기술을 도입하였다. 모든 벡터를 비교하지 않고, 데이터를 구조화해 검색 범위를 좁히도록 하였다.

일반적으로 텍스트 임베딩은 고차원이므로 트리 기반은 선택하지 않았다. HNSW는 가장 강력하지만 그래프 구조라 직접 구현하기엔 복잡하고 메모리 부담도 크다. 이번 목표는 "라이브러리 없이 원리를 이해하는 것"이므로, 성능 최상보다 손으로 구현하며 동작을 완전히 이해할 수 있는지를 기준으로 삼아 LSH을 선택하여 도입했다.

LSH도 해시 함수를 무엇으로 만드느냐로 갈린다. 유사도 기준에 따라 MinHash(Jaccard), SimHash(코사인), Hyperplane LSH(유클리드)가 있다. 거리 척도를 코사인으로 정했으니, 코사인 기반인 SimHash(랜덤 투영)가 일관된 선택이다. 고차원 공간에 무작위 초평면 여러 개를 세우고, 벡터가 각 평면의 어느 쪽에 있는지를 내적의 부호로 정해 시그니처를 만든다.

// 벡터를 밴드별 해시로 변환: 각 초평면과의 내적 부호가 1비트
#getBandHashes(vector) {
  const bandHashes = [];
  for (const band of this.#hashTables) {
    let bandHash = '';
    for (const plane of band.planes) {
      const dotProduct = math.dot(Array.from(vector), Array.from(plane));
      bandHash += (dotProduct >= 0) ? '1' : '0';
    }
    bandHashes.push(bandHash);
  }
  return bandHashes;
}

무작위 초평면은 정규분포에서 뽑은 벡터로 만든다. k비트 시그니처를 b개 밴드로 나누고, 한 밴드라도 해시가 같으면 유사 후보로 본다. 여기서 밴드 수 b와 밴드당 비트 수 r이 정확도와 속도를 조절하는 손잡이다. 밴드를 많이 나누면(b 크게, r 작게) 유사한데 놓치는 false negative는 줄지만 false positive가 늘고, 적게 나누면 그 반대다. 결국 목표하는 정밀도와 재현율에 맞춰 조정하는 문제다.

constructor(options = {}) {
  this.#numHashBits = options.numHashBits || 64; // k: 전체 비트 수
  this.#numBands    = options.numBands    || 8;  // b: 밴드 수
  this.#bitsPerBand = this.#numHashBits / this.#numBands; // r
  // ...
  if (this.#numHashBits % this.#numBands !== 0) {
    throw new Error('numHashBits는 numBands로 나누어 떨어져야 합니다.');
  }
}

추가와 검색은 이 밴드 해시를 버킷 키로 쓴다. 벡터를 넣을 땐 밴드별 버킷에 id를 담고, 검색할 땐 쿼리의 밴드 해시와 일치하는 버킷의 id만 후보로 모은다.

// 검색: 쿼리와 같은 버킷에 담긴 id만 후보로 수집
getCandidateIds(queryVector) {
  const candidateIds = new Set();
  const queryBandHashes = this.#getBandHashes(queryVector);
  queryBandHashes.forEach((hash, bandIdx) => {
    const bandKey = `band_${bandIdx}`;
    if (this.#index[bandKey] && this.#index[bandKey][hash]) {
      this.#index[bandKey][hash].forEach(id => candidateIds.add(id));
    }
  });
  return Array.from(candidateIds);
}

이제 검색은 전수 비교가 아니라 후보군만 도는 구조로 바뀐다. LSH가 후보 id를 좁혀주면, 그 벡터들만 파일에서 읽어 코사인 유사도를 계산한다.

async findNearestNeighbors(queryVector, k) {
  // LSH로 후보군을 먼저 좁힌다
  const candidateIds = new Set(this.#lsh.getCandidateIds(queryVector));

  const searchResults = [];
  const fileHandle = await fs.open(this.#vectorsPath, 'r');
  try {
    for (const id of candidateIds) {
      if (!this.#metadata[id] || this.#metadata[id].isDeleted) continue;

      // 후보 벡터만 오프셋으로 읽어 유사도 계산
      const offset = this.#index[id];
      const buffer = Buffer.alloc(this.#vectorByteSize);
      await fileHandle.read(buffer, 0, this.#vectorByteSize, offset);
      const candidateVector = new Float32Array(buffer.buffer, buffer.byteOffset, buffer.length / 4);

      const similarity = this.#cosineSimilarity(new Float32Array(queryVector), candidateVector);
      searchResults.push({ id, similarity });
    }
  } finally {
    await fileHandle.close();
  }

  searchResults.sort((a, b) => b.similarity - a.similarity);
  return searchResults.slice(0, k).map(r => ({ ...this.#metadata[r.id], similarity: r.similarity }));
}

LSH의 본질은 "약간의 정확도를 내주고 속도를 크게 얻는" 교환이고, 밴딩 파라미터가 그 교환의 손잡이다. 도입 후 검색은 후보군만 비교하는 구조로 바뀌었고, 만 건 규모에서도 쿼리가 약 1~2초에 끝난다.

마치며

이 미션을 수행하면서 벡터 DB가 어떻게 동작하여 데이터를 저장하는 지에 대해서 학습하게 되었다.

RAG는 검색(retrieval)과 생성(generation)으로 나뉘고, 그 검색 계층의 심장이 이번에 구현한 벡터 검색이다. 임베딩이 검색 품질을 좌우하는 이유, 데이터가 커지면 인덱싱이 필수가 되는 이유, 인덱싱이 필연적으로 정확도와 속도의 교환인 이유를 토이 프로젝트를 통해 측정하며 확인했다. 그래서 Pinecone, Weaviate, FAISS, Chroma 같은 실제 벡터 DB의 설명을 볼 때, HNSW냐 IVF냐, 어떤 거리 척도를 쓰느냐가 내부에서 무엇을 조절하는지에 대해서 이해할 수 있게 되었다.

하지만 여전히 실제로 문장을 하나의 벡터로 압축할 떄 발생하는 의미 손실을 어떻게 줄일 수 있는지, 데이터 구조가 1TB 이상으로 커지면 어떤 부분을 개선해야 할지에 대해서 고민할 필요가 있다.