今回の論文
今回取り上げるのは、Sebastian Borgeaud、Arthur Mensch、Jordan Hoffmann、Trevor Cai、Eliza Rutherford、Katie Millican らによる論文「Improving Language Models by Retrieving from Trillions of Tokens」です。発表は ICML 2022、公開元は PMLR です。研究分野は、検索拡張言語モデル、外部メモリ、長期知識の取り込みです。URL は https://proceedings.mlr.press/v162/borgeaud22a.html です。
この論文を選んだ理由は、いま広く使われている RAG の原型に近い発想を、推論時の後付けではなくモデル設計そのものに埋め込んでいるからです。特に、巨大モデルをさらに巨大化する代わりに、外部検索メモリを付けて性能を伸ばすという考え方は、開発コストと性能の両方を見直すヒントになります。
どんな技術か
RETRO は、言語モデルが次のトークンを予測するときに、直前までの文脈に似たテキスト断片を巨大データベースから検索し、その検索結果を参照しながら生成する技術です。
普通の Transformer は、学習時に重みへ知識を詰め込んだら、推論時にはその重みだけで答えます。RETRO はそこを変えて、モデルの外にある非パラメトリックな記憶領域を毎回参照するようにしました。論文では 2 兆トークン規模の検索データベースを使い、7.5B パラメータの RETRO が、The Pile で GPT-3 や Jurassic-1 と同程度の性能を示したと報告しています。
ポイントは、検索した文書をそのままプロンプトに貼る通常の RAG ではないことです。RETRO は、入力を小さな chunk に分け、各 chunk ごとに近傍テキストを取り出し、専用の chunked cross-attention でモデル内部に統合します。つまり、検索を「周辺機能」ではなく「予測機構の一部」として扱っています。
課題
この技術が解決しようとしているのは、巨大言語モデルが知識を持つために、ひたすらパラメータ数と学習コストを増やし続ける構図です。
何が難しいのかというと、言語モデルには大きく二つの仕事があるからです。一つは、文脈を理解して次の単語を自然につなぐことです。もう一つは、固有名詞、事実、定型表現、コード断片のような大量の知識を保持することです。従来はこの両方を、同じパラメータ空間に押し込む設計が主流でした。
既存の方法では、知識量を増やしたいならモデルを大きくし、より多くのデータで学習する必要がありました。ただしこの方法は、学習コストも推論コストも重くなります。また、学習後に新しい知識を入れたい場合でも、再学習や追加学習の負担が大きくなります。
なぜこの課題を解く必要があるのかというと、実際の AI システムでは「言語理解が足りない」より先に「知識の鮮度やカバレッジが足りない」が問題になる場面が多いからです。社内文書検索、法務文書、論文支援、コード支援、サポート FAQ などでは、モデルの一般知能だけでなく、必要な情報をその場で引けるか が重要です。
さらに、推論時に毎回似た情報を引けるなら、すべての知識を重みに焼き込まなくて済みます。RETRO はこの発想で、パラメータを増やす方向とは別のスケーリング手段を示しました。
用語解説
- パラメトリックメモリ
- モデルの重みに埋め込まれた知識のことです。RETRO はこの知識の持ち方を完全には捨てず、不足分を外部検索で補う設計として理解すると読みやすいです。
- 非パラメトリックメモリ
- 重みではなく、外部データベースとして保持される知識です。RETRO の本質は、巨大な非パラメトリックメモリを next-token prediction の中へ直接持ち込むことにあります。
- k近傍検索（k-NN）
- 入力に似たベクトルをデータベースから上位 k 件探す手法です。RETRO は各 chunk に対して近い chunk を検索し、その続きまでまとめて取り出します。
- Cross-Attention
- ある系列の表現が、別系列の表現を参照する attention です。RETRO では現在の文脈表現が検索結果を見るための窓口になっており、普通の自己注意だけでは届かない外部知識を取り込みます。
- Chunk
- 長いトークン列を小分けにした単位です。RETRO ではシーケンス長 2048 を 64 トークンごとの chunk に分け、chunk 単位で検索と統合を行います。これが大規模検索を現実的な計算量で回す鍵です。
技術の仕組み
RETRO の中核は、入力系列を chunk ごとに見て、前の chunk に似たテキスト断片を検索し、その結果を次の chunk の予測へ使う流れです。ここでは、通常の RAG と混同しやすい部分を分けて見ていきます。
基本アイデア
論文では、長さ 2048 の入力系列を 64 トークンずつの chunk に分割します。ある chunk を読むとき、モデルはその chunk 自体から検索するのではなく、直前までに見えている chunk を手がかりに検索します。これによって自己回帰性を保ち、未来の情報を先読みしないまま外部知識を使えます。
直感的には、文章を 64 トークン単位の小さな窓で読み進めながら、「今までに読んだ内容と近い段落」を外部コーパスから毎回探し、その情報を次の 64 トークンの予測に使うイメージです。
検索データベースの作り方
RETRO の検索データベースは、単なる文章一覧ではありません。各エントリは二つの連続 chunk [N, F] からなります。N は近傍検索のキーになる chunk、F はその続きです。
キーは、N を frozen BERT で埋め込み、そのトークン表現を平均したベクトルです。検索時には、現在の chunk も同じ frozen BERT で埋め込み、L2 距離で近いものを k 件探します。論文が frozen BERT を使うのは重要で、retriever を学習で更新しないため、データベース全体の埋め込みを何度も再計算しなくてよい設計になっています。
なぜ chunk 単位なのか
トークン単位で近傍検索すると、検索対象の数も索引サイズも巨大になりすぎます。RETRO はここを chunk 単位に落とすことで、保存量と検索コストを大きく削減しています。
しかも、取得するのは近傍 chunk N だけではなく、その続きを含む [N, F] です。これはかなり実務的な工夫です。知りたいのは「いまの 64 トークンに似ている断片」だけではなく、その断片の続きに何が書かれていたか だからです。質問応答やコード補完のようなタスクで効きやすい理由もここにあります。
モデル構造
RETRO は標準 Transformer に検索経路を足した構造です。大きく見ると、次の三つの要素があります。
- 自己回帰デコーダ
- 取得した近傍系列を処理する軽量エンコーダ
- 検索結果をデコーダへ流し込む chunked cross-attention
検索結果はそのまま生トークンで使うのではなく、まず neighbour encoder で表現へ変換されます。その後、デコーダ内部のいくつかの層で chunked cross-attention を行い、現在の文脈表現が検索結果へ注意を向けます。
論文では、cross-attention を全トークンに無差別にかけるのではなく、各 attending chunk ごとに処理しています。これにより、取得 neighbour 数が増えても計算量が線形に伸びる範囲へ抑えやすくなっています。
Chunked Cross-Attention の流れ
RETRO のわかりにくい部分はここです。各 chunk に対して取ってきた neighbour 群は、次の chunk を予測する局面で使われます。論文では、chunk の最後のトークンと次 chunk の先頭側トークンを含む形で attention をかけ、因果性を崩さず retrieval を混ぜています。
つまり、ある chunk の情報を見て検索し、その検索結果を「その chunk 自身の末尾」と「次の chunk」の予測に反映する構造です。これにより、現在の文章の流れと検索結果の対応を chunk 単位で合わせやすくしています。
学習方法
学習目標そのものは通常の next-token prediction です。ただし条件付き確率は、過去トークンだけでなく、過去 chunk に対応して取得された近傍集合にも依存します。ここが、後付け RAG ではなく検索込みの言語モデル学習になっている点です。
論文では MassiveText をベースに 5 兆トークン超のデータを扱い、学習中は 600B トークン規模、評価時は 1.75T トークン規模の retrieval database を用いています。面白いのは、推論時により大きなデータベースへ差し替えると性能が上がる点です。これは、知識の一部がパラメータではなく外部メモリ側に乗っていることを示しています。
RETROfitting という実装上の工夫
RETRO はゼロから学習するだけでなく、既存の causal Transformer を後から RETRO 化できます。論文では、ベースモデルの重みを凍結し、chunked cross-attention と neighbour encoder だけを追加学習する RETROfitting も検証しています。
この意味は大きいです。検索拡張を入れたいからといって、毎回最初から事前学習し直す必要はありません。既存の生成モデルに retrieval path を後付けできるなら、企業内の独自モデルや OSS ベースモデルにも応用しやすくなります。
実験と結果
論文では、RETRO が本当に「巨大モデルの代替スケーリング手段」になりうるかを、言語モデリング、検索データベースのスケール、後付け学習、質問応答で広く検証しています。
何を検証したのか
主な検証ポイントは次の通りです。
- モデルサイズを増やさず retrieval でどれだけ性能が伸びるか
- retrieval database を大きくすると性能がどう変わるか
- 近傍数を増やすと改善するか
- ゼロから RETRO を学習しなくても、既存 Transformer を RETRO 化できるか
- 下流の知識集約タスクでも効くか
言語モデリングでは「10倍大きいモデル級」の改善
論文では、150M から 7B 級までのモデルで RETRO の利得がほぼ一貫して出ると報告しています。特に Figure 1 では、RETRO の改善幅がモデルサイズを上げても消えず、パラメータを約 10 倍に増やしたのに近い利得に相当すると説明されています。
これはかなり示唆的です。通常は「性能を上げたければ model scale を上げる」が基本ですが、RETRO はそこに対して「検索メモリを増やす」という別の軸を提示しています。
The Pile では 7.5B RETRO が大規模モデル群と競合
論文の要約では、2 兆トークンの検索データベースを使った RETRO が、The Pile で GPT-3 や Jurassic-1 と同程度の性能を示したとされています。本文でも、7.5B RETRO は Pile の多数のサブセットで 178B の Jurassic-1 や 280B の Gopher と競合し、多くの集合で上回ったと報告されています。
ここで重要なのは、「すべてのデータで一方的に圧勝」ではないことです。dm_mathematics や ubuntu_irc のように、取得近傍が役に立ちにくい集合では改善が弱い、あるいは悪化する例もあります。つまり RETRO は万能ではなく、検索対象に有用な近傍が存在する領域で強い技術だと読めます。
Wikitext103 では retrieval database を大きくすると大きく改善
Wikitext103 の実験では、Wikipedia だけを引く設定より、より大きい MassiveText を検索対象にしたほうが大幅に性能が改善しました。論文中の表では、Wikipedia 検索の RETRO より MassiveText 検索の RETRO のほうが良い perplexity を示しています。
この結果は、RAG や検索支援システムでも重要です。retriever の精度だけでなく、検索母集団の広さと質そのものがモデル性能を左右することを示しているからです。
RETROfitting は少量追加学習でも効く
既存 Transformer に retrieval path を足す RETROfitting では、追加学習するのは主に cross-attention と neighbour encoder です。論文では、事前学習全体の約 3% にあたる 600 万系列の学習で、スクラッチ学習した RETRO に近い性能へ到達できると報告しています。
これはプロダクト開発の観点ではかなり扱いやすいです。基盤モデルを丸ごと作り直さずに、外部検索メモリを使う能力だけ後から足せる可能性があるからです。
質問応答では closed-book を大きく改善
Natural Questions では、7B の closed-book baseline が 30.4 の test accuracy だったのに対し、7.5B RETRO は DPR retrieval を使って 45.5 まで改善しました。FID や FID+Distillation には届かなかったものの、RAG や DPR よりは高い数値です。
この結果から言えるのは、RETRO は「どんな retrieval task でも最強」ではないが、自己回帰モデルでも retrieval を深く統合すれば知識集約タスクをかなり押し上げられるということです。
データリークだけで説明できるわけではない
検索系論文では、近い文が学習データ中にそのままあり、コピーで勝っているだけではないかという疑問が出ます。RETRO 論文はここも検証していて、評価文と訓練文の overlap を制限した filtered evaluation を行っています。
結果として、overlap を強く制限しても RETRO の利得は残りました。つまり、性能向上の一部は近傍コピーですが、それだけではなく、検索結果を足場に一般化している成分もあると解釈できます。
何に使える？
RETRO の発想は、そのまま現代の AI アプリ開発にも応用できます。特に「モデルを大きくするより、必要な知識を外から引くほうが合理的」な場面で効きます。
社内ナレッジ検索付きアシスタント
社内規程、仕様書、議事録、チケット履歴のように、知識量は多いが更新も多い領域では、重みへ全部覚えさせるより検索メモリ化したほうが現実的です。RETRO 型の考え方を使うと、単に検索結果を貼るだけでなく、検索情報を生成器の中間表現へどう混ぜるか を設計テーマにできます。
コード補完と開発支援
コードは定型パターン、類似 API 利用、過去実装の続きが多いため、近傍検索と相性がよい領域です。特に RETRO の「近い断片だけでなく、その続きも一緒に使う」設計は、次に来る処理や周辺文脈の予測に向いています。コードベース全体から類似関数を引き、その続きまで参照しながら補完する設計は十分考えられます。
ドメイン特化 LLM の軽量化
法務、医療、製造、論文検索のような専門領域では、巨大モデルを自前で再学習するより、比較的小さめのモデルに強い retrieval memory を付けるほうが現実的な場合があります。RETRO はその方向性の技術的な根拠になります。
知識更新が速いプロダクト
FAQ、商品情報、社内マニュアル、リサーチ支援のように更新頻度が高い場面では、重みの更新より索引の更新のほうが速いです。RETRO 的な設計思想は、知識更新をモデル再学習から分離するアーキテクチャに向いています。
開発や事業へのヒント
この論文から得られる最大のヒントは、AI プロダクトの競争力を「モデルサイズ」だけで見ないことです。検索メモリ、索引設計、統合方法の出来が、そのまま性能やコストに効きます。
RAG はプロンプト結合だけで終わらせない
多くの実装では、検索結果をそのままコンテキストへ足して終わりになりがちです。しかし RETRO は、retrieval をモデル内部の attention 経路に組み込むことで性能を出しています。自前プロダクトでも、rerank、chunk 統合、attention 制御、回答位置ごとの retrieval など、検索をどこで効かせるか を細かく設計する余地があります。
検索基盤はモデル基盤と同じくらい重要
RETRO では、frozen BERT 埋め込み、chunk 設計、近傍索引、巨大データベース管理が性能の前提になっています。これは実務でも同じで、埋め込みモデル、chunk 長、メタデータ、更新頻度が悪いと、上にどれだけ良い LLM を載せても伸びません。RAG プロダクトの差は、LLM 本体より検索基盤側で付くことが多いです。
小規模チームでも勝負しやすい
巨大基盤モデルをゼロから作るのは難しくても、特定領域で強い retrieval memory を整備することは比較的現実的です。たとえば、自社独自の文書群、コード資産、問い合わせ履歴を高品質に索引化し、生成系モデルと深く結合すれば、小さなチームでも高価値な支援機能を作れます。
注目すべき方向性
今後注目すべきなのは、retrieval を推論前の 1 回処理ではなく、生成の途中で動的に入れる流れです。これは推測を含みますが、RETRO のような chunk 単位 retrieval は、長文生成、エージェント計画、コード修正のような逐次的タスクと相性がよいです。将来的には、段階ごとに別の検索器やメモリを使い分ける構造がさらに重要になるはずです。
限界
RETRO にも明確な限界があります。
まず、検索データベースの構築と維持が重いです。論文では 2 兆トークン規模を扱っており、索引サイズも無視できません。プロダクトへ持ち込む場合でも、埋め込み生成、ANN 検索基盤、更新パイプラインが必要です。
次に、近傍検索の質に強く依存します。類似 chunk が本当に役立つ領域では強い一方、数式、対話ログ、特殊記法のように、ベクトル近傍が次トークン予測に直結しにくい領域では利得が弱い可能性があります。論文でも一部データセットで改善が小さい例が見られます。
また、データリークの影響を完全には無視できません。論文は overlap 制限付き評価を行っており、リークだけではないと示していますが、それでも retrieval 系では常に評価設計への注意が必要です。特に社内データや過去回答ログを使う場合、テストに近い表現が混ざると見かけの精度が上がりやすいです。
実装面では、通常の RAG より複雑です。単に検索結果を貼るだけならすぐ作れますが、RETRO のように生成器内部へ検索を組み込むには、モデル改造や学習が必要になります。そのため、現実の開発では「RETRO の思想をどこまで簡略化して入れるか」がポイントになります。
最後に、retrieval が入れば何でも強くなるわけではありません。検索可能な知識には効きやすいですが、純粋な推論、計画、世界モデル的な能力までは直接保証しません。
よくある質問
Q. RETRO と普通の RAG は何が違うのですか？
A. 普通の RAG は、検索結果をプロンプトや入力文へ追加してから生成することが多いです。RETRO は、検索結果を chunked cross-attention でモデル内部に取り込み、next-token prediction の一部として使います。つまり、検索が後付け補助ではなく、生成機構の中に入っています。
Q. なぜ大きいモデルを使う代わりに検索を使うのですか？
A. 知識の多くは、推論時に参照できるなら重みへ全部圧縮しなくても済むからです。RETRO は、外部メモリを使えば比較的小さいモデルでも大規模モデルに近い性能を出せる場面があることを示しました。
Q. RETRO はそのまま今の業務システムへ入れられますか？
A. そのままは重いです。論文の形を完全再現するには、専用モデル構造、検索索引、追加学習が必要です。ただし発想自体は実務向きで、段階的には「高品質な chunk 設計」「動的 retrieval」「生成途中の再検索」のような形で応用できます。
Q. RETRO は質問応答でも最先端だったのですか？
A. 論文時点では、Natural Questions で closed-book baseline を大きく改善しましたが、FID+Distillation などの専用 QA 手法には届いていません。つまり、RETRO は汎用言語モデルへ retrieval を深く統合する技術として価値があり、QA 専用最適化とは少し評価軸が違います。
Q. 小規模プロダクトなら何を真似するとよいですか？
A. 一番真似しやすいのは、知識をモデル重みへ寄せすぎず、検索メモリ側へ逃がす設計思想です。特に、chunk をどう切るか、何を近傍として採用するか、検索結果を回答生成のどの段階で使うかは、小規模プロダクトでもそのまま改善ポイントになります。
今日の学び
この論文は、知識を増やすたびにモデルを巨大化しなければならないという課題を扱いました。そこに対して、RETRO は巨大な外部テキストメモリを検索し、chunk 単位 cross-attention で次トークン予測へ統合する技術で解こうとしました。
ここから得られるヒントは明確です。AI プロダクトでは、性能向上の手段はモデル拡大だけではありません。検索メモリをどう設計し、生成器へどう深く結合するか も、十分にプロダクト差別化の源泉になります。