메모리 수치를 명시한 top-k 유사도 측정을 위한 패킹된 행렬(Packed matrix) 및 SIMD 내적(dot products).

저는 약 20,000개의 짧은 문서와 각 문서당 384차원의 임베딩(embedding)을 가진 Flutter 앱을 가지고 있었습니다. 사용자의 쿼리 임베딩(query embedding)을 가져와 코사인 유사도(cosine similarity)를 통해 가장 가까운 5개의 문서를 찾아 보여주는 방식입니다. 서버 왕복 없이 온디바이스(On device)에서 처리합니다.

첫 번째 버전은 아주 당연한 방식입니다. List<List<double>>에 임베딩을 저장하고, 모든 행의 점수를 계산하는 루프를 돌린 뒤, 정렬하여 상위 5개를 뽑는 방식입니다.

double cosine(List<double> a, List<double> b) {
  var dot = 0.0, na = 0.0, nb = 0.0;
  for (var i = 0; i < a.length; i++) {
...

이 방식은 정확합니다. 하지만 제 기기(Apple Silicon, Dart 3.11)에서 20,000행 인덱스를 대상으로 쿼리당 15.9ms가 소요됩니다. 15.9ms라면 사용자가 타이핑할 때 매 키 입력마다 순위를 재조정(re-rank)하고 싶을 때 그 지연을 느낄 수 있으며, 코퍼스(corpus)의 크기에 따라 선형적으로 증가합니다.

비용이 발생하는 원인은 알고리즘이 아닙니다. 이 정도 규모에서는 선형 스캔(linear scan)이 올바른 알고리즘입니다. 문제는 레이아웃(layout)입니다. List<List<double>>은 별도의 힙(heap) 객체들을 가리키는 포인터들의 리스트이며, 각 double은 박싱(boxed)된 64비트 부동 소수점입니다. 또한 내부 루프는 포인터를 추적하며 순차적으로 읽도록 설계되지 않은 메모리를 읽게 됩니다.

패킹된 레이아웃 (The packed layout)

vector_kit는 전체 인덱스를 하나의 연속된 Float32List로 저장하고 SIMD로 점수를 계산합니다. 동일한 선형 스캔을 사용하며 결과도 정확히 같지만, 메모리 구조가 다릅니다.

import 'package:vector_kit/vector_kit.dart';

// rows: 임베딩 모델로부터 얻은 20000 x 384 크기의 List<List<double>>.
...

20,000행 인덱스에 대해 상위 5개를 찾는 topKCosine은 1.4ms 만에 실행됩니다. 이는 앞서 언급한 15.9ms의 작업과 동일한 작업이지만 약 11배 더 빠르며, 아무것도 근사(approximated)하지 않았기 때문에 랭킹 결과도 동일합니다.

래퍼(wrapper)를 줄여준 한 가지 디테일은 query가 일반 List<double>이라는 점입니다. 모델에서 나오는 임베딩은 이미 List<double> 형태이므로, 호출 시점에 Float32List로 변환할 필요 없이 바로 topKCosine에 전달됩니다. 행렬은 매 쿼리마다 만드는 것이 아니라, 구축할 때 한 번만 패킹됩니다.

그 밑단의 기본 연산은 하드웨어 SIMD (Single Instruction, Multiple Data)로 매핑되는 dot (내적) 연산입니다. 768차원 기준으로 호출당 142 ns가 소요되며, 이는 두 개의 List<double>을 사용하는 동일한 루프의 665 ns와 대조적입니다. 만약 내적 값만 필요하다면 다음과 같이 원시 곱(raw product)을 사용할 수 있습니다:

final a = Float32List.fromList(embA);
final b = Float32List.fromList(embB);
final d = dot(a, b); // 768차원에서 142 ns

이 격차는 더 큰 규모에서도 유지됩니다. 100,000개의 행에 대해 k=10인 topKCosine은 13.3 ms가 소요됩니다. 동일한 데이터에 대해 전체 스캔 후 정렬(full-scan-and-sort)하는 베이스라인 방식은 82 ms가 걸립니다. 패킹된 버전은 100,000개의 점수 항목을 리스트로 생성하여 정렬하지 않습니다. 대신 실행 중인 top-k (running top-k)를 유지하므로, 할당(allocation)도 적고 산술 연산량도 더 적습니다.

인덱스의 메모리 비용

float32 형식의 20,000 x 384 인덱스는 29.3 MB입니다. 이 수치는 예산에 반영해야 합니다. 왜냐하면 이 인덱스는 기능이 유지되는 동안 메모리에 상주하며, 스마트폰 환경에서 29.3 MB는 결코 공짜가 아니기 때문입니다.

vector_kit에는 동일한 인덱스를 float32 크기의 4분의 1인 7.6 MB에 담을 수 있는 int8 QuantizedMatrix (양자화 행렬)가 있습니다.

final quant = QuantizedMatrix.from(index);
final hits = quant.topKCosine(query, 10);

양자화 (Quantization)는 손실이 발생하는 방식이므로, 문제는 랭킹 품질(ranking quality) 측면에서 어떤 비용이 발생하는가입니다. 저는 이를 가정하는 대신 float 랭킹과 비교하여 재현율 (recall)을 측정했습니다. 데모 데이터에서 int8 인덱스는 100%의 recall@10을 반환했습니다. 즉, 상위 10개 세트가 float 방식의 상위 10개와 정확히 일치했다는 의미입니다. 이 수치는 데이터에 따라 달라집니다. 임베딩 (embeddings) 간의 간격이 좁을수록 재현율을 일부 잃게 되므로, 제품을 출시하기 전에 본인의 코퍼스 (corpus)에서 동일한 측정을 수행하십시오. 이를 위한 도구는 패키지에 포함되어 있지만, 성능을 보장해 주는 것은 아닙니다.

사용하지 말아야 할 때

이 방식은 선형 스캔 (linear scan)입니다. 모든 쿼리에 대해 모든 행의 점수를 계산합니다. 이것이 결과가 정확한 이유이며, 행렬을 패킹하는 것 외에 구축하거나 튜닝할 인덱스가 없는 이유이기도 합니다. 또한, 이는 성능의 한계점(ceiling)을 설정하기도 합니다.

약 10만 개에서 100만 개 사이의 벡터까지는 스캐닝(scanning)이 괜찮으며, 위에서 언급한 수치들이 얻을 수 있는 결과입니다. 그 이상의 수천만 개 단위의 벡터로 넘어가면, 내부 루프(inner loop)가 아무리 타이트하더라도 선형 스캔(linear scan)은 잘못된 구조이며, 정확도를 희생하는 대신 서브리니어(sublinear) 쿼리 시간을 제공하는 근사 최근접 이웃 (Approximate Nearest Neighbor, ANN) 인덱스 (HNSW, IVF)를 사용해야 합니다. vector_kit은 그렇게 작동하지 않으며, 그렇게 작동하는 척하지도 않습니다. 그 정도 규모라면 실제 ANN 라이브러리나 벡터 데이터베이스를 사용하여 그에 따른 재현율(recall) 튜닝 비용을 지불해야 합니다.

또한, 이 라이브러리는 완전한 선형 대수(linear algebra) 패키지도 아닙니다. 패킹된 행렬(packed matrix)에 대해 내적(dot product), 코사인 유사도(cosine similarity), 그리고 top-k 연산을 수행합니다. 일반적인 행렬 곱셈(matrix multiply), 분해(decompositions), 또는 자동 미분(autograd)을 위한 도구는 아닙니다.

이 라이브러리가 목표로 하는 사례, 즉 장치의 메모리에 들어가는 코퍼스(corpus)에 대한 정확한 top-k 연산의 경우, 중요한 두 가지 수치는 쿼리당 1.4ms와 인덱스 크기로 29.3MB 또는 7.6MB 중 무엇을 선택하느냐입니다. 여러분의 데이터에 대해 가장 먼저 확인해봐야 할 수치는 바로 이것들입니다.