ONNX Runtime Web + 蒸留 student でブラウザ完結 intent search を作る

検証日: 2026-05-23

使用バージョン: onnxruntime-web@1.21.0 / @huggingface/transformers@3.0.0 / sentence-transformers@3.2.1 (蒸留)

対象: ブラウザだけで「やりたいこと → 関連コンテンツ」型の semantic search を組みたい、 タグ + BM25 では届かない open-vocabulary なクエリを扱いたい、 サーバ推論コストを 0 にしたい

タグや事前 score では届かない「ユーザの意図 → コンテンツ」の翻訳をブラウザ完結でやる実装。 商用 embedding (Gemini / OpenAI 等) を教師に小型 student へ MSE 蒸留 → ONNX 化 → onnxruntime-web で推論、 という流れ。

触って試す #

蒸留 student は rakuten_books 2000 件 (IT 系) を Gemini embedding-2 で teacher 化 したものを mimic。 overlap@5 = 0.26 (random の 104 倍)、 検索 latency は warm で ~30ms。

何を解こうとしたか #

カタログ系サイトの「関連コンテンツ」UX は普通 タグ + BM25 で組む。 これで届かない領域:

• 「Excel が遅くて困ってる」→ Excel マクロ / VBA / Python 自動化本
• 「新人エンジニアに最初の 1 冊」→ 入門書群 (ジャンル横断)
• 「LLM を業務システムに組み込みたい」→ LLM 入門 + システム設計 + 業務自動化

ユーザは 意図 (やりたいこと) を投げるが、 書籍タグは 内容 (この本は何について) しか持たない。 意図 → 内容の翻訳が必要で、 これが semantic embedding 唯一の存在意義。

サーバに API 化して Gemini に毎回投げる選択肢もあるが、 ブラウザ完結なら:

• レイテンシ削減 (1st 加味で 30ms encode + 5ms rank vs API 数百 ms)
• プライバシー保護 (クエリがサーバを通らない)
• インフラコスト 0 (キャッシュ後は CDN のみ)
• オフライン耐性

代償は精度。 蒸留 student は商用 embedding の精度を完全には再現できない (本稿の検証では top-K の 1/4 程度が teacher と一致)。 用途は 広めの推薦 に限定する。

アーキテクチャ #

[user input]
   ↓
[tokenizer.json] ──→ input_ids, attention_mask
   ↓                       ↓
   └──→ [model_f32.onnx (450MB)] ──→ 768d normalized vector
                                            ↓
                                  cosine 内積 (INT8 corpus)
                                            ↓
                                   top-K [{i, score}]
                                            ↓
                                      [book metadata]

実装パターン #

1. PyTorch wrapper で 1 つの ONNX graph に焼く #

sentence-transformers の Transformer + Pool + Dense + Normalize を Python 側で 1 つの forward に統合し、 torch.onnx.export する。 ブラウザ側で pooling / dense / normalize を JS 実装するより断然楽 + 速い。

class StudentWrapper(torch.nn.Module):
    def __init__(self, st_model):
        super().__init__()
        self.transformer = st_model[0].auto_model
        self.dense_linear = st_model[2].linear  # 384 -> 768

    def forward(self, input_ids, attention_mask):
        hidden = self.transformer(
            input_ids=input_ids, attention_mask=attention_mask
        ).last_hidden_state
        mask = attention_mask.unsqueeze(-1).to(hidden.dtype)
        pooled = (hidden * mask).sum(1) / mask.sum(1).clamp(min=1e-9)
        projected = self.dense_linear(pooled)
        return torch.nn.functional.normalize(projected, p=2.0, dim=1)

torch.onnx.export(..., dynamo=False, opset_version=17) の legacy TorchScript exporter を指定。 dynamo ベース exporter は shape inference でハマるケースがある (Linear の入出力次元が違うと quantization で Inferred shape and existing shape differ in dimension 0: (384) vs (768) エラー)。

2. corpus は per-vector INT8 量子化 #

書誌ベクトル 2000 × 768 = 1.5MB を per-vector INT8 化:

abs_max = np.abs(vectors).max(axis=1, keepdims=True).clip(min=1e-9)
scales = (abs_max / 127.0).astype(np.float32).reshape(-1)
quantized = np.round(vectors / abs_max * 127.0).astype(np.int8)
# 量子化誤差 (cos 差): mean 0.00015, max 0.00024 — 無視可能

ブラウザでの cosine 内積:

// dot(q_f32, c_i8 * scale_i) = scale_i * sum(q_j * c_i8[j])
for (let i = 0; i < N; i++) {
  let inner = 0;
  const off = i * DIM;
  for (let j = 0; j < DIM; j++) {
    inner += queryVec[j] * corpusInt8[off + j];
  }
  scores[i] = { i, s: inner * corpusScales[i] };
}

corpus binary layout は [int8 N*D bytes][float32 N scales] を 1 ファイルに連結。 fetch 1 回で読める。

3. onnxruntime-web の prefix を encode と一致させる #

E5 系 backbone は学習時に "passage: " prefix を期待する。 学習・評価・ブラウザ encode の 3 箇所すべてで完全一致 させること。 一致しないと意味空間が変わって overlap が崩れる。

const inputs = await tokenizer(PREFIX + text, {  // PREFIX = "passage: "
  padding: true,
  truncation: true,
  max_length: 256,
});
const feeds = {
  input_ids: new ort.Tensor(
    "int64",
    BigInt64Array.from(inputs.input_ids.data, BigInt),
    inputs.input_ids.dims
  ),
  attention_mask: new ort.Tensor(
    "int64",
    BigInt64Array.from(inputs.attention_mask.data, BigInt),
    inputs.attention_mask.dims
  ),
};
const out = await session.run(feeds);
return out.sentence_embedding.data; // Float32Array(768)

E5 系の query 側は "query: " prefix が本来は asymmetric retrieval の作法だが、 蒸留 student は passage 側だけで学習しているので両側 passage で揃える。 (asymmetric 化したい場合は teacher embedding 生成時に RETRIEVAL_QUERY task type で別途学習が必要)

数値 (検証結果) #

gemini-embedding-distill skill を使った 4 段階の検証:

• 絶対 overlap@5 は N=500 まで停滞 (0.17 帯)、 N=2000 で初めて 0.26 に jump。 量が効くのは指数ではなくしきい値 (~1000 件) 付近で
• ratio は monotone 改善、 N=500 で skill 閾値 10x を突破
• N=2000 で unfreeze (--unfreeze-base) も試したが、 MPS OOM (e5-base) もしくは 60 分 (e5-small) で session 内未完。 future work

ブラウザ実測 #

WebGPU は INT8 quantized + 短系列ではほぼ恩恵なし。 むしろ kernel launch overhead で warm latency が悪化する。 fp32 + WASM で十分。

つまずいたところ #

• gemini-embedding-001 が Vertex AI で 404: 旧 SDK (vertexai.language_models) は新モデル gemini-embedding-2 を解決できない。 google-genai>=0.8 の embed_content 経由に切り替え。
• MPS OOM で e5-base + unfreeze + N=2000 が落ちる: 64GB shared memory でも足りず。 batch を 4 に下げる or gradient accumulation で対応、 もしくは e5-small で妥協。
• INT8 dynamic quantization が N=2000 student で精度劣化: query encoding が collapse (任意の query 間 cos > 0.96)、 全 query で同じ top-K に。 N=500 では問題なかったので、 weight magnitude 分布の違いと推測。 f32 (450MB) で運用、 production 化時は per-channel quantization か HF Hub 配信を検討。
• torch.onnx.export(dynamo=True) (新 default) で shape inference 失敗: Inferred shape and existing shape differ in dimension 0: (384) vs (768) → dynamo=False (legacy TorchScript) で回避。 quantization 互換性も legacy の方が高い。
• WebGPU で速度差が出ない: kernel ops の一部が CPU fallback、 短系列での kernel launch overhead が支配的。 INT8 + 256 token 程度なら WASM の方が安定。

Cloudflare Pages 25MB/file 制約への対応 #

model_f32.onnx は 450MB あるので Cloudflare Pages の 25MB/file 制限に引っかかる。 選択肢:

1. Hugging Face Hub に置く (推奨): 無料、 CDN 付き、 ML model の de facto。 URL を component の executionProviders に渡すだけ
2. Cloudflare R2 (有料、 $0.015/GB/月): 自前ホスト、 アクセス制御も可能
3. 静的 quantization (per-channel INT8 + calibration): 学習 corpus でキャリブレーションして INT8 化、 動的 dynamic より精度落ちにくい。 113MB → CF Pages の上限 25MB はまだ超えるが、 R2 でも安い
4. モデル分割: ONNX の external data 機能で複数ファイルに分け、 各 25MB 以下にする

本検証では (1) を本番化の前提に置き、 開発時のみ public/ に直置きしている。

関連書籍 #

まとめ #

• 蒸留 student + onnxruntime-web で「やりたいこと → 関連書」型の semantic search はブラウザ完結で動く
• 精度は N=2000 で overlap@5 = 0.26 / 104x random、 broad な関連性は捉える、 シリーズ判別等の precision は無理
• WASM で warm 28ms、 WebGPU は短系列 + INT8 では恩恵薄い
• INT8 dynamic quantization は student size により精度劣化、 f32 + HF Hub 配信が production 想定
• 用途は タグ + BM25 で届かない open-vocabulary なクエリ (意図 → 内容の翻訳) に限る