検索せずに外部知識を使う：LLMの中間層に「予測した記憶」を溶かし込むDFM

どんな論文か

上海交通大学 人工知能学院の LUMIA Lab と、上海人工知能実験室（Shanghai Artificial Intelligence Laboratory）の研究チームによる論文です。事前学習で覚えきれなかった外部知識をどう LLM に持たせるか、という課題に挑んでいます。

これまでの手法には、それぞれ弱点がありました。継続事前学習（CPT）や LoRA でモデルを追加学習すると、もともとできていたことを忘れる「破滅的忘却」が起きることがあります。一方、RAG のように推論時に検索する方式では、検索や資料の処理に時間がかかり応答が遅くなります。

そこで著者らが提案するのが Deep Fusion Memory（DFM）です。基盤モデルは凍結したまま、外付けの「メモリ予測器」と「融合モジュール」だけを学習します。予測器は、前の文脈から「本来なら検索で取ってくるはずのチャンク埋め込み」を予測する仕組みです。そして予測した記憶を、デコーダの中間層へゲート付きクロスアテンションで流し込みます。

この論文の肝

一言でいうと、検索結果の埋め込みそのものを当てる予測器を学習し、それを出力層ではなく中間層に混ぜることです。

これまでのパラメトリックメモリ（MLP Memory や Memory Decoder）は、トークン1つごとに検索し、そこから作った「次トークンの確率分布」を予測するよう学習していました。そのため学習データを作るときに、コーパス内のすべてのトークン位置で検索を走らせる必要がありました。さらに、メモリの出力を最後の出力確率に混ぜ合わせる（補間する）だけにとどまり、論文の言葉でいえば「浅い融合」になっていました。

DFM は、この2点を両方とも塗り替えました。教師信号は32トークンのチャンク単位で作り、チャンク内のトークンでその記憶を共有します。予測した記憶は中間層の隠れ状態に注入するため、後続の層が文脈と一緒に記憶を深く処理できます。

図1　トークン単位のメモリとDFMの違い。従来のパラメトリックメモリ（MLP Memoryなど）とDFMの、教師データの作り方と記憶の混ぜ方を比べたもの。

使いどころと価値

意味: この研究が明らかにしたのは、粗い単位（チャンク）の教師信号でも、パラメトリックメモリは十分に学習できるという事実です。実験結果では、チャンク単位の教師で作った記憶が、トークン単位の MLP Memory を QA の平均点で上回りました。さらに、記憶を「どこに」混ぜるかが大きく効くことも分かりました。最終層だけに混ぜると、平均点が大きく落ち込みます。「検索の代わりに、学習済みの記憶を使う」という潮流の中で、作る手間と使い方の両方を一歩進めた研究といえます。

価値: 最も際立つのはコストの削減です。4B トークンのコーパスで比べると、学習データ作成時の検索クエリ数は 4B から 125M に、取得する件数は 4.096T から 2B にまで減ります。16 枚の A800 での学習時間も、MLP Memory の約 64 時間に対して約 24 時間で済みました。追加パラメータも 1.5B から 0.75B に半減しています。精度の面でも、6種類の QA の平均点で、最強の比較手法を Mistral-7B-v0.3 で 12.71%、Qwen3-8B-Base で 11.19% 上回りました（相対値）。

使いどころ: 筆者の見立てでは、次のような場面にうってつけです。

• 社内の技術文書や法令集のように、決まったコーパスの知識を自前のモデルに持たせたい。ただし、全パラメータを追加学習して元の能力が落ちるのは避けたい。
• 応答の遅さを抑えるため、リクエストのたびにベクトル DB を引く構成をやめたい。論文の推論では、データストアの検索も、取得した文章も、外部の埋め込みエンコーダも一切使いません。
• コード・数学・法律のように、分野ごとの特化モデルを用意したい。論文では分野ごとに別のモデルを学習し、3分野すべてで基盤モデルを上回る成果を出しています。

今日からできること: 本文によると、コードは「近日公開」とのことです。公開されれば、凍結した Qwen3 や Mistral にモジュールを付けて学習する道が開けます。コードがなくても今から取り入れられるアイデアは2つあります。1つは、自前の検索拡張やメモリ学習の教師データを、トークン単位ではなく固定長チャンク単位で作ること。もう1つは、外付けモジュールを足すときに、出力層の近くではなく中間から後ろの層へ分散させて組み込むことです。なお、筆者は手元で小さな実験を試せませんでした（理由は後述）。

向かない場面・前提: 学習には相応の計算資源が要ります。論文の QA 設定では、A800 80GB を 16 枚使って約 24 時間かかっています。加えて、コーパス全体を埋め込んで索引を作る作業も必要です。記憶は学習時のコーパスで決まるため、頻繁に更新される知識は再学習なしには反映できないと考えられます（論文はこの点を評価していません）。また、すべての課題で検索に勝てるわけではありません。いくつかのベンチマークでは、RAG や xRAG のほうが強い結果を出しています。

仕組み

DFM は3段階で動きます。オフラインで教師データを作る、学習する、推論する、の3つです。

図2　DFMの学習と推論の流れ。オフラインでの教師作りから、検索を使わない推論までの手順。

1. 教師データ作り（オフライン）: 学習用の系列を、重ならない32トークンのチャンクに区切ります。各チャンクについて、1つ前のチャンクを検索クエリにして、データストアから類似度上位16件のチャンクを取ってきます。取ってきたチャンクは、凍結したテキストエンコーダ（Qwen3-Embedding-0.6B）であらかじめ 1,024 次元の埋め込みに変換しておきます。これを検索順位の順に16個並べたものが、そのチャンクの「目標メモリ」になります。最初のチャンクには前のチャンクがないため、教師から外します。

2. メモリ予測器: 予測器は因果的 Transformer（前の位置しか見ない Transformer）です。入力には、基盤モデルのトークン埋め込みをチャンクごとに平均したものを使います。そこから、16個のメモリスロットそれぞれに対応する出力ヘッドで、次のチャンク用の連続ベクトルを出力します。論文の設定では、予測器は28層、隠れ幅は 1,024 です。

3. メモリ融合: 予測したメモリを、共有の射影器で基盤モデルの隠れ空間の幅に変換します。それを、選んだデコーダ層（Mistral では 0 始まりで 9・16・24・31 層）のクロスアテンションで参照します。クエリは各トークンの状態、キーと値はメモリスロットです。ヘッドごとにシグモイドのゲートを設け、そのトークン・そのヘッドがメモリをどれだけ使うかを決めます。この出力は自己注意と並列に計算し、同じ残差に足し合わせます。ゲートのバイアスは -2 で初期化します。こうすることで最初はメモリの寄与が小さくなり、事前学習済みのモデルを乱しにくくなります。

4. 学習: 損失は2つの和です。1つは通常の次トークン予測のクロスエントロピー、もう1つは予測メモリと目標メモリの平均二乗誤差（MSE）です。言語モデルが参照するのは、学習中も一貫して「予測した」メモリであり、目標メモリは教師としてのみ使います。基盤モデル、検索用エンコーダ、目標メモリは固定し、予測器・射影器・クロスアテンション・ゲートだけを1エポック学習します。

たとえるなら、「毎回図書館で調べる」のではなく、「この話の流れなら、こういう資料が出てくるはず」という勘を予測器に叩き込み、その勘を推論の途中に差し込む仕組みです。

5. 推論: 目標メモリは作りません。予測器が入力から直接メモリを生み出します。QA の評価では、見えているプロンプトからメモリを計算し、答えを生成する間は固定したままにしています。

結果

QA は、2018年12月の英語版 Wikipedia（約 4B トークン）で学習し、単一ホップ（NQ、WebQA）と複数ホップ（HotpotQA、2Wiki、MuSiQue、Bamboogle）の6課題で評価しています。比較相手は、非パラメトリックの RAG・xRAG・RETRO と、パラメトリックの CPT・LoRA・MLP Memory です。

Mistral-7B-v0.3 では、6課題の平均が 27.21 に達しました。最も強かった比較手法 xRAG の 24.14 を上回っています[引用1]。基盤モデルのみ（13.63）と比べると 13.59 点の上積みです。Qwen3-8B-Base でも 26.46 を記録し、最も強かった MLP Memory の 23.80 を 11.19% 上回りました。Bamboogle では、Mistral の基盤モデルの 3.91 から 22.65 まで跳ね上がっています。

図3　Mistral-7B-v0.3での6課題QA平均。6種類のQAベンチマークの平均点（論文の表1の値）を手法ごとに並べたもの。

一方で、勝てなかった課題もあります。Qwen3 の WebQA は、基盤モデルの 41.78 から 39.47 に下がりました。著者自身も、記憶による適応がどのモデルにも一様に効くわけではないと述べています[引用4]。Mistral の 2Wiki では RAG（31.02）が DFM（25.00）を上回り、Qwen3 の HotpotQA では xRAG（27.74）が DFM（26.46）を上回りました。

分野適応（コード・数学・法律）の平均は、Qwen3-4B-Base で 64.69（基盤から +6.51）、Qwen3-8B-Base で 67.52（+4.38）でした。どちらも比較手法の中で最高スコアです。LawBench を見ると、RAG はむしろ基盤モデルより下がったのに対し（4B で 24.35、8B で 26.89）、DFM はそれぞれ 40.41 と 46.61 を記録しました。ただし 4B では、CPT（41.66）が DFM を 1.25 点上回っています。

効率の面では、学習時間が約 64 時間から約 24 時間に短縮されました[引用2]。教師データ作りの検索クエリも、4B から 125M に減っています[引用3]。

アブレーション（部品を1つずつ変える実験）では、融合する層の位置が大きく効いていました。9層目から始めて2〜8層に融合した設定では 27.09〜27.54 で推移しました。これに対して、最終層だけだと 21.11、最後の4層に集めると 22.07 まで落ちます。チャンク長（16・32・64 で 26.51・27.21・27.07）とスロット数（2・4・8・16 で 26.88・27.04・27.09・27.21）は、平均点への影響が小さい結果でした。ゲートのバイアスを 0 で初期化すると 27.90 になり、既定値の -2 より高くなっています。

限界と注意

• アブレーションはすべて1回ずつの実行で、ばらつきは示されていません。Bamboogle は 125 問しかなく、平均点の差はこの課題のスコアの上下に左右されている箇所があります。
• MLP Memory の補間の重みは、報告している評価用データの上で課題ごとに一番良い値を選んでいます。比較相手に有利な条件ですが、読むときには押さえておきたい点です。
• QA の推論では、メモリのクエリを「いま見えているプロンプトのチャンク」から作っています。学習時の「1つ前のチャンク」とは扱いが異なります。生成した答えの部分ではメモリを更新しません。
• 検索件数の比較は、検索の出力の個数を数えたものです。実際にかかった構築時間や、FLOPs、保存容量を測ったものではないと、著者自身が断っています。
• 再現には、A800 80GB を 16 枚程度と、Wikipedia 全体の埋め込みおよび FAISS の索引作りが必要です。コードは「近日公開」とされており、本文の時点では公開先の記載がありません。
• 事例紹介（HotpotQA の6問など）は、DFM が正解した例を選んだものです。頻度を示すものではないと明記されています。

筆者は論文を読んで解説しています。手元で試したのは下の小さな実験だけで、論文の結果の追試はしていません。

手元で試した結果

今回は手元で試せませんでした。この論文の肝は、学習済みの予測器が出す連続値のメモリを、凍結したモデルの中間層へゲート付きクロスアテンションで混ぜることにあります。Claude Haiku への問いかけでできるのは文章を文脈に足すことだけで、それでは論文の比較相手である RAG そのものになってしまい、肝の部分を確かめる形を作れませんでした。

手元で試すなら

本文にはコードの公開先の記載がありません（「The code will be released soon」とだけあります）。学習に使われた公開モデルは Mistral-7B-v0.3、Qwen3-8B-Base、Qwen3-4B-Base で、埋め込みには Qwen3-Embedding-0.6B が使われています。

原文からの引用

[引用1] On Mistral, it raises the six-task mean from xRAG’s 24.14 to 27.21, a 12.71% relative improvement over the strongest baseline.

（訳）Mistral では、6課題の平均を xRAG の 24.14 から 27.21 に引き上げた。最も強い比較手法に対して相対で 12.71% の改善である。

[引用2] Trained on 16 A800 GPUs with the same corpus and number of epochs, MLP Memory and DFM take approximately 64 and 24 hours, respectively, reducing training time by approximately 62.5%.

（訳）同じコーパスと同じエポック数で 16 枚の A800 GPU で学習すると、MLP Memory は約 64 時間、DFM は約 24 時間かかり、学習時間は約 62.5% 減る。

[引用3] Figure 2(c) shows that, for a 4B-token corpus, token-level construction requires 4B retrieval queries, whereas 32-token chunks require 125M, a 96.88% reduction.

（訳）図2(c)が示すように、4B トークンのコーパスでは、トークン単位の構築には 4B 回の検索クエリが要るが、32トークンのチャンクでは 125M 回で済み、96.88% の削減になる。

[引用4] The contrasting WebQA outcomes show that memory adaptation is not uniformly beneficial across backbones, even when the original weights remain frozen.

（訳）WebQA で結果が分かれたことは、元の重みを凍結したままであっても、記憶による適応がどの基盤モデルにも一様に効くわけではないことを示している。