💬 NLP

MesaNet: Sequence Modeling by Locally Optimal Test-Time Training

原著者： Johannes von Oswald, Nino Scherrer, Seijin Kobayashi, Luca Versari, Songlin Yang, Sarthak Mittal, Maximilian Schlegel, Kaitlin Maile, Yanick Schimpf, Oliver Sieberling, Alexander Meulemans, Rif A. Sau

公開日 2026-06-04

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Johannes von Oswald, Nino Scherrer, Seijin Kobayashi, Luca Versari, Songlin Yang, Sarthak Mittal, Maximilian Schlegel, Kaitlin Maile, Yanick Schimpf, Oliver Sieberling, Alexander Meulemans, Rif A. Saurous, Guillaume Lajoie, Charlotte Frenkel, Razvan Pascanu, Blaise Agüera y Arcas, João Sacramento

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

全体像：現在のAIが抱える問題

非常に賢い司書（現在の標準であるTransformer）が、あなたの質問に答えるために本を読んでいる場面を想像してください。

優れた点: この司書は、本の最初から最後まで詳細を覚えていることが非常に得意です。
欠点: これを行うために、司書は机の上に増え続けるインデックスカードの束を保持しなければなりません。本が長くなればなるほど、束は高くなります。やがて、その束があまりに巨大になりすぎて部屋中を占領してしまい、答えを見つけるためにすべてのカードをめくる作業に、司書は圧倒されてしまいます。これにより、長いテキストを扱う際に動作が遅くなり、コストも高くつきます。

最近、科学者たちは異なる種類の司書（MambaやDeltaNetのようなRNNと呼ばれるもの）を作ろうとしました。これらの司書は、決まったサイズの小さなメモ帳だけを保持します。彼らは高速で、大きな部屋を必要としません。しかし、長い物語の最初の方を覚えておくのに苦労したり、物語が複雑になりすぎるとミスをしたりすることがあります。

解決策：MesaNetの登場

この論文の著者たちは、MesaNetと呼ばれる新しい司書を作り上げました。彼らは、小さなメモ帳を持つ司書のスピードと、大きなカードの束を持つ司書の記憶力の両方を求めていました。

これを行うために、彼らはMesaレイヤーという特別なツールを導入しました。

比喩：「即席のエキスパート」 vs 「学習の遅い学習者」

あなたが今見た数字のリストに基づいて、数学の問題を解こうとしていると想像してください。

旧来のRNN（学習の遅い学習者）: 新しい数字が現れるたびに、司書は小さな推測を行い、自分が間違っていたかどうかを確認し、メモ帳を少しずつ調整します。これを一歩ずつ行います。効率的ではありますが、単に「推測して調整している」だけなので、完璧な答えにすぐには辿り着けないかもしれません。
MesaNet（即席のエキスパート）: 新しい数字が現れたとき、MesaNetの司書はただ推測するだけではありません。彼らは、これまでに見たすべての数字に基づいて、メモ帳を調整するための完璧な方法を瞬時に計算します。彼らは、現在のコンテキストに対して答えが最適なフィットとなるよう、最適化問題を一度に解き明かします。

仕組み（「テスト時トレーニング」）

論文ではこれを**「テスト時トレーニング（Test-Time Training）」**と呼んでいます。

通常、AIモデルは工場で一度トレーニングされると、あとはただ実行されるだけです。ユーザーと会話している最中に新しいことを学ぶことはありません。

MesaNetのトリック: MesaNetは新しい単語を読むたびに、その瞬間に特化して内部メモリを「再トレーニング」するために、ほんの一瞬だけ停止します。そしてこう問いかけます。「この正確な瞬間までに読んだすべての内容を踏まえると、この情報を保存するための絶対的な最善の方法は何だろうか？」
コスト: この「再トレーニング」には、従来の方法よりも少し多くの計算能力（素早い数学ソルバーを実行するようなもの）が必要です。
メリット: 毎回「最適」な答えを解くため、従来のRNNよりも長く複雑な物語をより深く理解できます。

「チャンク（塊）」による革新

このアイデアのオリジナル版（以前の論文によるもの）は、本を一ページずつ読むように計算を一つずつ行う必要があったため、トレーニングが遅すぎました。

新しい工夫: 著者たちは、これらの計算を**チャンク（塊）**単位で行う方法を編み出しました。例えば、一ページずつ読む代わりに、司書が章全体をまとめて掴み、強力なコンピュータチップを使用して章全体の数学を一度に解き、それから次に進むようなイメージです。これにより、膨大な量のデータに対して高速にトレーニングできるようになりました。

彼らが発見したこと

チームは巨大なデータセット（数十億語）でMesaNetをテストし、Transformerや他の高速なRNNと比較しました。

始まりに強い: MesaNetは、文章や物語の始まりを理解することに長けています。最初の数百語においては、巨大なTransformerさえも凌駕することがよくあります。
長いコンテキストに強い: 他の高速なRNNは、物語が長くなるにつれて忘れたり混乱したりし始めますが、MesaNetはより強く踏みとどまります。長い文書を他の同等のモデルよりも正確に理解できます。
トレードオフ: MesaNetは「無料」ではありません。数学の問題を解くために、読書プロセス中に多くの計算量（FLOPs）を消費します。しかし、著者らは、この追加の努力が、特に深い理解を必要とするタスクにおいて、より高い精度という形で報われることを見出しました。
動的なスピード: システムは、いつ懸命に働くべきかを知るほど賢明です。もし文章が単純であれば、数学のステップを少なく済ませるかもしれません。もし文章が複雑であれば、より多くのステップを実行します。このように、タスクの難易度に応じて動的に努力を割り当てます。

まとめ

MesaNetは、読みながら物語を記憶するための完璧な方法を常に再計算し続ける司書のような、新しいタイプのAIアーキテクチャです。現在のコンテキストに最も適したメモリを見つけるために、あらゆるステップで複雑な数学問題を解くことで、他の高速でメモリ効率の良いモデルよりも優れた理解レベルを実現しています。これは、より優れた知能を得るために、多少の計算能力をトレードオフにする仕組みです。

技術要約: MESANET: LOCALLY OPTIMAL TEST-TIME TRAINING によるシーケンスモデリング

問題提起

シーケンスモデリングは現在、ソフトマックス自己注意（softmax self-attention）を利用した因果的トランスフォーマー（causal transformer）アーキテクキテクチャに支配されています。これらは効果的ですが、推論時におけるシーケンス長に対するメモリおよび計算量の要求が線形に増加するという課題を抱えています。近年、この問題に対処するためにソフトマックス操作を線形化し、定数メモリおよび計算コストを持つ再帰型ニューラルネットワーク（RNN）モデル（例：DeltaNet, Mamba, xLSTM）を実現する試みがなされています。これらの既存モデルは、その再帰的なダイナミクスがインコンテキスト回帰目的関数から派生しており、オンライン学習ルール（例：勾配降下ステップ）によって近似的に最適化されているという観点から統一されることが多いです。しかし、これらの既存のアプローチは通常、増分的な更新（一次手法）に依存しており、各タイムステップで最適な重みを解くわけではないため、過去の情報を保持しながら新しい証拠に迅速に適応する能力が制限される可能性があります。さらに、以前の最適テストタイムトレーニング層（オリジナルのMesa層など）は、忘却ゲート（forget gates）を使用する際の数値的不安定性と並列化の欠如に制限されており、スケーラブルな学習を妨げていました。

手法

著者らは、新規で数値的に安定し、チャンク単位での並列化が可能な Mesa層 を中心とした言語モデルアーキテクチャである MesaNet を導入しています。この層は、すべてのタイムポイントにおいて 局所的に最適なテストタイムトレーニング（locally optimal test-time training） を実行するように設計されています。

Mesa層

標準的な線形RNNが単一の勾配降下ステップ（例：DeltaNet）によって状態を更新するのに対し、Mesa層は $\Phi_t$ という線形写像をテストタイム最適化問題の解として定義します。具体的には、これまでに見たすべてのデータ（ $e_{1:t}$ ）に対する累積的な正則化付き二乗誤差を最小化します：
$\hat{\Phi}^{mesa}_t = \arg \min_{\Phi} \left( \frac{1}{2} \sum_{t'=1}^t \zeta_{tt'} \|v_{t'} - \Phi k_{t'}\|^2 + \frac{1}{2} \text{Tr}(\Phi \Lambda \Phi^\top) \right)$
ここで、 $\zeta_{tt'}$ は累積的な忘却係数を表し、 $\Lambda$ は対称正定値の正則化項です。
出力は $\Delta^{mesa}_t = \hat{\Phi}^{mesa}_t q_t$ として計算されます。極めて重要な点として、著者らは（忘却ゲートによる数値的不安定性や並列化の困難さを伴う）逐次最小二乗法を用いる代わりに、共役勾配法（Conjugate Gradient, CG法） を用いて最適なクエリ応答 $q^*_t = (H_t + \Lambda)^{-1} q_t$ を解きます。

並列化: 著者らは、複数のタイムステップにわたるCG法とゲート付き線形アテンション（GLA）の等価性を利用しています。状態変数 $G_t$ と $H_t$ を保持し、シーケンスのチャンクごとに線形システムを並列に解くことで、この層は学習時に $O(T)$ の複雑さを、推論時に $O(1)$ の複雑さを達成します。
動的計算量: 停止基準に達するために必要なCG反復回数（ $k$ ）は、シーケンスやトークンごとに異なります。これにより、モデルは推論コストと性能をトレードオフしながら、テストタイムの計算量を動的に割り当てることが可能になります。

アーキテクチャ (MesaNet)

MesaNetは、標準的なデコーダーのみのトランスフォーマー・バックボーン（Llama 2と同様）を採用していますが、シーケンス混合コンポーネント（通常はマルチヘッド・アテンション）をMesa層に置き換えています。このアーキテクチャには以下が含まれます：

残差ブロック (Residual Blocks): シーケンス混合層（Mesa）とチャンネル混合層（SwiGLU MLP）を含む $N$ 個のブロックをスタックしています。
前処理 (Preprocessing): キー（Keys）とクエリ（Queries）は正規化され、Mesa層の前に小さなウィンドウで畳み込みが行われます。
ハイブリッド・バリアント: 短いシーケンスのモデリング能力と長いシーケンスのモデリング能力を組み合わせるために、LRU (Hawk) 層とMesa層を交互に配置した「Hawk-Mesa」モデルも探索されています。

主な貢献

並列化可能で安定したMesa層: 本論文は、元の定式化における逐次的ボトルネックと数値的不安定性を克服したバージョンのMesa層を提示しています。これは、効率的なハードウェア（GPU/TPU）上での効率的な学習を可能にする、チャンク単位の並列化を備えたCGソルバーを利用しています。
動的なテストタイム計算量の割り当て: この層は、与えられたシーケンスの線形システムの難易度に応じて反復回数が変化する外部ソルバーを明示的に呼び出します。これにより、入力の複雑さに応じて推論計算量を変化させるための原理的なメカニズムを提供します。
広範な実証的評価: 著者らは、SlimPajamaデータセットを用いて、140M、440M、および1BパラメータのスケールでMesaNetを訓練しました。アーキテクチャ、トークナイザー、ハイパーパラメータを制御した上で、Mamba2、xLSTM、DeltaNet、Gated DeltaNet、およびTransformerとの厳密な「アップル・トゥ・アップル」の比較を行っています。

結果

言語モデリング性能: MesaNetは、すべてのスケールにおいて他のRNNベースライン（Mamba2, xLSTM, DeltaNet）よりも低い検証パープレキシティ（PPL）を達成しています。1Bパラメータのスケールでは、MesaNetは同等のアーキテクチャを持つTransformerベースラインと同等か、わずかに上回る性能を示しています。
トークン位置の分析: トークン位置に対する負の対数尤度（NLL）の詳細な分析により、明確なパターンが明らかになりました：
- 初期トークン: 多くのRNN（MesaNetを含む）は、シーケンス内の初期トークンにおいてTransformerよりも優れた性能を発揮します。
- ロングコンテキスト: Transformerは後半のトークンで優れていますが、MesaNetおよびHawk-Mesaは、他のRNNと比較して、その優位性を（512トークンを超えて）大幅に長く維持します。ただし、非常に長い深度においては最終的にTransformerに後れを取ります。
ダウンストリーム・ベンチマーク:
- グローバルな推論: MesaNetは、長期依存関係を必要とするベンチマーク（Lambada, HellaSwag, RACEなど）において、他のすべてのRNNを上回っていますが、依然としてTransformerベースラインには及びません。
- インコンテキスト・リコール: MesaNetは、他のRNNや、1024トークンのウィンドウを持つSliding Window Attention (SWA) モデルをも凌駕しますが、フルアテンションのTransformerには及びません。
- フューショット学習: MesaNetは、単語の並べ替え（word scrambling）や翻訳タスクにおいて強い性能を示し、並べ替えタスクではTransformerを上回ることがよくあります。
効率性: 推論時にCGソルバーによる追加のFLOPsを必要とするにもかかわらず、MesaNetは、特に適度な数のCGステップ（例：10-30）を使用する場合、H100 GPUやTPUv5において競争力のあるトークンスループットを維持しています。

意義と主張

本論文は、MesaNetをニューラルネットワークにおける最適なテストタイムトレーニングへの重要な一歩として位置付けています。著者らは、推論時に逐次的な最適化問題（具体的には、最適な線形結合メモリの重みを見つけること）を解くことで、単純な勾配更新を用いるモデルよりも、新しい証拠に対してより効果的に適応できると主張しています。

この研究は、トレードオフを浮き彫りにしています：パープレキシティとダウンストリーム推論の性能向上は、推論計算量（FLOPs）の増加という代償を伴います。著者らは、これが性能向上のためにテストタイムの計算量を増やすという最近の傾向と一致していると述べています。彼らは、Mesa層がこの計算量を動的に割り当てるための柔軟なメカニズムを提供し、ネットワーク内で複雑な逐次的問題を解決できる可能性があると主張しています。

結論として、MesaNetは現在の実装において計算コストという課題を抱えているものの、長いコンテキストの理解やグローバルな推論タスクにおいて他のRNNを上回る能力は、局所的に最適なテストタイムトレーニングが効率的なシーケンスモデリングのための有望な方向性であり、RNNの効率性とTransformerの性能の間の溝を埋める可能性があることを示唆しています。著者らは、現在の実装における計算コストの欠点を認め、これを軽減するためにウォームスタート最適化やハイブリッド推論戦略などの将来の研究の可能性を示唆しています。