2024年12月18日 12時45分ソフトウェア

日本に拠点を置くAI企業「Sakana AI」がキャッシュメモリの使用量を最大75％削減する技術「NAMM」を開発

Sakana AIが「進化したユニバーサルトランスフォーマーメモリ」という論文を発表しました。論文によると、ニューラルアテンションメモリーモデル(NAMM)という新たな仕組みを導入することでコンテキストにおいて冗長な情報を削減でき、推論時のメモリ効率を高めることができたとのことです。

An Evolved Universal Transformer Memory
https://sakana.ai/namm/

[2410.13166] An Evolved Universal Transformer Memory
https://arxiv.org/abs/2410.13166

AIの本体である大規模言語モデル(LLM)では、ユーザーからの入力であるプロンプトをコンテキストとして受け取り、新たな返答を生成しています。現在主流のモデルは数十万から数百万トークンという非常に長いコンテキストを受け取ることが可能ですが、コンテキストの長さに応じてコンピューティング費用がかかったり、パフォーマンスが低下したりするという問題がありました。

Sakana AIはLLMのメモリに保存されている各トークンを「記憶する」か「忘れる」かを決定するシンプルなニューラルネットワークとしてNAMMを開発。プロンプトのうち、重要な部分を残す一方で冗長な部分を積極的に「忘れる」ことで効率とパフォーマンスを同時に向上させることに成功しました。

NAMMはモデル本体とは別にトレーニング可能なため、容易に導入できるものの、モデルの内部にアクセスする必要があるためオープンソースモデルにしか適用できないとのこと。また、トレーニングでは勾配ベースではなく進化アルゴリズムベースの最適化が行われていると述べられています。

NAMMはアテンションレイヤーで動作し、アテンションの値に応じてコンテキストウィンドウ内のトークンを保持するか廃棄するかを選択していきます。アテンションベースのメカニズムを採用することで、トレーニング済みのNAMMを別のモデルで使用する事が可能になっているとのこと。

Sakana AIはLlama 3 8Bモデル上でNAMMをトレーニングし、ベンチマークを行いました。タスクの実行に必要なキャッシュメモリを最大で75％削減しつつ、H2OやL2という他のメモリ削減技術と比較してパフォーマンスで大きく上回ることに成功しています。

NAMMをテキストのみのデータでトレーニングした場合でも、調整不要で画像や動画を扱うマルチモーダルなモデル上で適切に動作させることができ、冗長な動画フレームを削除するなどの働きを確認できました。

Sakana AIによると、NAMMはタスクに応じてどのトークンを削減するかを変化させているとのこと。例えば、プログラミングのタスクであればまとまったコードやコメント、空白などを削除するのに対し、自然言語タスクでは文法的に冗長な部分を文の途中であっても削除しています。

Sakana AIは「モデルのトレーニング時にNAMMを使用する事で、非常に長いデータシーケンスでも効率的にトレーニングできるようになる可能性がある」と将来の展望について語りました。