Sakana AIが汎用記憶機構を発表、不要な情報を忘れることで性能向上

人間は身の回りの無駄な情報の99パーセントを忘れることで思考を保っています。しかし、現在のTransformerはすべての入力トークンを記憶領域に抱え込み続けるため、コンテキストが長くなるほど推論コストが跳ね上がるという課題を抱えています。

経済産業省のGENIACプロジェクトの採択を受けたSakana AIは、この問題を解決する新技術「NAMMs（Neural Attention Memory Models）」を発表しました。研究論文「An Evolved Universal Transformer Memory」において、生体の記憶と忘却の仕組みに着想を得た普遍的なメモリシステムを提案しています。

勾配降下法を捨て、進化計算で忘却を学ぶ

従来のメモリ削減手法（H2OやL2など）は、直近の単語や特定の重みを残すといった手動のルールに依存しており、メモリ効率と引き換えにモデルの精度を落とすことが一般的でした。

Sakana AIは、トークンごとに「記憶するか忘れるか」を判定するニューラルネットワークを学習させようと試みました。しかし、トークンを完全に破棄する二値の判断は微分不可能であり、通常の誤差逆伝播法では最適化できません。

そこで彼らが採用したのが、得意分野である進化アルゴリズムです。モデルの変異と選抜を繰り返す試行錯誤を通じて、勾配を用いずに最適な忘却パラメータを獲得することに成功しました。

処理の流れとしては、メモリ内のトークンのアテンション値をスペクトログラムに変換し、指数移動平均で圧縮したうえで、進化した分類器が各トークンの要不要をスコアリングして剪定します。

言語で学んだ忘却術を動画やロボットへ転送

このシステムの最大の特徴はゼロショット転移性です。NAMMsはどのTransformerにも共通するアテンション行列のみを入力とするため、モデルの種類や入力モダリティに依存しません。

Llama 3 8Bのテキストデータだけで進化学習させたNAMMsを、追加学習なしで70Bモデルに適用できただけでなく、動画理解モデル（Llava Next Video）の余剰フレーム削除や、強化学習モデル（Decision Transformer）の非最適アクション削除にもそのまま適用できました。

層ごとの挙動を分析すると、浅い層ではプロンプトなどの大域的情報を保持し、深い層ではそれらを忘れて局所的な詳細情報に集中することが確認されています。

また、コードタスクでは空白やコメント、ボイラープレートを連続して削る一方、自然言語では文法の冗長な単語を削り、人名や重要概念を保持していました。同社は全コードをGitHubで公開し、日本語の長文評価ベンチマーク「ChouBun」も同時にリリースしています。

なぜ重要か

長文コンテキストの処理コストに悩む開発者にとって、モデル全体の再学習なしに後付けでメモリ消費を削れるNAMMsは非常に魅力的なアプローチです。コードとベンチマークが公開されているため、長文処理パイプラインを持つ現場なら試す価値があります。

業界全体への影響としては、KVキャッシュの肥大化に対する解決策の幅を広げた点が挙げられます。言語モデルの注意機構から抽出した忘却ルールが、動画やロボティクスといった別領域のTransformerにもそのまま通用したという事実は、モデル構造を超えた共通の認知パターンの存在を示唆しています。

私たちは、微分不可能な離散的忘却という壁を進化計算で突破した設計は極めてクレバーだと評価しています。一方で、進化学習プロセスそのものは試行錯誤の計算コストが極めて高く、今回は国家プロジェクト級のスパコン資源があってこそ成立した側面が見受けられます。学習済みの汎用重みを利用するのは手軽ですが、自前の独自モデル向けに一からNAMMsを進化させるのは、一般的な開発チームにはまだハードルが高そうです。