Atlas: 検索機と読解機の共同事前学習により、11Bパラメータで540BパラメータのLLMを凌駕

Atlas は Izacard と Grave による、自身の Fusion-in-Decoder 論文の続編であり、FiD を検索機と読解機がゼロから共同でトレーニングされる完全に共同トレーニングされたシステムへと拡張したものです。私は今この論文を読んでいます。なぜなら、これはオリジナルの RAG 論文から FiD を経て共同トレーニングされた検索へと至るアーキテクチャの系譜を完結させるものであり、あらゆる元帳 QA システムがナビゲートする必要のある決定空間そのものだからです。

論文の概要

「Atlas: Few-shot Learning with Retrieval Augmented Language Models」(Izacard et al., JMLR 2023) は、検索拡張モデルが、知識集約型の少数学習（few-shot）タスクにおいて巨大なパラメータを持つ LLM に匹敵できるかどうかを問うています。中心的な貢献は、Contriever ベースの高密度検索機と T5 ベースの Fusion-in-Decoder 読解機を共同で事前学習させた、注意深く設計された検索拡張システムです。重要な洞察は、アーキテクチャではなく、共同事前学習こそが少数学習における知識性能を促進するものであるということです。このシステムは上位 20 個のドキュメントを検索し、それぞれをエンコーダーで独立してエンコードした後、デコーダーのクロスアテンションでそれらを融合させます。これは著者らの 2021 年の論文と同じ FiD 設計です。

主要なアイデア

• Atlas-11B は、わずか 64 個の学習例で Natural Questions において 42.4% の精度を達成し、パラメータ数が 50 分の 1 でありながら PaLM (540B パラメータ) を約 3 ポイント上回りました。
• TriviaQA (64-shot) において、Atlas-11B はフィルタリング済みセットで 74.5%、未フィルタリングの隠しテストで 84.7% に達し、検索コンポーネントが限られたタスク監視を強力に補完することを示しました。
• 検索機のトレーニング目的として、Attention Distillation (ADist)、EMDR2 (検索されたドキュメントを潜在変数として扱う)、Perplexity Distillation (PDist)、LOOP (leave-one-out) の 4 つが評価されました。それらの性能差は小さく、計算効率のために PDist が採用されました。
• ラベルなしテキストに対する共同事前学習が最大の要因です。すべての検索拡張事前学習構成は、検索拡張ファインチューニングのみのベースラインを大幅に上回りました。
• ドキュメントインデックスは、モデルを再学習することなく学習後に更新可能です。これは動的な知識ベースにとってアーキテクチャ上重要です。時間の経過とともに一致しなくなったインデックスは、性能を目に見えて低下させます。
• MMLU (5-shot) において、Atlas-11B は 47.9% に達し、パラメータ数が約 16 分の 1 であるにもかかわらず、GPT-3 の報告値 43.9% を超えました。

評価できる点と課題

検索によって、パラメータ数の一部で少数学習の知識性能が可能になるという主な主張は、説得力を持って維持されています。64 個の例で 42.4% という NQ の数値は驚くべき結果であり、当時最先端のスケールベンチマークであった PaLM との比較も妥当です。

しかし、3 つの懸念事項があります。第一に、共同トレーニング後であっても検索精度はそれほど高くありません。独立した分析によると、Contriever は約 85% のケースで少なくとも 1 つの正解ステートメントを見逃しており、QA 検索精度は約 47% にとどまっています。共同トレーニングは非共同トレーニングのベースラインよりも検索を改善しますが、読解機が不完全な検索を補うために膨大な作業を行っています。見出しを飾る少数学習の数値はシステムの上限を反映しており、検索コンポーネントの品質を反映しているわけではありません。第二に、インフラコストが現実的です。事前学習中にドキュメントインデックスを更新すると、計算オーバーヘッドが約 30% 増加し、Wikipedia と CommonCrawl のフルインデックスには fp16 で 587GB が必要です。これは研究環境では管理可能ですが、本番環境への導入においては真の運用上の制約となります。第三に、データ漏洩が認められていますが解決されていません。MMLU の質問の 2.8% が事前学習に使用された CCNet コーパスにそのまま出現しており、MMLU の結果を不明なマージンで膨らませています。

また、論文が完全には踏み込んでいない微妙なアーキテクチャ上の制限もあります。FiD は融合前に検索された各パッセージを独立してエンコードするため、並列処理には役立ちますが、エンコーダーにパッセージ間のアテンションがないことを意味します。複数のパッセージにわたる情報を接続する必要がある長いマルチホップ推論チェーンは、そのすべての作業をデコーダーで行う必要があります。20 個の検索されたパッセージがある場合、デコーダーのクロスアテンションは重い負荷を背負うことになります。

なぜこれが金融 AI にとって重要なのか

Beancount 元帳 QA にとって、Atlas の最も関連性の高い貢献は、共同の検索機-読解機のトレーニングが少数学習設定で報われることを実証したこと、そして、そうでない場合についての正直な説明です。数年分にわたる取引履歴を照会する Beancount エージェントは、まさに動的インデックスの問題に直面します。新しいエントリは毎日追加され、1 ヶ月古いインデックスは間違った回答を生成します。Atlas は、再学習なしでインデックスをホットスワップできることを示しており、これはアーキテクチャ上励みになります。

しかし、検索精度の数値は深刻です。もし Contriever が一般的なテキストでの共同学習後であっても、検索試行の 53% で関連する元帳エントリを見逃すのであれば、ドメイン固有の商品名、勘定科目階層、Bean ディレクティブを扱う Beancount 元帳を操作する金融ドメインエージェントには、ドメイン適応型の検索機トレーニングか、構造化クエリ手法（正確な勘定科目の一致、日付フィルタリング）によって拡張された検索のいずれかが必要になります。RAG スタイルの検索だけでは、たとえ共同トレーニングされていたとしても、高精度の元帳操作には不十分でしょう。

PaLM との比較は、アーキテクチャ上のトレードオフも明確にしています。検索を利用することで知識をより少ないパラメータに圧縮でき、推論コストを下げることができます。推論コストが大規模な運用で重要となる Beancount.io のようなプロダクトにとって、Atlas の設計思想は魅力的です。しかし、587GB のインデックスコストは負担をストレージと検索インフラに転嫁します。これはベンチマークの数値には現れない、別の種類の運用上の制約です。

次に読むべきもの

• REALM: Retrieval-Augmented Language Model Pre-Training (Guu et al., arXiv:2002.08909, ICML 2020) — Atlas が拡張した、より初期の共同検索機-読解機事前学習フレームワーク。Atlas が実際に何を改善し、何を変更せずに残したかを理解するために不可欠です。
• RA-DIT: Retrieval-Augmented Dual Instruction Tuning (Lin et al., arXiv:2310.01352, ICLR 2024) — ゼロからの共同事前学習ではなく、インストラクションチューニングを使用して Atlas と競合する性能を達成。共同トレーニングと独立トレーニングの差が、インフラコストをかけずに埋められる可能性を示唆しています。
• RETRO: Improving Language Models by Retrieving from Trillions of Tokens (Borgeaud et al., arXiv:2112.04426, ICML 2022) — 異なるスケールでの事前学習中の検索に対する DeepMind のアプローチ。元帳 QA のアーキテクチャを選択する前に、検索拡張事前学習アプローチの全体像を完結させます。