技術的概要：言語モデルの記憶と言語のための記憶モデル

問題定義

本論文は、大規模言語モデル（LLM）の「記憶」能力に関する理解における根本的な欠如に対処する。機械学習モデルは入力情報を隠れ層のベクトル埋め込みに格納するが、因果言語モデルにおけるこれらの埋め込みの情報的内容は十分に特徴付けられていない。著者らは、次のトークン予測に最適化された現在の因果言語モデルは、本質的に高忠実度の記憶形成に適していないと仮定する。その理由は、次のトークン予測が多対一かつ非可逆な関数であるため、モデルは入力シーケンス全体ではなく、直後のトークンを予測するために必要な情報のみを保持すればよいからである。その結果、これらのモデルが下流タスク（検索や拡張コンテキストウィンドウなど）のために入力シーケンスを記憶埋め込みに圧縮して使用される場合、一般的な言語タスクに必要な「任意の情報アクセス」をサポートできない。

手法

1. 情報保持の測定

モデル埋め込み内にどの程度の入力情報が保持されているかを定量化するために、著者らは逆変換に基づく測定フレームワークを提案する：

• アーキテクチャ：テスト対象のモデルである固定化されたエンコーダーと、学習可能なデコーダーを組み合わせる。デコーダーは、エンコーダーの最終隠れ層埋め込みから元の入力シーケンスの再構成を試みる。
• 指標：
  • エントロピー比（$H_r$）：再構成の交差エントロピー損失を理論的限界（一様分布対完全再構成）と比較することで、入力情報の保持割合を測定する。
  • トークン精度：デコーダーによって正しく識別された入力トークンの割合。
• 逆変換の実現可能性：著者らは、入力再構成のために実用的に逆変換関数を見出せる場合（計算コストとして約 96 V100 時間以内）、モデルを「良好な記憶」を持つと定義する。

2. 記憶モデルアーキテクチャ

本論文は、情報保持と情報抽出を分離するように設計された並列化可能な記憶モデルアーキテクチャを導入する：

• エンコーダー - デコーダー構造：エンコーダーは入力トークンのチャンクを固定サイズの記憶埋め込みに処理する。デコーダーは、これらの埋め込み（およびオプションの新しいトークン）を処理して、後続のトークンを予測する。
• 並列化：再帰的記憶モデル（例：Bulatov et al., 2022）とは異なり、このアーキテクチャでは、エンコーダーがトレーニング時と推論時の両方で、すべての入力チャンクを並列に処理することを可能にする。
• トレーニング戦略：
  • 組み合わせ目的：次のトークン予測損失と「コピー」損失（モデルが記憶からシーケンスを再現しなければならない）の線形結合によるトレーニング。
  • カリキュラム学習：2 段階のアプローチ。まずデコーダーが「ブランクコピー」タスクを通じてトークン入力なしで情報豊富な埋め込みを処理することを学習し、その後、標準的な次のトークン予測でトレーニングを行う。
  • 固定化されたエンコーダー：まず高忠実度のオートエンコーダーエンコーダーをトレーニングし、それを固定化した後、デコーダーがこれらの埋め込みを利用するようにトレーニングする。

3. 実験設定

• データセット：FineWeb-edu、FineMath、および各種標準ベンチマーク（LongBench、IFEval など）。
• モデル：Mixers および Transformer の各種サイズ（小規模なカスタムモデルから Llama 3.1、Qwen、BERT Large まで）でテスト。
• タスク：因果的次のトークン予測、入力再生（オートエンコーディング）、コピータスク、および空白コピータスク（記憶埋め込みからのみトークンを再構成）。

主要な結果

1. 因果モデルは記憶が劣る

• 低い情報保持：因果言語モデルと検索モデルは、非常に低い情報保持を示す。モデルサイズと計算量を増やしても、因果的埋め込みからの入力再構成におけるエントロピー比（$H_r$）とトークン精度は低いままである（例：$n_{ctx}=512$ の場合、トークン精度は約 5%）。
• 非可逆性：因果モデルからの埋め込みは入力シーケンスを再構成するには不十分であり、次のトークン予測目的が包括的な記憶の形成を促さないことを確認する。

2. オートエンコーダーはほぼ完璧な記憶を達成

• 高忠実度：入力再生のために特別にトレーニングされたオートエンコーダーは、トークン精度が 90% を超え、エントロピー比が 1.0 に近い、ほぼ完璧な記憶形成を達成する。
• データ依存性：これらの高忠実度の記憶は「自明」ではない（つまり、単にトークンをベクトルスロットにマッピングするわけではない）；分布外データに一般化し、一様ランダムノイズでは失敗するため、意味のあるデータ依存表現を学習していることが証明される。

3. 組み合わせ目的とカリキュラム学習の成功

• 組み合わせトレーニング：因果的およびコピー目的の組み合わせでトレーニングされたモデルは、因果的予測能力を維持しながら、情報豊富な記憶の形成と復号化を学習することに成功する。
• カリキュラム学習：「ブランクコピー」カリキュラム（トークン入力なしで記憶の復号化をまず学習すること）は決定的である。これにより、デコーダーは固定化された高忠実度オートエンコーダー埋め込みを効果的に活用できる。このカリキュラムがなければ、デコーダーは記憶に富んだ埋め込みを活用できず、トークンベースの予測に依存しがちである。
• 性能：組み合わせ目的モデルは、純粋な因果モデルと比較して、情報関連のベンチマーク（LongBench、マルチドキュメント QA など）で顕著な改善を示すが、一般的な言語理解ベンチマーク（HellaSwag など）では結果が混在するか中立的である。

4. アーキテクチャの効率性

• 並列化：提案された並列エンコーダー - デコーダーアーキテクチャは、推論の計算複雑性を $O(n^2)$ から約 $O(n^{4/3})$（完全並列化の場合のエンコーダーでは $O(n^{2/3})$）に削減し、最初のトークンまでの時間とメモリ使用量において大幅な利益をもたらす。
• エンコーダー - デコーダーのミスマッチ：本論文は、「Mixer」エンコーダーと「Transformer」デコーダーはしばしば互換性のない学習多様体を持ち、組み合わせると効率性が低下することを指摘しており、アーキテクチャ固有の整合性が必要であることを示唆している。

意義と主張

本論文は、記憶モデルが歴史的にフルコンテキスト Transformer に劣っていた理由に対する理論的および実証的基盤を提供すると主張する：因果的トレーニングは本質的に非可逆であり、したがって高忠実度の記憶形成には不適切である。

主な貢献は以下の通りである：

1. 特徴付け：標準的な LLM 埋め込みは、規模に関係なく、比較的一般的な入力情報をほとんど含んでいないことを実証すること。
2. パラダイムシフト：正確な記憶形成には、次のトークン予測を超えた目的、具体的には情報保持を明示的に要求する組み合わせ目的またはカリキュラム学習が必要であることを提案すること。
3. アーキテクチャ：計算効率と任意の入力情報の保存・検索能力のバランスを取る並列化可能なエンコーダー - デコーダー記憶モデルを導入すること。
4. 検索への示唆：これらの知見は、現在の埋め込みベースの検索システムが、埋め込み自体の低い情報密度によって制限されており、理論的限界と比較して任意の情報検索タスクで失敗することが多い理由を説明することを示唆する。

著者らは、次のトークン予測のみでは記憶に不十分であるが、モデルは「情報豊富」な記憶を形成し復号化するようにトレーニング可能であり、計算効率と任意の情報アクセス能力の両方を提供する新しいクラスの記憶モデルを可能にすると結論づける。