7,736 文字
Googleリサーチが、現在のAIブームの火付け役となった「Attention Is All You Need」論文の後継となる可能性のある新しい論文を発表しました。この論文では、TITANSと呼ばれるアプローチについて詳述しています。これは、人間の記憶の仕組みに近い形でモデルに長期記憶を与え、推論時に興味深い驚き機構を用いて機能するものです。これは非常に複雑な内容ですが、できるだけ分かりやすく説明していきたいと思います。
この論文は「TITANS: Learning to Memorize at Test Time（TITANs：テスト時における記憶学習）」というもので、Googleリサーチから発表されました。論文の要旨では、Transformersの問題点について説明しています。それは、Transformersのアーキテクチャの性質上、コンテキストウィンドウを拡大すると大きなペナルティが発生するため、コンテキストウィンドウが制限されるという問題です。しかし、もしこの制限がなく、無限のトークンをコンテキストウィンドウに入れても高いパフォーマンスを維持できたらどうでしょうか。それを実現しようとしているのがTITANSです。
論文では、より正確な依存関係のモデリングにはコスト面で二次的な制約があり、モデルは固定長のコンテキストに制限されると説明しています。主要なモデルのコンテキストウィンドウを考えてみると、GPT-4oは128Kトークン、Geminiは最大で200万トークンと非常に大きいですが、私たちのニーズが増大するにつれて、このコンテキストウィンドウを超える方法が必要になってきます。
この論文で説明されているのは、TITANベースのモデルが実際に長期記憶を持ち、何に注意を向けるべきかを理解する方法です。言語モデリング、常識的推論、ゲノム解析、時系列タスクにおける実験結果は、TITANSがTransformersよりも効果的であることを示しています。さらに、200万以上のコンテキストウィンドウサイズまで効果的にスケールでき、ベースラインと比較して「干し草の中の針」タスクでより高い精度を示しています。
もしこれが本当なら、これは非常に重要な論文となります。導入部では、まずTransformersと注意機構が重要である理由について説明しています。Transformersの純粋な注意ベースのアーキテクチャは、シーケンスモデリングにおける最先端のモデルとして確立されています。基本的に、今日使用されている言語モデルの大多数、OpenAIからAnthropic、Llama、DeepSeekに至るまで、すべてがTransformersアーキテクチャと注意機構に基づいており、そのスケール学習能力により最先端となっています。
しかし、Transformersアーキテクチャには大きな欠点があります。それは、コンテキスト長に関して二次的な時間と記憶の複雑さを伴うということです。つまり、より長いコンテキスト長で苦戦し始め、推論時やプロンプト内でより多くの情報を与えるほど、パフォーマンスが低下します。
もちろん、200万トークンというのはかなりの量で、動画全体を読み込むことさえできますが、これらのモデルにより多くの情報を入力する必要性が高まるにつれて、この問題はますます大きくなっています。ビデオ理解や長期時系列予測などの複雑な実世界タスクでは、コンテキストウィンドウが非常に大きくなる可能性があり、これらのダウンストリームタスクにおけるTransformersの適用が課題となっています。
これが彼らがTITANSで解決しようとしている問題です。TITANSが本当に目指しているのは、人間の脳の働き方により近いアーキテクチャをモデル化することです。人間の脳には複数のタイプの記憶があります。短期記憶、長期記憶、メタ記憶があり、これらの異なる記憶タイプが協調しながらも独立して機能する方法は、人間の脳の働き方と思考プロセスの重要な特徴です。
これは現在のTransformersモデルにはない機能であり、TITANSが目指しているのは、複数のタイプの記憶を与え、これらの異なるタイプの記憶が協力して機能できるようにすることです。彼らは、人間の脳に似た効果的な学習パラダイムには、個別でありながら相互に接続されたメモリモジュールが存在し、各モジュールが学習プロセスに不可欠な要素を担当していると主張しています。
ここで1万フィートの視点から見てみましょう。記憶とは何か、なぜそれが重要なのでしょうか。記憶は基本的な精神プロセスであり、人間の学習と切り離せない要素です。学習と記憶は異なりますが、相互に接続されています。もし記憶がなければ、人間や動物は基本的な反射や定型的な行動に制限されてしまいます。最も基本的な生物のことを考えてみてください。彼らは単に環境に反応しているだけで、それ以上のことはありません。
人間の場合もそうで、ただ複雑すぎて実際には見えないだけなのか、それとも私たちには本当に記憶があるのか、それは別の機会の議論とします。神経心理学文献における記憶と学習の一般的な定義からヒントを得て、既存のアーキテクチャのほとんどは、記憶を入力によって引き起こされる神経の更新として考えています。何かを見たり、聞いたり、嗅いだりすることが入力であり、それが記憶が形成される方法です。
また、記憶とは別に学習を定義し、目的に応じて効果的で有用な記憶を獲得するプロセスとしています。この論文では、いくつかの質問に答えることを目指しています：
良い記憶構造とは何か
適切な記憶更新メカニズムとは何か
良い記憶検索プロセスとは何か
異なる相互接続された記憶モジュール（短期、長期、メタ記憶）を組み込む効率的なアーキテクチャをどのように設計するか
過去の長期的な記憶を効果的に保存・記憶するために深い記憶モジュールは必要か
これは、Transformersが長期記憶を扱うのに苦労する欠点の一つです。そこで、もしモデルに直接記憶を組み込むことができたらどうでしょうか。
ここで重要な部分があります。この論文では、上記の5つの質問に答えることを目指し、テスト時に効率的かつ効果的に記憶を学習できる長期神経記憶モジュールを設計しています。テスト時というのは事前学習時ではなく、モデルが実際に実行されている時、つまりプロンプトを与えられて応答を考えている時に、モデルに新しい記憶を与えたいということを強調したいと思います。
これによって、モデルはテスト時にデータをパラメータに記憶する方法を学習します。もしこれが聞き覚えがあるとすれば、それはそうかもしれません。以前、テストタイムトレーニングと呼ばれる論文を取り上げました。これは最も単純な言葉で言えば、プロンプトに基づいて学習できるモデルを可能にし、実際にそのプロンプトと推論時にパラメータを更新できるようにするものでした。
私の考えでは、最も興味深い部分は次のことです。彼らはこの記憶モジュールを、期待に反する出来事、つまり驚きがより記憶に残るように設計しました。人間がどのように機能するか考えてみましょう。本当に退屈なことをしている時や、何百万回も行ったこと、非常に反復的なことをしている時のことを考えてみてください。
これらの行動は、あなたのやり方に深く根付いているので、ただ意識が遠のいて、考えることなく行動してしまいます。運転を例に考えてみてください。目的地に着いて「あれ、運転してきた記憶がない」と思うことはありませんか？それは、長年運転しているので、考えなくても、考える必要もなく運転できるからです。
これは、何かに驚いた時とは異なります。運転中に突然割り込まれたり、目の前で大きな事故が起きたり、空からタイヤが落ちてきたりした時のことを考えてみてください。それは記憶に残るでしょう。なぜなら、それは驚きであり、普段遭遇しないことだからです。
このような機能をモデルに与えようとしているのです。この驚きのメカニズムを実際にメモリの働き方のアーキテクチャに組み込んでいます。非常に魅力的です。モデルが驚いた時、「これは記憶する必要がある」と認識するのです。
入力の驚きを測定し、記憶サイズとデータの驚きの量の割合を考慮する減衰メカニズムを提示することで、より良い記憶管理を実現します。基本的に、何かに驚いた時、最初はとても高い記憶要因を持ちますが、時間とともにその注意や優先度を下げていくというものです。
再び、人生で驚くような出来事が起きた時のことを考えてみてください。驚きの瞬間には、それについてよく考え、記憶に深く刻まれます。しかし時間が経つにつれて、その記憶は減衰し始め、詳細を忘れ始めます。そして日々、週、月、年と経つにつれて、その記憶はより抽象的になり、現在の時点では驚きではなくなるため、重要性が低下します。
この減衰メカニズムは、実際、現代の再帰モデルにおける忘却メカニズムの一般化です。では、TITANSのアーキテクチャとは何でしょうか？それは、このタイプの学習、このタイプの記憶をAIモデルに組み込む能力です。
残された重要な質問は、メモリを深層学習アーキテクチャに効果的かつ効率的に組み込む方法です。彼らは、3つのハイパーヘッドからなるTITANSという深層モデルのファミリーを提示しています：コア、長期、持続的記憶です。
コアは短期記憶として考えてください。それはその瞬間に最も重要な記憶であり、データの主要な流れを担当します。次に長期記憶があり、これは長期にわたって記憶を保存し、思い出す責任があります。そして持続的記憶は、彼らが説明するように、タスクに関する知識を符号化する、学習可能だが日付に依存しないパラメータのセットです。
これは短期記憶でも長期記憶でもなく、モデル自体に組み込まれているものです。彼らは実際に、それぞれに異なるトレードオフを持つTITANSアーキテクチャの3つの異なるバリアントを提供しています。彼らはメモリを文脈、レイヤー、ゲート付きブランチとして組み込んでいます。これらについては後で説明を試みます。
彼らは、TITANSアーキテクチャが包括的なベンチマークセット全体で、すべての現代の再帰モデルとそのハイブリッドバリアントを上回るパフォーマンスを示し、Geminiモデルの現在の最先端限界である200万以上のコンテキストウィンドウサイズまでスケールできることを観察しています。
次のセクションは、テスト時における記憶の学習についてです。繰り返しになりますが、テスト時とは推論時、つまりモデルにプロンプトを与えて応答を得る時、その応答を考えている間の期間を指します。これが興味深い部分です。なぜなら、テスト時には非常に迅速に行われる必要があるからです。
また、数週間前にテストタイムトレーニングと呼ばれる論文を取り上げましたが、これは同様の技術を使用しています。基本的に、推論時、テスト時にモデル自体を更新するというものです。
彼らは、テスト時に記憶を学習するメタモデルである神経長期記憶モジュールを提示しています。まず長期記憶は、過去の履歴の抽象化をそのパラメータに符号化します。再び繰り返しになりますが、これは人間の脳の働き方に似ています。長期的な記憶では、記憶のすべての主要な細部や些細な詳細を覚えているわけではありません。それは人間の脳の働き方ではないのです。そのため、彼らは新しいTITANモデルでも人間の脳の働き方を模倣しようとしています。それは抽象化であり、長期にわたって起こったことの大まかな図です。
しかし、記憶は常にニューラルネットワークにおいて望ましくない現象として知られてきました。なぜならモデルの一般化を制限するからです。すべてを記憶すると、一般化する能力が低下します。そのため、何を記憶するかを知ることが、これらのモデルの重要な特徴となります。また、プライバシーの懸念も引き起こします。すべてを記憶すると、IPの問題が発生し、開示されるべきでない個人情報が開示される可能性があり、テスト時のパフォーマンスが低下します。
私のお気に入りの部分である驚きの指標について説明します。先ほど議論したように、期待に反する出来事、つまり期待していたことと異なることが起こった時に驚きが生じ、それがより記憶に残ります。これにヒントを得て、モデルにとっての驚きの単純な定義は、入力に関する勾配として表現できます。つまり、期待していたものとどれだけ異なるかということです。勾配が大きいほど、入力データは過去のデータとより異なるということです。
しかし、驚きの指標は大きな驚きの瞬間の後に来る重要な情報を見逃す可能性があります。驚きの瞬間に過度に注意を向けすぎると、その後に起こることを記憶できない可能性があります。再び人間のことを考えてみてください。人間の記憶の観点からすると、出来事は長期間にわたって一貫して私たちを驚かせるわけではありません。しかし、それは記憶に残ります。
その理由は、最初の瞬間が十分に驚くべきものであり、私たちの注意を引くからです。そして、その時間枠全体を記憶することにつながります。上記の驚きの指標を改善するために、彼らは驚きの指標を過去の驚きと瞬間的な驚きに分けています。過去の驚きは非常に最近の過去の驚きの量を測定し、瞬間的な驚きは入ってくるデータの驚きを測定します。つまり、新しい驚きと、たった今起こった驚きを区別しているのです。
しかし、モデルに忘却する能力も与える必要があります。すべてを記憶することはできません。数百万のトークンという非常に長いシーケンスを扱う場合、どの過去の情報を忘れるべきかを管理することが重要です。このため、彼らは適応型忘却メカニズムを使用しています。これにより、もはや必要のない情報を記憶から忘れることができます。
この忘却メカニズムはいくつかの要素を考慮しますが、要するに、驚きと利用可能な記憶容量を考慮して、何を忘れるかを決定します。
次のセクションでは、記憶を組み込むことができる3つの異なる方法について説明し、それぞれにトレードオフがあることを先ほど言及しました。まず、文脈としての記憶（MAC）について説明します。これは非常に複雑なので、AIの助けを借りて類推で説明します。
文脈としての記憶は、会議で過去の議論の詳細なメモを取る個人アシスタントを持っているようなものです。長期記憶が必要な時に関連情報を耳元でささやき、過去の知識と現在の議論の両方を組み合わせて決定を下すのを助けます。基本的に、誰かが多くの記憶を記録し、その記憶を文脈として決定を下すために伝えるようなものです。
次にゲートとしての記憶（MAG）があります。この実装では、頭の中に2人のアドバイザーがいます。1人のアドバイザーは今起こっていることだけに焦点を当て（短期的な注意）、もう1人のアドバイザーは何年もの経験（長期記憶）から情報を引き出します。そして、ゲートキーパーが各アドバイザーの意見をどの程度聞くかを決定します。
つまり、1人は今この瞬間に起こっていることに超集中し、もう1人のアドバイザーは歴史だけを見て、現在起こっていることは見ません。そして3番目の人が、必要な決定に各アドバイザーの意見をどの程度使用するかを決めます。
そして、他の2つとは異なるレイヤーとしての記憶があります。これは本質的に、情報を異なるレイヤーを通して送り、各レイヤーが記憶のタイプとなっています。例えば、最初のフィルターは長期記憶を通してすべてを処理し、2番目のフィルターは注意を通して即時の文脈を見ます。そして各レイヤーは次のレイヤーに渡す前に情報を洗練させます。
先ほど述べたように、トレードオフがあります。それらについて説明しましょう。MAC（文脈としての記憶）は、詳細な歴史的文脈を必要とするタスクに最適です。MAG（ゲートとしての記憶）は、短期と長期の焦点を切り替えるのにより柔軟です。そしてMAL（レイヤーとしての記憶）は最も効率的ですが、やや力が弱いです。
3つの実装があり、すべてを使用するか、1つを選択するか、必要なトレードオフを考慮する必要があります。
では、パフォーマンスはどうでしょうか？上部には、テストされた異なるベンチマークがあります。Arc-E、Arc-C、Wikiなどです。上部には、私たちが話してきた異なるアーキテクチャがあります。このチャンネルでMambaについて話したことがあり、TTT（テストタイムトレーニング）とTransformersについても話しました。そしてここにTITANモデルがあります。
全体を見渡すと、基本的にTITANモデルが勝利しています。それは本当にシンプルです。これは3億4000万パラメータのモデルについてです。そして4億パラメータと7億パラメータのモデルがあり、ご覧のように、ハイライトされているものが異なるベンチマークで最高のパフォーマンスを示し、TITANSは信じられないほど良い成績を収めました。
次に、干し草の中の針テストです。これは、非常に長いコンテキストウィンドウがある場合に、忘れたり混乱したりすることなく、そのコンテキストウィンドウの深いところにある情報を記憶し、取り出すことができるかどうかを意味します。
ここでy軸は長いコンテキストからの検索の精度を示し、x軸にはシーケンス長があります。ご覧のように、ここにはGPT-4とGPT-4o-miniがグレーで示されています。これら2つのグレーのものと、上部にあるTITANs MACを見ると、コンテキスト長が増加してもパフォーマンスはかなり一貫していますが、他のモデルは急速に低下します。特にここにある他のモデルはそうです。これはフューショットセットアップについてで、こちらはファインチューニングセットアップです。
再び、ご覧のように、TITANSモデルは干し草の中の針という点で、全体的に他の異なるアーキテクチャよりも優れたパフォーマンスを示しています。
結論として、彼らはメタインコンテキスト学習者として、テスト時に記憶を学習する神経長期記憶を提示しています。これは本質的に再帰モデルであり、より驚くべきトークンと驚くべきトークンに近いトークンを適応的に記憶します。私は研究者たちのこの直感を絶対に素晴らしいと思います。
多様なタスクにおける実験的評価は、TITANSがTransformersや最近の現代的な線形再帰モデルよりも効果的であることを検証しています。特に長いコンテキストについてそうです。
以上が、この驚きメカニズムを用いてモデルにより良い長期記憶を与える方法として彼らが提案している内容です。私はこれを非常に魅力的だと思います。この論文の著者に祝福を。
コメント