Fastly Compute と AI (1) HNSW を使ったベクトル検索
この記事は Fastly Compute 一人アドベントカレンダー 2024 23 日目の記事です。本稿では Fastly Compute と AI と題して、Compute を使った AI ワークロードのハンドリングについて書いてみたいと思います。第一回目の本稿では Compute 上でベクトル検索の実装例を見ていきます。
Compute と機械学習/LLM
以前から Fastly Compute では機械学習を利用した例として次のデモが公開されていました。ソース のコミットログを見ると 2021~2022 年頃に開発・公開されたもののようです。
上記デモサイトで実際に画像をクリックしてみると分かるのですが、Compute インスタンス上で実際に軽量なモデル（MobileNetV2）を使って画像認識を行うデモになっています。エッジコンピューティングのようなリソースが限られた制約の多い環境下においてどのように推論を行うことができるかを示した例（Proof of Concept）として現在でも参考になる点も多い一方で、同じ要領で近年目覚ましい進化を続けている LLM を使った推論や、その際に必要になるテキストや画像を Embedding モデルを使ってベクトル化するといったタスクを実行できるかどうかについて気になる人もいるのではないかと思います。
著者が調べた限りでは、本稿を執筆している 2024 年 12 月時点ではリソースの制約の多い環境向けの軽量と呼ばれるような LLM や Embedding のモデルであっても、実用に耐えるレベルのモデルの大きさや実行に求められるメモリのオーダーが小さいものでも GB レベルであり、Fastly Compute でプログラムを走らせる際の制約条件である 128MB メモリ（ヒープ領域）や 100 MB までのパッケージサイズなどの前提条件を満たすことは現状だと難しいように見えます。
対話型 UI などで応答速度が求められるシーンでエッジを活用
それでは Compute のようなエッジ環境は LLM や Embedding などの AI のワークロードに関してどのようにその特徴を引き出して活用すると良いのでしょうか。私が観測している範囲では、近年加速度的に LLM や Embedding のユースケースが増えてきたことで、それらにまつわる関連技術領域の裾野が広がっており、特に次のような推論やベクトル化の前工程や後工程を担当する処理基盤として特に活用できるシーンが増えているように思います。
- AI API Gateway
- 簡単に API Key を秘匿化しつつ限定ユーザに対してサービス提供が可能
- REST だけでなく、WebSocket のパススルー接続などが可能。応答速度を重視する場合ストリーミングでのレスポンス送信なども標準的に対応
- WebSocket の具体的な例についてはこちらのブログ参照
- AI Accelerator のような AI に特化したキャッシュ基盤の連携
- ベクトル検索の実行基盤
- 特徴ベクトル自体の生成は Embedding モデルを用いてエッジ外で事前に準備
- エッジでは近似最近傍探索（ANN）等のアルゴリズムを用いてベクトル検索を実行
前者については Gateway というコンセプト自体それなりに認知されているかと思うのでどちらかというとイメージがしやすいのではないかと思います。一方で後者のベクトル検索については、本稿を執筆している 2024 年 12 月時点で、Fastly では例えば KV Store などのデータストアによって標準的にベクトル検索のメソッドが提供されているわけでもないので、どのように実現するのか疑問に思われる方もいるかもしれません。そこで後者のベクトル検索の例として、本稿では以下のブログで紹介されているデモとそのソースコードを参照しながら、実際にどのようにベクトル検索の実行エンジンを Compute 上に構築できるかについてサンプルを紹介していきます。
HNSW を使ったベクトル検索
今回取り上げる例では、HNSW (Hierarchical Navigable Small World) と呼ばれる近似最近傍探索のアルコリズムを使いベクトル検索を自前で実装しています。検索の対象はニューヨークのメトロポリタン美術館の 50 万点近い（正確には N = 484,956 個の）作品に関する概要（説明文）のテキストデータで、ユーザが直近で訪れた美術品に類似している別の美術品のページをリアルタイムで推薦するのが今回のベクトル検索の目的です。
ここでの技術的なチャレンジの一つとして、仮に事前に説明文テキストのベクトル化や HNSW インデックスの生成をローカル環境等エッジ以外の環境で済ませたとしても、肝心のリアルタイムのベクトル検索処理をナイーブに実装するとエッジ上で 50 万個近い数百次元のベクトルを扱うことになってしまい、とうてい 128MB のエッジ環境でのメモリには乗らなくなってしまう点が挙げられます。そこで本デモでは下記 2 種類の方法で計算量とメモリ使用量両方の削減を行い、リソース制限の制約を回避しています。
方法(1) PCA による次元削減
今回のデモではローカル環境で軽量な sentence-transformers/all-MiniLM-L12-v2 モデルを使って説明文テキストをベクトル化しており、これにより 384 次元のベクトルが N = 484,956 個生成されます。こうして得られたベクトルを主成分分析(PCA)により 5 次元に次元削減することで処理時間とメモリ使用量を削減します。下記 json は削減後のベクトルのイメージです。
{
"id": 10826,
"vector": [-0.26523229479789734, -0.0947713553905487, -0.1279277801513672, 0.013157757930457592, -0.045752234756946564]
}
方法(2) k-means 法によるクラスタリング
仮に上記 PCA によってベクトルを 5 次元まで圧縮したとしても、現状のままでは 50 万件近いベクトル数の HNSW インデックスはまだメモリに乗らずリソース制限に引っかかってしまうため、更に踏み込んだメモリ使用量と処理時間の短縮を行う必要があります。そこで本デモでは k-means 法によるクラスタリングを行います。具体的には、次元削減後のベクトル約 50 万個を 500 個のクラスタに分割します。実際にクラスタに分割した際のベクトル数の分布の例は以下の通りです(クラスタ当たり158~3184個のベクトルを保持、平均960個/クラスタ)。
このとき、後で近傍探索が実施できるようにクラスタ内の全ての点の平均値は事前に算出しておきます(centroid)。ここで作られた 500 個のクラスタはローカル環境でそれぞれ個別の HNSW グラフとしてコンパイルされて、最終的に Compute サービス上から呼び出せるように KV Store にそれぞれの centroid をキーとする形でコンパイル結果をバイナリ形式で保存しておきます。また、KV Store が標準でベクトル検索に対応していないため、その点を補うために KV Store の各エントリに先ほど割り振った 500 個の centroid の集合に対しても同様に HNSW グラフを構築し事前コンパイルしておきます[1]。この 2 種類の階層化された HNSW グラフを用意し検索に利用することで、全ベクトル数に対する HNSW インデックスをメモリに読み込むことを回避し、効率的に検索することを意図しているようです。
実装のコード例
Compute サービス上で実際にベクトル検索している部分についても見てみましょう。実装の箇所としては recommender/src/recommender_kv.rs の下記コードとなります。本デモでは instant_distance クレートで提供される Search や MapItem 構造体を使いながら検索対象として受け付けたベクトルの近傍探索を実装しています。
...
// Find nearest cluster.
println!("🔍 Searching for nearest cluster:");
let start = time::Instant::now();
let mut search = Search::default();
let MapItem {
value: cluster_id, ..
} = vec_db
.centroid_map
.search(&median_v, &mut search)
.next()
.unwrap();
println!("\tFound cluster {} in {:?}", cluster_id, start.elapsed());
...
(中略)
...
// Find nearest neighbours within cluster.
println!("🔍 Searching for approximate nearest neighbour(s):");
let start = time::Instant::now();
let mut search = Search::default();
let nearest_neighbors: Vec<u32> = cluster
.search(&median_v.into(), &mut search)
.skip(offset)
.take(recs) // Limit the results to MAX_RECS
.map(|MapItem { value: obj_id, .. }| *obj_id)
.collect();
println!("\tNearest neighbours found in {:?}", start.elapsed());
...
備考
本稿を執筆するにあたって今回紹介したリポジトリの手順に沿ってローカル環境で 50 万個近いベクトルに対して HNSW インデックスの生成を行ってみましたが、軽量なモデルのおかげで GPU を搭載していない非力な Ubuntu (Core i7-1355U、Mem 8GB割り当て)での実行だったので 50 万個のベクトル化に少し時間はかかりましたが（1時間強程度で完了）、GPU 搭載環境であれば十分短い時間でベクトル化が完了できると思われます。HNSW インデックスの生成や、グラフの更新時の現実的な処理時間なども含めると現実的な更新を伴うワークフローまでは今回の記事執筆段階では精査できていないため、また別の機会があればその部分についても考察を追加したいところです。
また、検索対象のベクトル数が少なければ今回のように HNSW グラフを 2 段階構成にせずに 1 段階のグラフのみで済ませたり、あるいは数に応じてクラスタ数（500）や PCA で削減した次元数（5）を変更していくことで、エッジで実行できるパッケージサイズやメモリのサイズの制約の中で最大限のパフォーマンスを発揮できるポイントを探していくこともできそうです。
最後に、本デモも含めてエッジ環境での AI の活用について紹介された動画やその文字起こしもありますので、本デモに更に興味がある方はぜひこちらも参照してみてください。
まとめ
本稿では Compute を使った AI ワークロードのハンドリングの例として、Compute 上で HNSW インデックスを用いたベクトル検索エンジンをどのように構築できるかについて、具体的なデモと例を通して実装イメージについて考えてみました。本デモはたまたま「ニューヨークのメトロポリタン美術館の 50 万点の説明文テキストを用いて、ユーザが訪れた美術品に類似している別の美術品を推薦する」というテーマで実装されていますが、このテーマの汎用性はさておき「テキストや画像をベクトル化しておき、エッジ環境上でベクトル検索を行う」というアイデア自体は広く汎用的に応用ができそうです。
本年のアドベントカレンダー最終回となる次回は Wizer を使った WebAssembly モジュールのローディング速度高速化の tips について書いてみたいと思います。
Discussion