新アーキテクチャ「Titans」発表、200万トークン超の文脈と長期記憶を両立

長大なドキュメントをAIに読み込ませて細部まで正確に答えさせる作業は、常に計算コストとの戦いでした。従来のAttentionは精度が高いものの二次関数的な計算コストがかかり、RNNのような再帰型モデルは履歴を固定長に圧縮するため過去の文脈を失いがちでした。

研究チームが新アーキテクチャ「Titans」をarXivで公開しました。人間の認知構造のように記憶の役割を分け、推論速度を維持したまま200万トークンを超えるコンテキスト長を実現しています。

Attentionを短期記憶にし、長期記憶を外付けする

過去10年間、ディープラーニングは再帰モデルとAttentionのトレードオフに悩まされてきました。Attentionは文脈全体のトークン間の依存関係を正確に捉えられますが、計算コストが膨大で窓のサイズに限界があります。

そこでTitansは、Attentionの役割をあえて限定し、局所的な依存関係を高精度に処理する「短期記憶」として位置づけました。

その上で、テスト時（推論時）に文脈を学習して記憶する「ニューラル長期記憶モジュール」を新たに導入しました。この長期記憶が過去の膨大な情報を保持し、現在の文脈を処理するAttentionを補助します。

この長期記憶モジュールは、推論が高速なだけでなく、訓練時の高速な並列処理にも対応しています。論文では、これら2つのモジュールを効果的に統合するための3つのバリアントが提案されています。

200万トークンの「針」も高精度に発見

この二重記憶構造は、実験でも明らかな成果を示しました。言語モデリングや常識推論タスクにおいて、Titansは従来のTransformerや近年の最新線形再帰モデルを上回る有効性を記録しました。

さらに自然言語にとどまらず、超長文処理が求められるゲノミクスや時系列データのタスクでも高い性能を発揮しています。

特に注目されるのが、大量のテキストから特定の情報を探し出す「Needle In A Haystack（干し草の中の針）」テストです。Titansは200万トークン以上のコンテキストサイズまでスケールし、既存のベースラインよりも高い検索精度を達成しました。

なぜ重要か

数千ページの仕様書や巨大なコードベース、ゲノム配列を扱う開発者にとって、Titansのアプローチは非常に実用的です。Transformer並みの検索精度を維持しながら、線形モデルに近い処理速度と超長文対応を両立できる可能性を示しています。

業界全体にとっても、コンテキスト長競争の新たな道を開くものと言えます。単純にAttentionを薄く引き延ばす工夫や、状態空間モデル単体に頼るのではなく、短期記憶とテスト時長期記憶をモジュールとして切り分けた設計は非常に合理的です。

長期記憶をテスト時に適応させるアプローチは筋が良いと感じられます。一方で、実際のプロダクト運用において長期記憶の状態管理やキャッシュ構造がどう機能するのか、今後のオープンソース化や実装検証が待たれます。