💬 NLP

DiffRetriever: Parallel Representative Tokens for Retrieval with Diffusion Language Models

DiffRetriever は、拡散言語モデル向けに、プロンプトに複数のマスクされた位置を付加し、単一の双方向パスで読み取る並列代表トークン検索手法を導入し、逐次自己回帰型マルチトークン手法および既存の対照微調整型検索器を上回る性能と効率を達成する。

原著者： Shuai Wang, Yin Yu, Shengyao Zhuang, Bevan Koopman, Guido Zuccon

公開日 2026-05-11

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Shuai Wang, Yin Yu, Shengyao Zhuang, Bevan Koopman, Guido Zuccon

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

巨大な図書館で、ナプキンに書いたたった一文に基づいて完璧な本を見つけようとしていると想像してください。これがコンピュータの「リトリーバー」が行うことです。あなたのクエリを受け取り、最も適合するドキュメントを見つけ出します。

長らく、これを行う最良の方法は、あなたの文章を読み、その全体のアイデアを表すたった一つの単語を書き出す「賢い司書」（AI モデル）を使うことでした。これは高速ですが、複雑な思考を捉えるには、時として一つの単語では不十分な場合があります。

研究者たちは、より正確にするために、司書に複数の単語を書かせることを試みました。しかし、従来の AI（「自己回帰型」と呼ばれるもの）では、司書はこれらの単語を一列に、一つずつ書き出す必要がありました。20 語を求めれば、まず一つ目を書き、待ち、二つ目を書き、待ち、という具合です。これはあまりにも遅く、追加の単語は待つ価値がないように見え、結果もあまり改善されませんでした。

DiffRetriever の登場：「集団思考」の司書

この論文は、拡散言語モデルと呼ばれる異なる種類の AI を使用する、DiffRetrieverという新しい種類の司書を紹介します。

ここには魔法のような仕組みがあります：
単語を列に一つずつ書く代わりに、拡散型の司書は20 個の空の枠（マスクされた位置）がある白紙のページを見て、20 個の枠すべてを完全に同時に埋めます。

次のように考えてみてください：

従来の方法（自己回帰型）：友人に映画について説明を頼むとします。彼らは「それは…」（間）「…」（間）「素晴らしい…」（間）「アクション…」（間）「映画だ」と言います。次の単語を始める前に、それぞれの単語が終わるのを待つ必要があります。20 語を求めれば、永遠に時間がかかります。
新しい方法（拡散型）：友人に 20 個の空の箱が並んだ用紙を渡します。「これらの箱に説明を書き込んで」と頼みます。彼らは用紙全体を見て、映画について考え、同時に、一瞬のうちに 20 個の箱すべてを埋めます。

論文が明らかにした点

速度対品質：研究者たちは、従来の「一つずつ」書き出す AI と、新しい「一度にすべて」書き出す AI の二種類でこれをテストしました。
- 従来の書き手に対して複数の単語を生成させると、速度はどんどん遅くなり、品質はほとんど向上しませんでした。
- 一方、新しい書き手に複数の単語を生成させると、一つの単語を書くのと同じくらい速く（すべてを一度に行うため）、品質は劇的に向上しました。同じ時間内で、はるかに優れた説明が得られたようなものです。
最高性能者：新しい「一度にすべて」の司書（具体的には「Dream」と呼ばれるもの）を訓練した後、それはテストにおいて関連するドキュメントを見つける能力が最も優れており、従来の複数単語の試行を含む他のすべての手法や、他の最新の AI モデルを凌駕しました。
「完璧な」予算：研究者たちはまた、異なる質問には異なる数の単語が必要であることに気づきました。「天気はどうですか？」のような単純な質問は 2 語だけで十分かもしれませんが、複雑な質問は 16 語必要かもしれません。
- 現在、彼らは一つの数字（例えば 4 や 16）を選び、それをすべての質問に対して使用しています。
- 彼らは、もし書き出す前に、各特定の質問に正確に何語が必要かを魔法のように知ることができれば、システムはさらに良くなり、訓練されたモデルさえも凌駕するだろうと発見しました。
- また、質問の長さのような単純な手がかりが、必要な単語数を予測できることも発見しました。これは、将来的に、各質問に対して書き出す単語数を自動的に調整し、両方の利点を享受できるシステムを構築できる可能性を示唆しています。

まとめ

この論文は、検索に複数の単語を使うことの問題点は、複数の単語を使うというアイデアそのものではなく、それらを（一つずつ）書くという古い方法にあったことを証明しています。すべての単語を一度に書く新しい AI に切り替えることで、彼らは複数単語検索を高速で、安価で、はるかに正確なものにしました。まるで、荷物を届けるために遅い列車を待つ必要はないと気づき、すべてを一度に投下するドローンが必要だということに気づいたようなものです。

技術的サマリー：DiffRetriever

問題定義

最近の研究、特にPromptReps（Zhuang et al., 2024）は、自己回帰言語モデル（AR LM）をゼロショットリトリーバーとして機能させ、クエリとパッセージの代表的なトークン（密ベクトルと疎なログオッズ）を生成させることで、その有効性を示しました。しかし、このアプローチをマルチトークン検索（1 つではなく $K$ 個の代表トークンを生成する）に拡張することは、AR モデルでは非効率的であり、しばしば効果的ではありませんでした。AR モデルはトークンを逐次生成するため、 $K$ 個の代表トークンを生成するには $K$ 回のフォワードパスが必要となり、エンコーディングの遅延が $K$ に比例して線形に増加します。この顕著な遅延コストにもかかわらず、以前のマルチトークン PromptReps の試みは、単一トークン復号化に対して一貫した性能向上を示すことができませんでした。そのため、この限界がマルチトークン概念そのものにあるのか、それとも逐次生成のボトルネックにあるのかは不明瞭なままでした。

手法

著者は、逐次生成のボトルネックを克服するためにDiffusion Language Models（DiffLMs）向けに設計された検索フレームワークDiffRetrieverを提案します。

中核メカニズム

トークンを 1 つずつ生成する AR モデルとは異なり、DiffLM は双方向注意機構の下でマスクされた位置を同時に埋めます。DiffRetriever はこれを以下のように活用します：

プロンプト構築： 代表的なトークン用プロンプト（例：「The words are [MASK]...」）に $K$ 個のマスク位置（[MASK]）を付加します。
並列復号： プロンプトと $K$ 個のマスクを含む完全なシーケンスを、単一の双方向フォワードパスで DiffLM に入力します。
表現抽出： $K$ 個のマスク位置から、 $K$ 個の密な隠れ状態と $K$ 個の疎なログオッズベクトルを同時に抽出します。

スコアリングとトレーニング

スコアリング： フレームワークはPromptRepsのハイブリッドスコアリングレシピを採用します：
- 密スコア： クエリとパッセージの隠れ状態間の ColBERT 風の遅延相互作用（MaxSim）によって計算されます。
- 疎スコア： 最大プーリングされ、ReLU 活性化されたログオッズによって計算されます。
- ハイブリッドスコア： 正規化された密スコアと疎スコアの線形補間です。
ファインチューニング： 著者は、AR および DiffLM の両方のバックボーンに適用可能な、密および疎の両方のスコアに対する対照的ファインチューニング（InfoNCE 損失）を採用します。

実験設定

本研究は、4 つのバックボーン（7B〜8B パラメータ）を比較します：

自己回帰（AR）： LLaMA3-8B、Qwen2.5-7B。
拡散（Diff）： Dream-7B（Qwen2.5 から初期化）、LLaDA-8B（ゼロからトレーニング）。
ベースライン： BM25、PromptReps（AR）、DiffEmbed（DiffLM を平均プーリングエンコーダとして使用）、および対照的にファインチューニングされた RepLLaMA。
データセット： MS MARCO、TREC DL 2019/2020、および BEIR-7（7 つの多様なドメイン外データセット）。

主要な結果

1. マルチトークンの効率性と有効性

拡散 vs 自己回帰： 試験されたすべての拡散バックボーンにおいて、マルチトークン DiffRetriever は単一トークン DiffRetriever を大幅に上回ります。一方、マルチトークン AR モデルは、単一トークン AR モデルに対して平坦な、あるいは負の利益しか示しません。
遅延： マルチトークン DiffRetriever のエンコーディングコストは $K$ に関わらずほぼ一定（単一フォワードパス）ですが、AR マルチトークンのコストは線形に増加します。ゼロショット設定（ $K \le 20$ ）において、AR マルチトークンは、一貫した性能向上を伴わずに、拡散モデルと比較して約 15 倍の遅延ペナルティを課します。
性能の逆転： ゼロショット設定では、AR モデルは $K=1$ で DiffLM よりも当初優位に立ちます。しかし、 $K>1$ になると、DiffLM（特に Dream と LLaDA）は AR 対応モデルを凌駕します。例えば、Dream のマルチトークンは、単一トークン版と比較して MS MARCO の性能をほぼ倍増させ、Qwen2.5 を追い抜きます。

2. 最先端の性能

ゼロショット： LLaDA上の DiffRetriever は、比較対象のゼロショットシステムの中で最高の BEIR-7 平均スコアを達成し、AR バックボーン上の PromptReps や DiffEmbed ベースラインを上回ります。
ファインチューニング： 教師ありファインチューニング後、Dream 上の DiffRetrieverは、ファインチューニングされた PromptReps、DiffEmbed、および対照的にファインチューニングされた RepLLaMA を凌駕し、全体として最も強力な BEIR-7 リトリーバーとなります。
バックボーンの相乗効果： Qwen2.5 から初期化された Dream バックボーンは、LLaDA に比べてゼロショットからファインチューニングへの性能変動が大きいことを示しており、AR 事前学習の事前知識と双方向拡散復号の組み合わせが検索に極めて有効であることを示唆しています。

3. バジェット分析とオラクルの余力

固定 vs 適応： 本研究は、すべてのクエリに対して固定バジェット $(K_q, K_p)$ を使用することが、大幅な性能の向上余地を放棄していることを特定しました。
オラクル分析： 各クエリに対して真の正解を用いて最適な $K$ を選択する「クエリごとのオラクル」は、固定バジェットにおける対照的ファインチューニングを、すべての組み合わせにおいて上回ります。
予測可能性： 最適なバジェットは、安価なクエリ特徴（長さおよびシャノンエントロピー）と正の相関を示しており、適応的バジェット選択のための学習されたルーターが、将来の有望な方向であることを示唆しています。

意義と主張

本論文は、マルチトークン検索を取り巻く曖昧さを解消すると主張しています：

ボトルネックは逐次生成にある： 以前のマルチトークン AR リトリーバーの失敗は、マルチトークン概念そのものによるものではなく、逐次復号のコストによるものでした。DiffRetriever は、拡散モデルを通じて逐次生成から切り離された場合、マルチトークン検索が効果的かつ安価であることを証明します。
トレーニングと整合した目的： 拡散モデルを事前学習された形式（並列にマスクされた位置を予測）でクエリすることは、DiffEmbed が行ったようにモデルを標準的な BERT 風エンコーダとして流用するよりも、検索タスクへの転移が優れています。
適応の可能性： 固定バジェットのパフォーマンスとクエリごとのオラクルのパフォーマンスの間に大きなギャップが存在することは、検索の効率性と有効性を向上させるための重要かつ未開拓の道として適応的バジェット選択を浮き彫りにします。

著者は、DiffRetriever が拡散ベースの検索における新たな最先端を確立し、並列的な代表トークンが、単一トークン拡散および逐次的自己回帰アプローチの両方を上回る性能を発揮しつつ、低遅延を維持することを示していると結論付けています。