Retrieve-for-Train、検索時の推論を5,390万パラメータの拡散検索器にコンパイル

Google Researchは9月15日、Retrieve-for-Train（R4T）を公開し、「検索結果を組み合わせて初めて良い結果になる」という集合検索の問題の解決を試みた。従来のランキングモデルは通常、関連性を項目ごとに評価する。しかし、コーディネート、プレイリスト、商品セットでは、元のクエリから逸脱せず、結果同士に多様性と相互補完性があることも求められる。大規模言語モデルに10個のサブクエリを直接展開させれば、こうした条件を扱えるものの、推論過程とテキストをトークン単位で順次生成する必要があり、検索欄に求められる低レイテンシーを満たすのは難しい。[公式技術記事](https://research.google/blog/bypassing-inference-bottlenecks-accelerating-complex-ai-search-with-retrieve-for-train/)では、R4Tを高コストな推論処理をオフライン学習へ移す3段階のパイプラインと説明している。

第1段階では、Soft-GRPOを用いてGemma 3 4BまたはQwen3 4Bをファインチューニングする。報酬では、クエリとデータベース間のgroundedness、元の意図とのalignment、集合のVendi Scoreによる多様性を同時に評価する。第2段階では、このfan-out言語モデルを凍結し、クエリごとに128組の学習ターゲットをサンプリングする。第3段階では、これらの合成ペアを使って5,390万パラメータのDiffusion Transformerを学習し、クエリembeddingから10個のターゲットembeddingを直接生成させ、最後に最近傍探索でデータベース内のコンテンツへマッピングする。これは、オンラインで推論モデルを呼び出す代わりに、RLが学習した集合戦略を小型の生成検索器へ「コンパイル」することに相当する。

[論文](https://arxiv.org/abs/2603.06397)では、Polyvoreの服飾データと産業用途の音楽プレイリストにおいて、単一クエリ、ゼロショット展開、Best-of-Nより高い品質を達成し、クエリ側の処理速度が自己回帰手法より12～20倍向上したと報告している。エンジニアリング上の注意点として、いわゆる1回の非自己回帰生成でも256ステップのSDE solverを使用しており、単一のネットワークforward passではない。実際のレイテンシーは、embeddingモデル、最近傍インデックス、集合サイズにも左右される。また、音楽データはプロプライエタリであり、オープンエンドな品質評価の一部にはLLM judgeが使われているため、ニュース、企業文書、変化の速い商品データベースでも同等の効果が得られるとはまだ断定できない。今後は、コードとデータが公開されるか、また新しい領域で報酬の重みを調整するために高コストな再学習が必要になるかを注視すべきだ。