Memory Decoderとは？記憶を分離し12Bモデルを39%減で凌駕する新手法

記憶を担う独立モジュールを本体とは別に大規模化し、410M本体に6.9Bの記憶を組み合わせてPythia-12Bに匹敵、総パラメータを39%削減しました
kNN検索の分布を蒸留で学習するパラメトリック記憶により、平均ベンチマークスコアが29.86から37.34へ向上しています
Qwen3（0.6Bから14B）へプラグインでき、ドメイン特化の1.7B記憶モジュールで全スケール平均9点以上の改善を実現しました

研究の背景

現在主流のデコーダ型大規模言語モデル（Large Language Model、以下LLM）には、構造上の悩みがあります。それは、膨大な知識を覚えておく「長期記憶」の役割と、覚えた内容を使って考える「推論」の役割を、同じ一つのパラメータの塊に詰め込んでいる点です。

この2つが同じ場所に混ざっているため、知識量を増やそうとするとモデル全体を大きくするしかありません。しかし、推論に必要な計算部分まで一緒に肥大化してしまい、効率が悪くなります。知識を増やしたいだけなのに、思考回路まで膨らませることになるわけです。

今回紹介する論文「Memory Decoder at Scale: A Pretrained, Parametric Long-Term Memory」は、この記憶と推論を切り離すという発想に立っています。記憶を専用のモジュールとして分離し、本体とは別に独立して大規模化できるようにした点が中心的な貢献です。モデル内部に記憶を統合しようとするMetisのような研究とは対照的に、記憶をあえて外付けの独立部品として扱うアプローチと言えます。

Memory Decoderの仕組み

Memory Decoderは、本体モデルの中身には一切手を加えず、外側から記憶を補うパラメトリックな記憶モジュールです。ここでいうパラメトリックとは、記憶の内容を学習可能な重みとして持つという意味で、従来のように巨大なデータベースを保持する方式とは異なります。

この手法のルーツは、kNN-LMと呼ばれる検索ベースの言語モデルにあります。kNN-LMは、学習データの中から現在の文脈に近い箇所を都度検索し、次の単語の分布を取り出して本体の予測に混ぜ込む方式です。精度は高い一方で、推論のたびに大量のデータから近傍検索を行うため、動作が重いという難点がありました。

Memory Decoderは、この検索によって得られる次単語の分布を、あらかじめ蒸留（教師の出力を真似るように学習させる手法）で学び取ります。つまり、検索結果を模倣するように記憶モジュールを訓練しておくことで、推論時には重い検索を行わずに、同じような分布を直接出力できるようにするわけです。

推論の際は、本体モデルが出した予測と、Memory Decoderが出した記憶由来の分布を、重み付きで足し合わせて最終的な予測を決めます。本体は文脈を読んで考える役割に集中し、記憶モジュールは事実や知識を供給する役割を担うという分業が成立します。

大規模化を支える技術基盤

記憶モジュールを独立に大きくするという発想自体はシンプルですが、実際に大規模な事前学習で実現するには工夫が必要でした。蒸留の教師信号として使うkNN検索の分布を、膨大な学習データに対して用意しなければならないためです。

そこで研究チームは、Faiss（大規模な近傍検索を高速に行うライブラリ）を使ったインデックス作成と検索を、複数マシンに分散して処理するパイプラインを構築しました。さらに、kNNの分布をバッチ単位でスパース（必要な部分だけ）に読み込む仕組みを取り入れ、計算とメモリのボトルネックを解消しています。

これらの基盤整備によって、6.9Bという大規模なパラメトリック記憶を持つモジュールの事前学習が現実的なものになりました。記憶を独立にスケールさせるという設計思想を、実運用可能な規模で裏付けた点に意味があります。

実験結果と性能比較

最も印象的な結果は、わずか410Mの本体モデルに6.9Bの記憶モジュールを組み合わせた構成が、17個のベンチマーク全体でPythia-12Bと同等の性能に到達した点です。両者の総パラメータを比べると、本手法は合計およそ7.3Bで、12Bのモデルより39%も少なくなっています。

記憶モジュールを加える効果も明確でした。410M本体単体では平均スコアが29.86だったところ、Memory Decoderを組み合わせると37.34まで上昇しています。小さな本体でも、大規模な記憶を外付けするだけで大幅に性能が底上げされることを示しています。

構成	総パラメータ	平均スコア
410M本体のみ	約0.4B	29.86
410M本体+6.9B記憶	約7.3B	37.34
Pythia-12B	12B	同等水準

汎用性の高さも確認されています。ドメイン特化で学習した1.7Bの記憶モジュールをQwen3シリーズにプラグインしたところ、3つの専門分野で平均9点以上の改善が得られました。しかもこの効果は、0.6Bから14Bまでのあらゆるモデルサイズで一貫して現れています。既存モデルの構造を変えずに後付けできるため、実用面での応用範囲が広い手法だと言えます。

まとめと今後の展望

Memory Decoderは、LLMが抱える記憶と推論の混在という構造的な課題に対し、記憶を独立部品として分離しスケールさせるという明快な解決策を示しました。小型本体と大規模記憶の組み合わせで大型モデルを上回るという結果は、効率的なLLM設計の新しい方向性を提示しています。

一方で、いくつか留意すべき点もあります。記憶モジュール自体が6.9Bと決して小さくはなく、推論時には本体と記憶の2つを動かして分布を補間する必要があるため、単純に軽量化できるわけではありません。また、蒸留の教師信号を用意する段階で分散Faiss検索という重い前処理を要する点も、導入のハードルとなります。

とはいえ、知識を増やしたい部分だけを独立に拡張できるという設計は、モデルの更新や専門分野への適応を柔軟にする可能性を秘めています。記憶と推論をどう分業させるかという問いは、今後のLLM研究において重要なテーマになっていくでしょう。