技術的概要：テンソルメモリ：長視野トランスフォーマーのための固定サイズ再帰状態

1. 問題提起

トランスフォーマーは、画像や動画を処理する際、空間次元と時間次元を長いトークン系列に平坦化します。アテンション機構と KV キャッシュは過去の特性を保持しますが、系列の視野が広がるにつれて、以下の 2 つの主要な制限に直面します：

1. メモリスケーリング: 実用的な推論は、過去のキーと値の保存に依存しており、トークン数に比例して線形に増加します。これは、長尺の動画や文書に対して、計算上およびメモリ上の制約を招きます。
2. 永続状態の欠如: 従来のトランスフォーマーには明示的な永続状態が存在しません。過去の特性にアテンションを向けますが、入力の変化に応じて更新・照会できる構造化された表現を維持しません。これにより、遮蔽や観測欠落に対して脆弱となり、常に成長する履歴からグローバルな文脈を繰り返し再構築する必要に迫られ、コンパクトで永続的な状態を維持することができなくなります。

既存のアプローチ（スパースアテンション、履歴圧縮、検索拡張コンテキストなど）は、多くの場合コンテキスト長を拡張しますが、情報をトークンベクトルやスロットの非構造化コレクションとして保存します。これらは空間構造を明示的に保持せず、動画理解のような状態ベースのタスクに適したコンパクトな「世界状態」を維持しません。

2. 手法：テンソルメモリ

著者らは、標準的なトランスフォーマーブロックに固定サイズの再帰的 3 次元メモリテンソルを付加する軽量モジュールであるテンソルメモリを提案します。このモジュールは、状態容量を入力長から切り離しつつ、空間的帰納バイアスを保持します。

中核アーキテクチャ

メモリは、隠れ状態 $h_t$ とセル状態 $c_t$ からなる一対の体積として定義され、次元は $C \times D \times H \times W$ です。ここで、空間次元は入力長に関わらず一定です。このモジュールはトークンチャンクごとに以下のサイクルで動作します：

1. チャンキング: 入力トークンをチャンクにグループ化します。「読み取りクエリ」はチャンクの最初のトークンから導出され、「書き込みサマリー」はチャンク全体からの学習済み射影です。
2. 連続アドレス指定: 共有座標予測器（結合重みを使用）が、ボクセルグリッド内での読み取り ($\mu_{read}$) と書き込み ($\mu_{write}$) のための連続的な 3 次元座標を生成します。
3. 読み取り（連続サンプリング）: モデルは、$\mu_{read}$ においてトリリニア補間を用いて隠れ状態 $h_t$ から読み取ります。取得された特徴はトークン次元に射影され、ゲート付き残差接続を介してトークンストリームに融合されます。学習されたスカラーゲート $\gamma$ がこの融合の強度を制御し、特定のタスクにメモリ経路が有益でない場合に抑制することを可能にします。
4. 書き込み（微分可能なソフト書き込み）: モデルは、内容ベクトルと拡散パラメータ $\sigma$ を予測することで情報を書き込みます。この内容は、予測された 3 次元位置 $\mu_{write}$ 周辺のガウス重み付き体積として蓄積されます。この「ソフト書き込み」は、ハードなボクセル割り当てを回避し、勾配が情報の配置場所を形成することを可能にします。
5. 再帰的更新: 書き込み体積は現在の隠れ状態と連結されます。軽量な因数分解された 3 次元演算子（深度方向 3 次元畳み込みに点ごとの混合が続く）が、ConvLSTM 風のゲート ($i_t, f_t, o_t, g_t$) を計算します。これらのゲートがセル状態と隠れ状態を更新し、情報が永続化し、局所的に伝播し、観測が欠落しても利用可能であることを可能にします。

主要な設計選択

• 固定サイズ状態: メモリテンソルのサイズは一定であり、KV キャッシュに伴う線形メモリ成長を防ぎます。
• 空間構造: 3 次元ボクセルグリッドは、非構造化トークンスロットとは異なり、空間的関係を明示的に保持します。
• 微分可能なソフト書き込み: ガウス重み付き体積を使用することで、離散的でハードな割り当てではなく、滑らかで連続的な更新を可能にします。
• ゲート付き融合: 学習されたゲート $\gamma$ により、このモジュールはトレーニングを不安定化させることなく事前学習済みモデルに接続可能であり、メモリが不要な場合はゲートが経路を抑制します。

3. 主要な貢献

• 新規メモリ機構: トークンが連続座標を介して書き込み・読み取りを行う、固定サイズで再帰的な 3 次元ボクセル状態の導入。これは明示的な空間状態として機能します。
• 統合: このモジュールは標準的なトランスフォーマートレーニングパイプラインにきれいに統合され、既存のブロックに追加または削除可能であり、アーキテクチャの変更は不要です。
• 診断スイート: 永続状態が有益なシナリオ（座標バインディング、遮蔽、地図構築など）と、そうでないシナリオを分離するように設計された制御された「玩具」診断スイートの開発。

4. 実験結果

著者らは、言語、画像、動画の分野における標準ベンチマークおよび診断スイートにおいて、テンソルメモリを評価しました。

• 言語モデリング（WikiText-2）: テンソルメモリは最大の絶対的改善を達成し、テストパープレキシティを130.30（ベーストランスフォーマー）から100.07に削減しました。Transformer-XL、ローカルアテンション、線形アテンションを含むベースラインを上回りました。
• 画像パッチ再構築（CUB-200-2011）: マスク付きパッチ再構築タスクにおいて、テンソルメモリ変種（Spatial ViT）は、ベース ViT、ローカル ViT、レジスタ ViT と比較し、マスク比率（25%、50%、75%）全体で PSNR と SSIM に一貫した改善を示しました。
• 動画アクション認識（UCF-101）: このモデルは、結合時空間トランスフォーマーベースラインおよび強力な「レジスタトークン」ベースラインを、Top-1（89.21% 対 82.12%）および Top-5 精度で上回りました。著者らは、ステップごとのボクセル再帰に起因するトレーニングウォールクロック時間の大幅な増加を伴うと指摘しています。
• 玩具診断:
  • 座標バインディング: テンソルメモリは書き込み数が増加しても高い精度（75–91%）を維持しましたが、アテンションベースのベースラインはほぼランダムな性能（7–14%）に崩壊しました。
  • 無害制御: メモリを必要としないタスクにおいて、学習されたゲート $\sigma(\gamma)$ は初期化値の近くに留まり（メモリ経路を抑制）、すべてのモデルが 100% の精度を達成しました。
  • 地図構築: 情報を時間的に保持する必要がある長視野グリッドナビゲーションタスクにおいて、テンソルメモリはベースラインを大幅に上回りました。

5. 意義と主張

本論文は、トランスフォーマーに明示的かつコンパクトで構造化された状態を提供することで、長視野推論を処理する実用的な方法を提供すると主張しています。

• 容量と長さの切り離し: 固定サイズテンソルを使用することで、KV キャッシュの線形メモリスケーリングを回避し、理論的には長系列に対してよりスケーラブルです。
• 遮蔽に対する頑健性: 永続状態により、モデルはシーンに関する「信念状態」を維持でき、トークンストリームからの文脈再構築のみに依存するモデルと比較して、観測欠落や遮蔽に対してより頑健です。
• 汎用性: 長視野動画に動機付けられましたが、この機構は言語（グローバル文脈の蓄積）および画像（構造化された空間集約）にも有益であることが示されています。
• 将来の作業の基盤: 著者らは、現在の実装を基盤と見なしています。将来の方向性としては、高次元メモリ（時間認識空間メモリのための 4 次元など）、より豊かなダイナミクス（光流ワーピングなど）、効率向上のためのスパース更新が含まれます。

著者らは、現在の限界について謙虚であり、高密度なボクセル更新がスループットのボトルネックであること、およびタスクが膨大な数の異なる詳細の保存を必要とする場合、固定サイズ状態が飽和する可能性があることを認めています。しかし、彼らは、少量の構造化状態を追加することが、より能力のある長視野知覚および推論システムへの viable な道であると主張しています。