埋め込みと再ランキングの分離

Rust、Axum、ort、tokenizersで、embedとrerankを別のAPIとして設計しました。OpenAI互換のローカルLLM基盤で、検索処理を再利用するための構成です。LLM導入時のRAGにも使う想定です。

当初は文書のベクトルを事前計算し、LAN内でkey -> binキャッシュを共有する案でした。LiteLLM ProxyやCPU/GPU分業の構成に合わせ、ONNX/CPUを前提にしました。

CPU側に置く処理

CPUで前処理、embedding、タスク整形を行い、GPUで最終生成を行う案です。embeddingとrerankを、生成とは独立した共通処理に分けます。

ColBERTのrerankはtoken単位のベクトルを扱います。毎回送るとLANの帯域と待ち時間が増えるため、初期案ではkeyと最小限のメタデータを送り、文書ベクトルをキャッシュから取得します。

設計方針

仕様はembed(256d)とrerank(ColBERT 64d token, MaxSim)の2系統です。どちらもRustとONNX/CPUを使います。LiteLLMやTask Routerを含む全体構成のうち、検索と再ランキングをここで定義します。

文書側は更新時に事前計算し、query側はリクエスト時に計算します。keyでキャッシュを共有し、複数ワーカーから同じ文書ベクトルを再利用する案です。

仕様詳細

設計目標

embedは近傍検索用の単一ベクトルを返します。rerankはColBERTのlate interactionで候補を並べ直します。別々に調整できる構成です。

Embed API

POST /embedはtexts: string[]を受け取り、256次元ベクトルを返します。tokenizersでtokenizeし、ortで推論します。モデルに合わせてpoolingし、256次元にtruncateしてL2 normalizeします。

1024次元級の出力も256次元に揃え、索引とcosine/inner productの計算を統一します。ただし、モデルカードと異なるpoolingやnormalizeは品質に影響します。Python参照実装との数値一致を確認する前提です。

Rerank API

POST /rerankはquery: stringとcandidates: Candidate[]を受け取り、scores: float32[]とorder: int[]を返します。Candidateにはdoc_key、またはkeyを解決するdoc_id + seq + chunk_idを渡します。

query token vectors (Tq x 64)は生成し、document token vectors (Td x 64)はキャッシュから取得します。queryの各tokenについて文書tokenとの最大類似度を求め、合計するMaxSimです。候補本文を毎回送らずにColBERTの計算を行います。

モデルとファイル構成

embedにはlightonai/modernbert-embed-large/onnx/*のINT8版（onnx/model_int8.onnx、約17MB）を採用しました。検索精度にも寄与し、レイテンシは約10msでした。軽量なモデルで並列処理を行う構成です。

同じリポジトリのtokenizer.jsonを使います。special tokenとnormalizationを参照実装に揃え、検索結果への影響を避けます。

rerankはmixedbread-ai/mxbai-edge-colbert-v0-32mを想定しました。ONNXが含まれなければortで直接実行できないため、lightonai/mxbai-edge-colbert-v0-32m-onnxを取得する案です。

キャッシュ設計（LAN 内、key/bin 方式）

キャッシュ案では、doc_id + content_seqなど、次の材料から不可逆キーを作ります。

  doc_id|content_seq|chunk_id|model_id|tokenizer_hash|max_len|chunk_ver|dtype|layout

blake3でまとめ、doc:を付けます。content_seqは文書更新時に進め、内容が同じなら再利用します。chunking変更時はchunk_verを変えて無効化します。

最初はfp32（per-row float32 vectors）で検証し、その後にbitpack（sign encoding + u64 array）を検討します。

ヘッダはdtype, dim (=64), n_tokens, scales_present, layout (token-major), versionです。MaxSimとデコードで、token-majorの配置を共有します。

gRPC 設計（バイナリ vec 直接送信）

実装時にはDB依存をなくす方針に変更しました。doc_keysだけを送る案から、embed済み256d vecを直接送る方式に切り替えています。

gRPCのMaxSimSearchはmap<string, VecF32>でkeyと256dベクトルを受け取ります。256 floats × 4 bytes = 1024 bytes/vecで、約1KB/候補です。サーバーでthreshold、top_k、top_pを適用し、keysとscoresの並列配列を返します。

  message MaxSimRequest {
  map<string, VecF32> candidates = 1;  // key -> 256d vec (100-200 typical)
  VecF32 query_vec = 2;               // query vector (256d)
  float threshold = 3;                // min score to include
  int32 top_k = 4;                    // max results
  float top_p = 5;                    // cumulative-score cutoff
}

ベクトルを直接送る方式は、次の条件で採用しました。

  • ローカル LAN 内の通信に限定される
  • 1 候補あたり約 1KB、100-200 候補でも 100-200KB 程度に収まる
  • DB 依存がなくなることで、embed/rerank サービスが完全にステートレスになる
  • この精度で検索品質は十分という検証結果がある

この変更でキャッシュ層とPostgreSQLのcontent_seq管理が不要になり、サービスをステートレスにして並列実行しやすくしました。

HTTPの/v1/rerankはqueryテキストとcandidate keysを受け取り、内部でONNX推論します。gRPCは、embed済みベクトルを持つagent-gatewayなどがMaxSimスコアリングを依頼する経路です。

Rust 実装要件

RESTはAxumを使い、内部gRPCは必要に応じて追加します。tokenizersでtokenizer.jsonとspecial tokenを揃えます。ortはintra_threads = num_cpus、optimization_level = Level3、CPU execution providerで始めます。

MaxSimは内積を使います。後からbitpack + Hamming approximationへ切り替えて検証できるよう、traitで分ける想定です。

実装順序

初期案ではembedとpool -> truncate256 -> normalizeを先に固めます。次にrerankのONNXを確認・取得して/rerankを作り、PostgreSQLのcontent_seq、文書の事前計算とcache set、最小構成のgRPCを順に追加する予定でした。

APIと配列形式を先に確認し、その後に通信を追加する順序です。DBとキャッシュを使う部分は、前述の実装変更で不要になりました。

注意事項・未知数

初期案の確認点は、rerankのONNX配布、poolingとnormalizeの一致、キャッシュの無効化です。max_lenやstrideを変える場合は、キーのchunk_verも変える必要があります。

結果

embeddingとrerankの責務を分けました。独立サービスにも、OpenAI互換Proxyと同じプロセスにも配置できる境界です。

gRPCの直接送信に切り替え、keyだけを運ぶ初期案のキャッシュとDB版管理を外しました。

今後の作業

embedは約17MBのINT8で約10msでした。rerankもlightonai/mxbai-edge-colbert-v0-32m-onnxを採用し、Rustのortから実行できる状態です。

残る判断は、/embedと/rerankを公開APIにするか、OpenAI互換Proxyの内部APIにするかです。認証、監視、ワークフローとの接続が変わります。