RAG-Retrieval：埋め込み・ColBERT・リランカーを1つの訓練経路にまとめる

RAG-Retrievalが埋めるのは「検索モデルの学習コードが散らばる」問題

RAGの検索品質を上げようとすると、埋め込みモデル、ColBERT系のlate interactionモデル、リランカーという別々のモデルを、それぞれ別のリポジトリの別々の学習スクリプトで扱うことになりやすい。データ形式も損失関数も推論インターフェースも揃わない。RAG-Retrievalはこの3種類を1つのリポジトリにまとめ、訓練・推論・蒸留という3つの工程を同じ語彙で扱えるようにすることを狙っている。READMEの冒頭は「end-to-end code for training, inference, and distillation of the RAG retrieval model」と述べており、対象読者は自前の検索モデルを学習したいエンジニアである。bge、bce、gteといった公開モデルの名前が互換対象として挙がっている点は、ゼロからモデルを作るのではなく既存モデルを自前データで微調整する使い方を想定していることの表れだ。

3つのモデル種別をサブディレクトリで分ける設計

リポジトリ構成でまず目に入るのは、訓練コードがモデル種別ごとのサブディレクトリに分かれている点である。READMEは埋め込みモデルの例としてrag_retrieval/train/embeddingを挙げ、他の種別も同様の構成だと説明している。つまり1つの訓練スクリプトが3種類を分岐して処理するのではなく、種別ごとに独立したディレクトリとREADMEを持ち、それぞれの手順に従って実行する形になっている。これは見通しの良さと引き換えに、共通部分の変更が複数ディレクトリに波及しうる構造でもある。READMEのFeaturesには「Rejects complex, with a simple and understandable code structure for easy modifications」とあり、拡張しやすさを優先した結果だと読める。推論側は逆に、リランカーに焦点を当てた軽量なPythonライブラリrag-retrievalとして分離されている。訓練はリポジトリをcloneして使うが、推論はpipで入れる。この非対称性は、訓練環境と本番環境を分けたい利用者には都合がよい。

訓練の実際：conda環境、requirements.txt、train_embedding.sh

READMEが示す訓練用のセットアップは次のとおり。condaでPython 3.8の環境を作り、torchを手動で入れてからrequirements.txtを適用する。READMEはこの手順について、自動インストールされるtorchとローカルCUDAの非互換を避けるため、互換バージョンのtorchを先に入れることを勧めている。pip install -r requirements.txtの前にこの作業を挟むかどうかで、その後の詰まり方が変わる。訓練の起動は種別ごとのディレクトリに移動してシェルスクリプトを叩く形で、埋め込みの例はcd ./rag_retrieval/train/embeddingの後にbash train_embedding.shである。パラメータはこのシェルスクリプトと、その先の設定ファイル側で調整することになる。マルチGPUはdeepspeedとfsdpがFeaturesに挙がっているが、README本文には具体的な起動コマンドや設定キーの記載はない。ここは各サブディレクトリのREADMEを開かないと確定できない部分である。

推論はrag-retrievalパッケージに寄せる：BaseRerankerの継承

推論側のpip install rag-retrievalは、訓練用の依存とは別に用意された軽量ライブラリである。対応するのはCross Encoder系のリランカーとDecoder-Only LLM系のリランカーで、複数のモデルを同じインターフェースから呼べることを狙っている。長文書への対応として、最大長で切り捨てる方式と分割して最大スコアを取る方式の2つが用意されているとREADMEは説明している。拡張の入口も明示されていて、新しいランキングモデルを足したい場合はBaseRerankerを継承し、rankとcompute_scoreの2つの関数を実装する。この2関数という境界は狭く、既存の推論コードを大きく触らずにモデルを差し替えられる。逆に言えば、ランキング以外の処理、たとえば検索そのものやインデックス管理はこのライブラリの守備範囲ではない。詳細はexamples/Reranker_Tutorial.mdに委ねられている。

蒸留とMRL：大きいモデルを小さいモデルに落とす

RAG-Retrievalの特徴としてREADMEが繰り返し挙げるのが蒸留である。LLMベースのリランカーや埋め込みモデルを、0.5BパラメータのLLMやBERT-baseといった小さいモデルに蒸留できるとされている。推論コストを下げたいが、教師信号は大きいモデルから取りたいという場面を想定した機能だ。埋め込み側にはMRL（Matryoshka Representation Learning）の損失も実装されており、出力ベクトルの次元を削減できる。次元を落とせばベクトルストアのサイズと検索時の計算量が下がる。ここで注意したいのは、蒸留もMRLも「精度を保ったまま小さくする」ことを保証する仕組みではなく、精度とコストのトレードオフを調整する手段だという点である。READMEには蒸留後の品質を測った表は載っていない。効果は自分のデータセットで測るしかない。

READMEの実験表が語らないこと

READMEにはリランカーモデルのMTEB Rerankingタスクでの結果表がある。bge-reranker-base、bce-reranker-base_v1、そしてrag-retrieval-rerankerの3つが、T2Reranking、MMarcoReranking、CMedQAv1、CMedQAv2の4列で並んでいる。rag-retrieval-rerankerはモデルサイズ0.41GBと、他2つの1.11GBより小さい。ただし表のAvg列はrag-retrieval-rerankerの行では空欄で、4列すべての値も与えられていない。bge-reranker-baseのAvgは67.03、bce-reranker-base_v1は66.33と記載されているが、3つ目は比較できない。つまりこの表だけから「小さいモデルが同等以上」と結論することはできない。また、この表はリランカーの評価であり、埋め込みモデルやColBERTの検索品質は示されていない。MRLで次元を削ったときの劣化幅も、蒸留後の性能も、この表からは読めない。導入判断の材料としては不足している。

向かないケースと、代わりの選択肢

向かないケースははっきりしている。既製モデルをそのまま使う予定で、自前データでの微調整も蒸留も不要なら、このリポジトリの訓練部分を入れる理由はない。推論だけが目的ならrag-retrievalパッケージ単体で足り、conda環境もrequirements.txtも不要である。もう1つの境界は計算資源で、deepspeedやfsdpを使うマルチGPU訓練を前提とした構成であり、CPUのみの環境で動かす話はREADMEには出てこない。代わりの選択肢として挙げられるのは、埋め込みモデルの訓練に特化したフレームワーク（たとえばsentence-transformersの訓練インターフェース）である。違いは扱う範囲で、sentence-transformersは埋め込みとCross Encoderの訓練に集中し、ColBERT系のlate interactionを同じ枠組みで扱うことは主目的にしていない。RAG-Retrievalは3種別を1つのリポジトリに並べる代わりに、種別ごとのディレクトリとスクリプトを自分で把握する必要がある。どちらが良いかは、ColBERTを選択肢に入れるかどうかで分かれる。

ライセンスと保守コストの見積もり

ライセンスはMITで、リポジトリのLICENSEファイルへのリンクがREADMEに置かれている。MITは商用利用を含めて寛容な条件だが、これはリポジトリのコードに対するものであり、学習に使う公開モデルの重みやデータセットの条件は別に確認する必要がある。蒸留では蒸留元モデルの出力を教師信号として使うため、蒸留元のライセンスが自分の用途を許すかを先に確認したい。保守の観点では、最新のリリースがRAG-Retrieval v0.1（rag_retrieval_only_train、2024年5月4日）である一方、READMEのNewsには2025年5月のMyopic Trapや2024年12月のStella/Jasper関連の公開が並んでいる。バージョン番号の付いたリリースと、Newsで告知される追加成果が一致していない。pipで入れるrag-retrievalパッケージを本番に組み込むなら、PyPI側のバージョンとリポジトリの状態がどこまで対応するかを確認しておく必要がある。訓練側はcloneして使う以上、追従は自分のペースで決められる。

編集部の結論

採用を検討すべきなのは、bgeやbce、gteといった公開の埋め込み・リランカーを自前データで微調整し、その結果を蒸留で小型モデルに落とすところまでを1つのリポジトリで回したいチームである。逆に、既製のモデルをAPI経由で呼ぶだけで足りている場合や、訓練環境にGPUを用意できない場合は、推論ライブラリであるrag-retrievalだけをpipで入れて様子を見るほうが無駄が少ない。導入前に確認するのは3点で、requirements.txtのtorchが手元のCUDAと噛み合うか、rag_retrieval/train配下の各サブディレクトリのREADMEが想定するデータ形式に自分のコーパスを変換できるか、そして蒸留の蒸留元モデルのライセンスが商用利用を許すかである。READMEの実験表はリランカーのMTEB Reranking上の数値であり、埋め込み側の検索品質はそこからは読み取れない。

公式情報源

1. Issues
2. License: MIT
3. NovaSearch-Team/RAG-Retrieval on GitHub
4. README
5. Releases