Learning to Remember, Learn, and Forget in Attention-Based Models

技術要約：アテンション・ベース・モデルにおける「学習、記憶、および忘却」の習得

問題提起
トランスフォーマーにおけるインコンテキスト学習（ICL）は、オンラインの連合記憶として機能し、複雑なシーケンス・タスクにおいて高い性能を実現します。しかし、ゲート付き線形アテンション・モデルや状態空間モデル（SSM）は、標準的なトランスフォーマー（KVキャッシュによる線形および二次的なスケーリングを必要とする）の計算およびメモリのボトルネックに対処するため、通常、固定サイズのメモリ状態に依存しています。この固定容量は「安定性と可塑性のジレンマ」を生み出します。すなわち、モデルは重要な既知の知識を上書きすることなく（破滅的忘却）、あるいは、古いデータの過剰な保持によって新しい情報を統合できなくなること（破滅的記憶）なく、新しい情報を学習しなければなりません。神経科学における既存のメタプラスティシティ（メタ可塑性）の手法は、明確なタスク境界への依存、メモリの増大、あるいは微分可能性の欠如といった理由から、トランスフォーマーには不適合であることが多いです。さらに、現在の線形アテンション・モデルは、シーケンス長が増加するにつれて、メモリ容量の制限と干渉の問題に直面することがよくあります。

手法：Palimpsa
著者らは、ICLを「ベイズ的メタプラスティシティ」を通じて解決可能な継続学習問題として再定義した自己アテンション・モデルであるPalimpsaを提案しています。核心となる洞察は、各アテンション状態の可塑性は、蓄積された知識を捉える事前分布に基づいた「重要度状態」に基づいて動的に調整されるべきであるという点です。

1. ベイズ的定式化: アテンション・ヘッドは、変分自由エネルギー目的関数を最小化することを目的としたベイズ推論問題としてモデル化されます。メモリ状態 $S$ に関する事後分布は再帰的に更新されます。破滅的記憶（事前分布が過度に集中してしまう現象）を防ぐため、著者らは、累積された事前重みを効果的に切り捨てる「重み付き事後分布」を導入しており、これは忘却ゲート $\alpha_t = \exp(-A d_t)$ によって制御されます。
2. メタプラスティック更新則: 共分散行列の対角近似を用いた変分推論（VI）を適用することで、著者らは閉形式の更新則を導出しています。メモリ状態 $\mu_t$ と、各状態の精度（または逆分散）を表す重要度行列 $I_t$ は同時に更新されます：
   • $I_t = \alpha_t I_{t-1} + (1 - \alpha_t)I_{prior} + \beta_t \otimes k_t^2$
   • $\mu_t = \alpha_t \frac{I_{t-1}}{I_t} \odot \mu_{t-1} + \frac{1}{I_t} \odot [(\beta_t \odot v_t) \otimes k_t]$
     ここで、$\beta_t$ は入力依存の重要度因子として機能し、$I_t$ は各状態の学習率を動的にスケールさせます。高い重要性（低分散）は安定性（低可塑性）をもたらし、低い重要性は高速な学習を可能にします。
3. アーキテクチャのバリエーション: 著者らは2つのバリアントを実装しています：
   • Palimpsa-D: ゲート付きデルタ・ルール・モデル（Deltanets）のドロップイン置換であり、$\beta_t$ のための射影と残差項を追加しています。
   • Palimpsa-M: Mamba2をベースとしており、入力ゲート $d_t$ が忘却ゲートと入力ゲートの両方の役割を兼ね備えています。
4. 理論的統一: 本論文は、様々なゲート付き線形アテンション・モデル（Longhorn, Deltanets, Mesanet）が、この共通のベイズ的枠組みにおける特定のアーキテクチャ上の選択肢、あるいは事後分布の近似として現れることを示しています。特に、Mamba2は、忘却率が非常に高いために重要度行列が一定（$I_t \approx I_{prior}$）となり、実質的にメタプラスティシティが無効化されているPalimpsaの特殊なケースであることが示されています。

主な貢献

• ICLにおけるメタプラスティシティ: 不確実性に基づいて状態の可塑性を動的に調整することにより、インコンテキストにおける破滅的な忘却と記憶を軽減する、ベイズ的メタプラスティシティの導入。
• Palimpsaモデル: 高速なGPUトレーニングのためのカスタムカーネルを用いて実装された、古い知識を段階的に解放することで破滅的記憶を防ぐ、スケーラブルなアテンション・ブロック。
• 統一された枠組み: 既存のゲート付き線形アテンションおよび状態空間モデルが、共通のベイズ的枠組みにおける特殊なケースであることを示す数学的導出を行い、Mamba2をPalimpsaの極限として明らかにしました。
• 変換パイプライン: 非メタプラスティックなモデル（学習済みのMamba2など）を、非メタプラスティックなカーネルの速度を利用して事前学習を行い、メタプラスティックな性質を利用して微調整（ファインチューニング）を行うことで、メタプラスティックなモデルへと変換する手法。

実験結果
著者らは、170M、760M、および2.7Bパラメータのモデルを用いて、合成、常識推論、およびロングコンテキストのベンチマークでPalimpsaを評価しています。

• 合成データ (MQAR): Multi-Query Associative Recallベンチマークにおいて、Palimpsaは非メタプラスティックなベースライン（Gated DeltanetやMamba2を含む）を一貫して上回っています。性能差はシーケンス長が増すにつれて拡大しており、メタプラスティシティが複雑な記憶タスクにおける上書きを効果的に軽減していることを示しています。
• 言語モデリングと常識推論: FineWeb-Eduデータセットにおいて、メタプラスティックなバリアント（Palimpsa-DおよびPalimpsa-M）は、LAMBADA、PIQA、WinoGrandeなどのタスクにおいて、非メタプラスティックな対応モデルや標準的なトランスフォーマーよりも優れたパープレキシティと精度を達成しています。具体的には、760Mスケールにおいて、ファインチューニングされたPalimpsa-Dは、Gated Deltanetのベースラインを平均精度で約0.8ポイント上回りました。
• ロングコンテキスト性能: 2.7Bパラメータのモデルを用いた実験では、最も忘却が弱い（実効的なメモリウィンドウが最大となる）8つのMamba2レイヤーをPalimpsaブロックに外科的に置換した結果、常識推論で+0.84ポイント、LongBenchベンチマークで+1.1ポイントの改善が見られました。メタプラスティック・モデルは短文脈の検索タスクではわずかに劣ったものの、コンテキスト長が増加するにつれて劣化が著しく遅く、32kトークンにおいてベースラインを上回りました。
• 学習ダイナミクス: 分析により、学習が進むにつれてメタプラスティシティ比（可塑的なシナプスと固定されたシナプスの差異）が増加し、重要な情報を安定させつつ、新しいデータに対しては可塑性を維持できることが示されました。

意義と主張
本論文は、神経科学（メタプラスティシティ）とベイズ推論の概念を統合することで、固定サイズ・メモリ・モデルにおける安定性と可塑性のジレンマに対する原理的な解決策をPalimpsaが提供すると主張しています。その意義は以下の通りです：

1. 理論的洞察: 異質な線形アテンション・アーキテクチャを単一のベイズ的枠組みの下に統一し、Mamba2を忘却優位の極限として特定したこと。
2. 実用的なスケーラビリティ: 事前学習（高速な非メタプラスティック・カーネルを使用）とファインチューニング（メタプラスティック・カーネルを使用してメモリ容量を拡張）を分離することにより、大規模なメタプラスティック・モデルの使用を可能にしたこと。
3. パフォーマンス: メタプラスティシティが、標準的なトランスフォーマーの二次的なメモリコストを必要とせずに、ロングコンテキストのシナリオにおけるメモリ保持能力と推論能力を一貫して向上させることを示したこと。

著者らは、このアプローチが情報の学習、記憶、および忘却の方法を最適化することで、厳格なメモリ制限内での高性能なシーケンス・モデリングの展開（エッジデバイスなど）を支援すると結論付けています。