オフラインで数百万のWikipediaページを「雰囲気」で検索する
(leebutterman.com)- サーバーなしでブラウザ内にオフラインのリアルタイム埋め込み検索を実装し、英語版Wikipediaの600万件の文書をローカルで対話的に探索できるようにする
- 384次元のfloat32埋め込みの元データは約9GBでブラウザアプリには収まらないため、all-minilm-l6-v2とページチャンクの平均埋め込みを基にサイズを削減
- Product Quantizationを8次元サブ空間単位で適用し、データセットを288MB程度まで縮小。96MBに約200万件の埋め込みを格納できる
- 検索では展開せずに圧縮ドメインで距離計算を行い、ONNXのGatherElements opcodeで変換結果を修正して計算を約4倍高速化
- 最新デバイスではMiniLMの推論と100k件の埋め込みの距離計算が短時間で終わるため、React UIが100〜300ms単位で結果を更新できる
ブラウザ内で動作するオフラインベクトル検索
- 目的は、複雑なサーバー検索インフラがなくても数百万のWikipediaページをブラウザでリアルタイム検索できることを示すこと
- 中核となる設計は次の3点に集約される
- ベクトル検索はProduct Quantizationと線形スキャンだけでも十分高速になり得る
- ONNX Runtimeを使いつつ、PyTorch変換で生じた非効率はopcodeレベルで修正できる
- ブラウザはWASMでリアルタイム推論を実行できるほど高速であり、WebGPUにも今後の可能性が残っている
- 検索条件には、本文埋め込みのベクトル類似度とタイトル先頭文字の等価条件を併用する
- 一部のカラムはベクトル類似度で、別のカラムは等価条件で探すデータベースクエリの例に近い
- オフラインのブラウザアプリという制約があるため、埋め込みは小さく、埋め込みモデルはオープンソースで、モデルサイズと計算量も軽量である必要がある
- 2023年時点で軽量モデルとは通常、100MB未満を意味する
英語版Wikipediaデータと埋め込みサイズ
- データセットは英語版Wikipediaで、all-minilm-l6-v2モデルで埋め込む
- 文書は約600万件で、ページをチャンクに分割したうえで各チャンクの埋め込みを平均し、ページ埋め込みを作る
- ページは長いものから先に配置する
- データベースを段階的にロードするとき、人がより多くの労力をかけた可能性のある長いページを先に表示するため
- 埋め込み次元は384次元
- float32で保存すると
6M * 384 * 4 = 9GB - 96MBには約64k件の埋め込みしか保存できない
- float32で保存すると
- 元のfloat32エンコーディングは、ブラウザベースのオフライン検索には大きすぎる
Product Quantizationによる埋め込み圧縮
- Product Quantizationは浮動小数点値を直接保存せず、複数のパレットを作って入力値をパレットのインデックスに置き換える方式
- パレットは通常256個以下の値を持つため、各インデックスは最大1バイトにできる
- 暗黙的なパレットは、-127〜127や0〜255のような数直線を移動・スケーリングして使う
- 明示的なパレットは、256個の数値を順番に保存する
- Product Quantizationは明示的なパレット化を使う
- 最小構成として384次元それぞれを1バイトのインデックスに置き換えると、データセットは
6M * 384 = 2.25GBになる- 96MBには約256k件の埋め込みを保存できる
- この方式でも目標に対して依然として非効率
- 2次元の点単位で量子化すると、384次元の埋め込みを192個のインデックスで表現でき、2倍の省スペース化が可能
- 実際の構成では8次元の点単位を使う
- 384次元を48個のインデックスで表現する
- パレットサイズは約384KB程度に保たれる
- データセットサイズは
6M * 48 = 288M - 96MBに約200万件の埋め込みを保存できる
- どの量子化レベルでもprecisionとrecallを評価できる
展開せずに直接検索
- Product Quantizationの利点は、解凍せずに圧縮された状態で検索できること
- n次元をn/2個の2次元点にまとめると、各パレット点とクエリの該当する2次元点との距離を事前に計算できる
- その後、各埋め込みはパレットインデックスで距離値を参照し、それらを足し合わせて全体距離を計算する
- 全n次元の点を復元してから各次元の距離を計算する方式より作業量が少ない
- 実際の8次元点構成では、非圧縮の埋め込みを扱う場合の1/8程度の作業量で距離計算が可能
データ交換形式としてArrowを使用
- データの移動と変換のコストが大きいため、圧縮埋め込みはすぐ使える形式でシリアライズする方が有利
- Arrowはこの用途に適したカラム指向形式
- 埋め込みとページタイトルをArrowテーブルとして保存する
- パレットインデックスは2次元配列のように、タイトルは1次元文字列配列のように扱う
- パースやロードよりも、ビットを所定の位置へコピーすることに集中できる
- Arrowの配列形式は1次元データしか保存しないため、埋め込み48次元を扱うために2つのスキーマを使う
- メタデータスキーマは10万行単位
- 埋め込みスキーマは
10万 * 48行単位 - ロード時に埋め込みを再度reshapeする
- safetensorsも同じ設計原則に従う
- JSONは48要素の配列を可変長ASCII文字としてシリアライズするため、ロード段階が複雑になる
- Protocol Buffersは32ビット未満の整数にbase 128可変幅整数を使い、この整数形式は現在のcompute kernelでは十分にサポートされていない
Wikipediaのパースと埋め込み生成
- Wikipediaマークアップのパースにはmediawiki parser from hell、つまりmwparserfromhellを使う
- OLMのWikipedia datasetは、特定言語のWikipediaの最新データダンプを取得し、タイトルや本文などの行にパースする
- このデータセットは無害なシリアライズ済みデータというより、実行する必要のあるコードに近く、実行に対する信頼が必要
- パースは1台のマシンの全コアで行う
- 最新のマシンは数十コアを備えており、現在の600万ページ規模の英語版Wikipediaのような1億文書未満の体系に適している
- 埋め込みモデルにはall-minilm-l6-v2を使う
- 英語向けに性能が高く軽量なsentence transformer
- クエリと文書を同じ空間に埋め込む
- 2,200万パラメータの非常に軽量なモデル
- ブラウザJavaScript実行のために、ONNXと
transformers.jsによる8ビット量子化版も活用する - モデルは128トークンのシーケンスで学習されており、平均的なページ長よりコンテキストウィンドウがはるかに短い
- ページをチャンクに分割し、各チャンクの埋め込みを平均してページ平均埋め込みを生成する
pq.jsと線形スキャンベースのファセット検索
- 数百万文書をローカルで検索する場合、複雑なインデックス構造が必ずしも必要ではない
- 目標は特定の点に近い上位いくつか、たとえば距離の上位12件を得ること
- 距離配列1,000万件とファセットカラム1,000万件、ファセット値1つがある場合、条件に合えば0、合わなければInfinityを加えたうえでtop-kを探す方式でフィルタリングする
- 最新のスマートフォンで上位100件を10ms未満で見つけられる
- 実装はfiltered-topkと
pq.jsにある
- Product Quantizationで圧縮した埋め込みの距離計算はPyTorchで書ける
subspaceCount個のパレットがあり、各パレットにはcodewordCount個のsubspaceDim次元点がある
- PyTorchからエクスポートしたONNXモデルのインデックス処理はやや不自然な形になる
- ONNXのGatherElements opcodeが必要な処理を直接実行する
- ONNX-modifierのようなツールで、エクスポート済みONNXモデルのデータフローグラフのノードを追加・削除できる
- 複数段階のインデックス処理を正しい単一opcodeに置き換えると、距離計算が約4倍速くなる
ストリーミング計算と対話型UI
- 距離計算対象の埋め込み数は固定されていない
- クエリと一部の埋め込みとの距離をストリーミングで計算できる
- 十分な時間が経過したら、それまでに計算した距離でtop-kを実行し、検索結果を更新する
- すべての埋め込みがローカルにあるため、検索インフラまでの遅延は実質的に0ms
- 目標とするUIの応答性は、操作後100〜300ms以内に結果が見えるレベル
- MiniLMは100ms未満で実行可能
- エッジデバイスによっては15ms未満の場合もある
- 100k件の埋め込みの距離計算は約10msで実行可能
- データセットを長い記事から短い記事の順に並べると、最終的な上位検索結果の大半が初期ストリーミング結果に早く現れる
- 長い記事は作成時により多くの関心と労力が注がれている場合が多く、多数のstub記事より検索クエリに合いやすいことが多いため
- グローバルに距離を計算した後でファセットフィルタとtop-kを適用するため、同じクエリでファセット値や検索結果数を変える場合はフィルタリングだけを再実行すればよい
- このフィルタリングは10ms未満なので即時に感じられる
再利用可能なpq.jsコンポーネント
- Wikipedia search app全体の多くのライブラリ関数は、再利用可能なpq.jsコンポーネントへ移せる
- 現在は多くのONNX shapeが事前に固定されている
- さまざまな量子化レベルとさまざまな埋め込み次元をサポートすれば、より広範な再利用が可能になる
1件のコメント
Hacker News のコメント
確かに興味深いが、哲学・心理学の用語をいくつか説明してみたところ、探していた項目はいずれも20位前後にしか出てこなかった
より有名だが精度の低い項目が上位に並び、たとえば特定の心理療法の様式を定義しようとして何を入力しても、「psychotherapy」が常に1位だった
逆に、名前を思い出せない狭い下位分野を ChatGPT で探したことがあるが、毎回当ててくれた
説明から物の名前を探してくれる AI サービスというアイデアは良いが、Wikipedia や Wikipedia の記事タイトルに限定するやり方が正しいのかは分からないし、汎用の大規模言語モデルがすでにかなりうまくやっているように見える
それでも概念実証であり、ブラウザ上でローカルに動かせる点は本当にすばらしい
調整できるパラメータは多そうだ。記事の最初の段落だけを使うのか全文を使うのか、特定の記事に近い範囲内で検索するのか、といったものだが、まだ手を入れていない
Wikipedia はデモ用データセットとして優れており、ほかのデータセットも追加してみたい。たとえば iPhoto で「mountain」を検索すると山が写った写真が出てくるように、CLIP のようなマルチモーダルモデルで複数のデータセットを検索すると面白そうだ
大規模言語モデルは最高のファジー検索エンジンのように見え、従来の検索エンジンとはかなり独特だが相補的な形で動作する
コンセプトは気に入ったが、結果はうまく出なかった
「weird looking monkey」と入力して、テングザルやキンシコウのような結果を期待したが、「Pet monkey」「List of individual monkeys」「Ethnoprimatology」「Monkey」といった記事しか出なかった
同じクエリを Google に入れるとまさに期待した結果が出るので残念だったし、知らなかった奇妙な見た目のサルを探してみたかった
デモの焦点は埋め込みデータベースを見せることにあったが、埋め込み自体も多少は役に立つ
ページ上で何が見つかり、何が見つからないかについての分析データをまったく残していないので、検索結果を改善する準備はできていない
そのため、この知識コーパスから良い結果を得るのが難しくなる
実装は本当に良く、オフラインでこれができる点はすばらしい。ただし埋め込みの品質はまだ不足しているように見える
役に立つかもしれないコツの一つは、Wikipedia 記事全体ではなく、定義に近い文や、通常は最初の文・最初の段落だけを埋め込むことだ。今どの部分を使っているのかはよく分からない
私のサイト OneLook も、2003年から https://onelook.com/thesaurus/ で、説明から単語や概念を探す似た機能を提供してきた
当初は純粋な逆引き辞書検索だったが、この20年で単語埋め込み、文埋め込み、最近では大規模言語モデルまで試しており、最近は自前では答えられない入力に対して GPT が候補を生成している
この作業では大規模言語モデルが以前の方式よりはるかに優れているため、OneLook のこの部分を改善しようという意欲が少し下がった。逆定義検索が ChatGPT を使う主な理由だという人をよく見かける
少し遅れて見たが、テキスト埋め込みは、少なくともこの記事で使われているものに関しては、通常雰囲気で検索するのにはあまり向いていない
概して、重複する単語を比較したり、クエリと似た内容を探したりする方向に近い
ただし、実際にこの問題に取り組んだ最近の論文がある: “Retrieving Texts based on Abstract Descriptions” (Ravfogel et al., 2023) https://arxiv.org/abs/2305.12517
論文には「建物を設計する建築家」「別の会社の一部である会社」「ジャンルの発展に影響を与えた本」といった抽象的な説明で検索する例が多く載っている
その埋め込みはこうした検索をずっとよく支援できるようなので、リンク先記事のオフライン Wikipedia 検索をこの新しい種類の埋め込みでやり直してみると興味深そうだ
今のところ、このページは私の環境では動作しておらず、
model_quantized.onnxが読み込まれていない入力している間にも 19.2MB までダウンロードされ、速度は約 50KB/s だった。訪問者ごとにこれが発生するなら、Lee Butterman の帯域料金にひどいことをしているかもしれない
やっていること自体は非常に印象的だが、検索結果の品質は良さそうに見えない
経験上、検索結果の品質を手で評価するのが本当に難しいことは分かっている。非常に良い結果にかなり近づいていながら、それよりはるかに悪いマッチを返すこともあり得る
より最近の文埋め込みを使えば結果は良くなりそうだし、もっとデータを集める必要がありそうだ
技術はとても印象的だが、結果はそうではなかった。
「pointy building in Paris」で検索したところ、Tourism in Paris、Bourse de commerce (Paris)、Grands Projets of François Mitterrand、List of tallest buildings and structures in the Paris region、List of tourist attractions in Paris、Palais des congrès de Paris、Landmarks in Paris、Palais de la Bourse, Lyon、Outline of Paris、Architecture of Paris が出てきた。
パリで最も有名な尖った建物はまったく見当たらなかった。
こうしたアプリケーションでは、文書全体の文埋め込みが最善ではないのかもしれない。
たった今文書を確認したが、「building」という単語は19回出てくるものの、ほとんどは動詞で、その次は「Chrysler Building」だった。
私が思い浮かべられていない別の有名な尖った建物があるのでなければ、という意味だ。
検索エンジンの魔法の一部は、そのページへリンクしているページの埋め込みや、従来型の情報検索キーワードを、クリック数や権威スコアで重み付けして混ぜ合わせるところにある。
このシグナルがないと、有用な情報の多くが無視され、結果が魔法のようには感じられない。
それでも印象的で興味深いデモだ。
好きになりたかったが、自分が試した検索では関連する結果がほとんど出なかった。
「The wizard in The Lord of the Rings」では Gandalf も Saruman もなく、LOTR 関連の本だけが出てきた。
「Protagonist of Scorsese's Taxi Driver」では Travis Bickle が出なかった。
「A person that plants trees for a living」では、なぜか gardener がリストになかった。
「Curly-haired painter on TV」では Bob Ross がまったく出なかった。
「Unusually shaped modern art museum in Spain」では Bilbao が4位に出てはくるが、残りは変わった形ではなかった。
「Dog shaped like a sausage」なら dachshund が上位結果にあるべきだと思う。
記事がなかったのなら漏れていてもまだ不思議ではないが、実際にはすべて存在する。
「Vibes」は「sentence embeddings」よりずっとしっくり来る表現だ。自分もその表現を使い始めるべきかも :)
元記事の著者はなぜその単語を選んだのか説明しておらず、私の知る「vibe」のどの用法にも合わない。
「gist」では十分に流行語っぽくなかったのだろうか、と思う。