pgvectorで実装するバイナリ量子化 — 768次元を96バイトに圧縮する2段階検索

課題: fp32ベクトルはスケールしない

Marianの埋め込みはGemini embedding(gemini-embedding-001)の768次元です。fp32なら1ベクトル3,072バイト。チャンク数が50万本で約1.5GB、5億本なら1.5TBになり、「全ベクトルをメモリに置いて総当たり」はもちろん、ANNインデックスの構築・保持も重くなっていきます。

解決の方向性は量子化です。ただし量子化は精度を落とすので、「どこで精度を捨て、どこで取り戻すか」の設計が本体になります。Marianの答えは2段階検索です。

• 粗探索: 1bit量子化(96バイト)のHamming距離で候補200件を高速に集める
• リランク: fp32(またはint8)の正確な距離で上位だけ並べ直す

量子化の設計とは精度を捨てる場所の設計である。全件に正確な距離は要らない。正確さが要るのは最終上位の数十件だけ。

1bit量子化: 符号だけ残す

量子化関数は「各成分の符号が非負なら1、負なら0」というだけの素朴なものです。768次元 → 96バイト、圧縮率32倍。

// lib/marian-vector/quantize.ts
export function quantizeBinary(vector: readonly number[]): Uint8Array {
  const bytes = new Uint8Array(Math.ceil(vector.length / 8))
  for (let i = 0; i < vector.length; i++) {
    if (vector[i] >= 0) {
      bytes[i >> 3] |= 0x80 >> (i & 7) // MSB-first: bit 0 はbyte 0の最上位
    }
  }
  return bytes
}

なぜこれで検索が成立するのか。高次元埋め込みでは、2つのベクトルのコサイン類似度と「符号が一致する次元の割合」に強い相関があるためです(ランダム超平面LSHと同じ原理)。Hamming距離から類似度への変換は次の式です。

export function binarySimilarity(a: Uint8Array, b: Uint8Array, bits: number): number {
  const agree = bits - hammingDistance(a, b)
  return (2 * agree - bits) / bits // [-1, 1] に正規化(コサインと同レンジ)
}

Hamming距離自体は、バイトXOR→256エントリの事前計算POPCOUNTテーブル引き、の累積で計算します。1比較あたり96回のXORとテーブル参照だけなので、TypeScript実装でも数十万候補の走査が現実的です。

int8リランク: 精度の取り戻し方

バイナリだけだと上位の順序が粗いため、中間段としてint8の対称スカラー量子化も実装されています。ベクトルごとに最大絶対値からスケールを決め、[-127, 127]に丸めます。

export function quantizeInt8(vector: readonly number[]): Int8Vector {
  let maxAbs = 0
  for (const value of vector) maxAbs = Math.max(maxAbs, Math.abs(value))
  const scale = maxAbs / 127
  const codes = new Int8Array(vector.length)
  for (let i = 0; i < vector.length; i++) {
    codes[i] = Math.max(-127, Math.min(127, Math.round(vector[i] / scale)))
  }
  return { codes, scale }
}

export function dotInt8(a: Int8Vector, b: Int8Vector): number {
  let acc = 0
  for (let i = 0; i < a.codes.length; i++) acc += a.codes[i] * b.codes[i]
  return acc * a.scale * b.scale // 整数で累積し、最後にスケール積を掛ける
}

PostgreSQL側: bit(768)とHNSW(bit_hamming_ops)

サーバー側はpgvector拡張だけで完結します。バイナリ符号はbit(768)列に持ち、Hamming距離のHNSWインデックスを張ります。

alter table public.marian_file_embeddings
  add column if not exists embedding_binary bit(768);

create index if not exists marian_file_embeddings_binary_idx
  on public.marian_file_embeddings using hnsw (embedding_binary bit_hamming_ops);

2段階検索はRPC関数1つで実装します。CTEでHamming距離(<~>演算子)の上位coarse_count件を取り、その中をfp32コサイン(<=>)で並べ直してmatch_count件を返します。

create or replace function public.match_marian_files_2stage(
  query_embedding vector(768),
  query_binary bit(768),
  match_user uuid,
  match_count int default 8,
  coarse_count int default 200
) returns table (file_id uuid, relative_path text, content text, similarity float)
language sql stable as $$
  with coarse as (
    select e.file_id, e.relative_path, e.content, e.embedding
    from public.marian_file_embeddings e
    where e.user_id = match_user
      and e.embedding_binary is not null and e.embedding is not null
    order by e.embedding_binary <~> query_binary  -- Hamming距離(HNSW)
    limit greatest(coarse_count, match_count)
  )
  select c.file_id, c.relative_path, c.content,
         1 - (c.embedding <=> query_embedding) as similarity
  from coarse c
  order by c.embedding <=> query_embedding       -- fp32コサインでリランク
  limit match_count;
$$;

従来の単段検索(IVFFlat・lists=100)はフォールバックとして残してあります。バイナリ列が未整備のユーザーは自動的に単段へ落ちるため、移行期間中も検索は止まりません。

バックフィル: 再埋め込みゼロで移行する

既存データへのバイナリ列追加は、LLM APIを一切呼ばずにSQLだけで完了します。fp32埋め込みは既にDBにあるので、符号を見てビット文字列を組み立てるだけだからです。

update public.marian_file_embeddings
set embedding_binary = (
  select string_agg(case when v >= 0 then '1' else '0' end, '' order by ord)
  from unnest(embedding::real[]) with ordinality as t(v, ord)
)::bit(768)
where embedding is not null and embedding_binary is null;

where embedding_binary is nullにより冪等で、何度流しても安全です。「新しい表現を導入するとき、既存表現から純粋に導出できる形にしておく」と移行コストが消える、という好例です。

まとめ

バイナリ量子化の実装は、(1)符号1bit量子化で32倍圧縮、(2)POPCOUNTテーブルによる高速Hamming距離、(3)pgvectorのbit(768)+HNSWbit_hamming_opsで粗探索、(4)fp32コサインでリランク、(5)SQLだけの冪等バックフィル、という構成です。専用ベクトルDBを足さずに、PostgreSQLの中で大規模検索への道を確保しています。