Rustでベクトルデータベースを自作して学んだこと

はじめに

最近、AIツールが急速に普及し、その多くがRAG（Retrieval Augmented Generation）をベースにしています。開発者として自分でもLLMを使って何か作りたいと思いつつ、embedding、RAG、アライメントといった概念がなかなか腹落ちしませんでした。

そこで、RAGの基盤であるベクトルデータベースを自作してみることにしました。本番運用を目的としたものではありませんが、動くものを作ることで、セマンティック検索の仕組みを実感レベルで理解できました。

本記事ではRust・embedding・ベクトル演算に関する概念が登場しますが、ベクトル演算の基礎知識があれば十分に読み進められます。

この記事では、実装の内容と、その過程で得た学びを共有します。

RAGとベクトルデータベースの関係

RAGは、LLMの上に「文脈提供レイヤー」を追加する仕組みです。ユーザーのクエリをそのままLLMに渡すのではなく、まず関連情報をデータベースから検索し、クエリと組み合わせてLLMに渡します。

この「関連情報の検索」に使われるのがセマンティック検索です。キーワードの一致ではなく、意味的に近い結果を返します。そのために、テキストをMLモデルで固定長の数値配列（ベクトル）に変換するembeddingが使われます。意味が近い文字列は、ベクトル空間上で近い位置にマッピングされます。

ベクトルデータベースは、このベクトル同士のマッチングを行う場所です。

ユーザークエリ
           │
           ▼
┌───────────────────────┐
│      Embedding        │
│  (クエリ → ベクトル)    │
└───────────┬───────────┘
            │
            ▼
┌───────────────────────┐
│  ベクトルデータベース    │  ← 今回自作した部分
└───────────┬───────────┘
            │
            ▼
┌───────────────────────┐
│  関連ドキュメントを      │
│  IDで取得              │
└───────────┬───────────┘
            │
            ▼
┌───────────────────────┐
│  クエリ + 取得した      │
│  コンテキストを結合      │
└───────────┬───────────┘
            │
            ▼
┌───────────────────────┐
│        LLM            │
│    (回答を生成)         │
└───────────┬───────────┘
            │
            ▼
        最終回答

実装内容

学習目的のプロジェクトなので、ANNインデックスや量子化（大規模データ向けの最適化技術）は入れず、シンプルな構成にしています。現状、100K件の786次元f32ベクトル（786要素で構成され、各要素が32ビット浮動小数点数の配列。言語モデルによく使われているスペック）を扱えます。

実装した機能は以下の通りです。

• Insert / Delete / Search
• Save / Load
• CLI

コード：

永続化やCLIは本プロジェクトの主題ではないため、この記事ではデータベース本体の実装、特に検索に焦点を当てます。

ベクトルデータベースの役割は、質問に対して最も意味的に近い答えを見つけることです。この検索は類似度をもとに行われます。本プロジェクトではシンプルな実装としてコサイン類似度を採用しています。さらに検索精度と効率を高めるために、L2正規化を導入しています。

L2正規化

すべてのベクトルは挿入時にL2正規化されます。各要素をベクトルの幾何学的長さで割る処理です。

L2_norm(x, y) = (x / sqrt(x² + y²), y / sqrt(x² + y²))

これには2つの利点があります。

1. 検索品質の向上: 正規化によりベクトルの大きさが統一され、純粋に方向（意味）の類似度で比較できます。正規化しないと、「大きい猫」が「虎」よりも「大きい犬」に高い類似度を持ってしまうような問題が起こる可能性があります。

2. 計算の効率化: L2正規化後、コサイン類似度の計算は単純なドット積に変換されます。CPUは乗算と加算の連続処理が得意なため、これは大きなメリットです。挿入時に正規化しておくことで、検索時のコストを下げています。

検索の流れ

クエリベクトル
     │
     ▼
┌──────────┐
│    L2    │
│  正規化   │
└────┬─────┘
     │
     ▼
┌──────────┐
│  線形     │
│  スキャン  │──▶ 各ベクトルとドット積を計算
└────┬─────┘
     │
     ▼
┌──────────┐
│  Top-K   │
│  選択     │──▶ 類似度スコア上位K件を保持
└────┬─────┘
     │
     ▼
   結果
 (ID + スコア)

クエリベクトルをL2正規化した後、保存済みの全ベクトルに対してスキャンを行い、ドット積を計算します。Top-K（Kはユーザー指定）をIDと類似度スコアと共に返します。

ブルートフォース（全件をしらみつぶしに調べる方法）ですが、小規模データでは十分実用的です。

パフォーマンス: 100ms → 35ms

最も学びが大きかったのはパフォーマンスチューニングです。

100K件の786次元ベクトルに対するTop-10検索で、初期実装では100ms以上かかっていました。最終的に35msまで改善できました（i7-13700KF, 64GB/2400MT/s）。

改善のポイントはデータレイアウトの設計でした。ベクトルデータをどうメモリ上に配置するかによって、CPUキャッシュラインの効率が大きく変わります。線形スキャンではデータに順次アクセスするため、キャッシュフレンドリーなレイアウトが直接的にパフォーマンスに影響します。

高レベルの設計判断がこれほど大きなパフォーマンス差を生むことを、身をもって体感しました。

Rustで学んだこと

Rustlingsの次に取り組んだ初めてのRustプロジェクトでした。

CLIや永続化の部分は外部クレートやOSとの結合が強く複雑だったため、Claude Codeを活用しました。一方、データベース本体とベクトル計算の部分はロジックが明確でほぼ独立していたので、ここに集中して自分で実装しました。スコープを明確に切ったことで、学習効率を最大化できたと感じています。

Rustを選んだことで、キャッシュラインやCPUパイプラインといった低レベルの知識を実践で活用する機会になりました。また、SIMD最適化の可能性も見えてきました。ASMを確認したところ、アンローリングとSIMDが十分に活用されていない箇所があるので、次のステップとして取り組む予定です。

今後の展望

SIMD最適化: ホットパスにSIMD命令を適用し、検索時間を20ms以下に改善することを目指しています。

セマンティック検索エンジンの構築: このベクトルデータベースをベースに、embedding処理を含むセマンティック検索エンジンを構築する予定です。レイヤーを一つずつ積み上げて作っていく過程は、理解を深める上で非常に効果的だと感じています。

まとめ

ベクトルデータベースの自作は、RAGの仕組みとRustの両方を学ぶ上で非常に良い方法でした。概念的に複雑に見えるものでも、実際に作ってみると実装はシンプルで、読むだけでは得られない確信が手に入ります。

質問やフィードバックがあれば、ぜひコメントをお願いします。