CRINN：強化学習による近似最近傍探索アルゴリズムの自動最適化
3つの要点
✔️ CRINNは強化学習を用いてANNSコードを自動最適化し、人手調整を不要にした
✔️ QPSと再現率のトレードオフを考慮し、速度を報酬信号として性能改善を実現
✔️ 6つのベンチマークで評価し、MNISTやGISTで大幅な速度向上を達成
CRINN: Contrastive Reinforcement Learning for Approximate Nearest Neighbor Search
written by Xiaoya Li, Xiaofei Sun, Albert Wang, Chris Shum, Jiwei Li
(Submitted on 4 Aug 2025 (v1), last revised 20 Aug 2025 (this version, v2))
Comments: Preprint Version
Subjects: Machine Learning (cs.LG); Artificial Intelligence (cs.AI); Computation and Language (cs.CL); Databases (cs.DB)
概要
本論文は、高次元ベクトル空間における近似最近傍探索（ANNS）の最適化をテーマとしています。
ANNSは、検索精度をわずかに犠牲にする代わりに検索速度を大幅に向上させる技術であり、近年ではRAG（Retrieval-Augmented Generation）やエージェント型LLM応用の基盤技術として不可欠な存在です。
従来の最適化は人間の専門家がプロファイリングを行い、キャッシュミスの分析やデータ構造の調整、パラメータの手動チューニングを繰り返すことで実現されてきました。
しかしこの方法は専門性と労力を要し、ハードウェアや応用環境の進化に追従するには限界があります。
そこで著者らは、LLMと強化学習を組み合わせた新しい最適化フレームワーク CRINN を提案。
CRINNは、コード実装の実行速度を報酬として扱い、対比学習に基づく強化学習で効率的なANNSコードを自動生成します。
これにより、人手による調整に依存せず、逐次的に改良された実装を生み出すことが可能となり、検索性能の新たなブレークスルーを実現しました。
提案手法
CRINNは、ANNSを強化学習による最適化問題として捉え、対比学習（contrastive learning）を組み合わせることで性能向上を図ります。
具体的には、既存の実装コードとその実行速度をプロンプトに組み込み、LLMに「なぜある実装が速いのか」を比較分析させる設計が特徴です。
これにより、モデルは速度向上につながるパターンを学習し、より優れた新しいコードを生成。
生成されたコードは実行され、速度と再現性を基に報酬が付与されます。
この報酬を用いてGroup Relative Policy Optimization（GRPO）に基づく強化学習を行い、モデルを逐次的に更新。
また、報酬設計ではQPS（Queries Per Second）と再現率（Recall）のトレードオフに注目し、[0.85,0.95]のRecall範囲における曲線下面積をスカラー報酬としました。
さらに、GLASSと呼ばれる既存のANNSライブラリを初期基盤とし、グラフ構築・探索・精緻化というモジュールごとに逐次最適化を実施。
この構造化された手法により、従来の専門家依存の調整を自動化し、効率的な探索アルゴリズムの開発を可能にしました。
実験
実験では、SIFT-128、GIST-960、MNIST-784、GloVe-25、GloVe-100、NYTimes-256の6つのベンチマークデータセットを用いて、CRINNの性能を検証。
比較対象として、ParlayANN、GLASS、NNDescent、PyNNDescent、Vearch、Voyagerといった代表的なオープンソースANNS実装を選定しました。学習はSIFT-128（ユークリッド距離）のみを用いて行い、その後他のデータセットに対して汎化性能を評価しています。
結果として、CRINNはMNIST-784やGIST-960などで最大85%の速度向上を達成し、特にグラフ構築モジュールで大幅な改善が確認されました。
一方でNYTimes-256のような一部データセットでは性能低下も見られ、距離尺度やデータ特性によって最適化が限定される可能性を示唆。
加えて、段階的なモジュール最適化の有効性が実証され、基盤となるGLASSから継続的な改善が可能であることが確認されました。
総じて、CRINNは既存手法を凌駕する速度と精度を両立し、強化学習を用いたコード最適化の新しい方向性を提示しています。
この記事に関するカテゴリー
