今回の論文
今回取り上げるのは、Yu Sun、Xinhao Li、Karan Dalal、Jiarui Xu、Arjun Vikram、Genghan Zhang、Yann Dubois、Xinlei Chen、Xiaolong Wang、Sanmi Koyejo、Tatsunori Hashimoto、Carlos Guestrin による論文「Learning to (Learn at Test Time): RNNs with Expressive Hidden States」です。2024年7月5日に arXiv で初版が公開され、2025年8月31日に改訂版が公開されています。公開元は arXiv、研究分野は系列モデリング、長文LLM、RNN系アーキテクチャ、推論時計算です。URL は https://arxiv.org/abs/2407.04620 です。
この論文を選んだ理由は、長文を扱うときの選択肢を「Transformerを工夫する」以外にも広げてくれるからです。隠れ状態をただのベクトルではなく小さな学習器にして、推論中にも更新するという発想はかなりユニークで、長文RAG、エージェント、ログ解析、コード理解のような長い履歴を扱う開発に応用のヒントがあります。
どんな技術か
TTT Layers は、RNNのようにトークンを順番に処理しながらも、単純な固定サイズ状態ではなく、隠れ状態そのものを小さなモデルとして持つ 系列モデリング手法です。
通常のRNNは、過去の情報を1本の状態ベクトルに押し込めます。そのため長文になるほど情報が潰れやすくなります。一方でTransformerは過去トークンを広く参照できますが、attentionの計算量やメモリ使用量が重くなりやすいです。
TTT Layers はその中間を狙います。各時刻で保持する状態を「学習対象の重み」にし、その重みを自己教師あり損失で少しずつ更新しながら 次のトークン表現を作ります。要するに、コンテキストを全部保存する代わりに、その場で小さな学習を回して文脈を圧縮していく設計です。
課題
この技術が解決しようとしているのは、長文コンテキストを扱うときに、Transformerは重すぎ、従来RNNは情報圧縮が弱すぎるという問題です。
何が難しいのかというと、長文では遠く離れた位置の情報を必要に応じて使い分ける必要があるからです。コード理解なら前半の定義と後半の利用箇所、RAGなら大量文書中の少数の重要断片、エージェントなら数十ターン前のツール実行結果が効いてきます。こうした依存関係を固定サイズの状態に押し込めるのは簡単ではありません。
既存の方法にも限界があります。Transformer系は長文でも情報を保持しやすいですが、系列長に対して計算負荷が増えます。逆にRNNや状態空間モデルは線形計算量で魅力的ですが、長くなるほど「追加の文脈を見ても性能が伸びにくい」ことがあります。論文でも、Mamba は 16k を超えるあたりから、追加コンテキストを与えても perplexity の改善が頭打ちになる傾向が報告されています。
なぜこの課題を解く必要があるのかというと、実際のAIシステムでは長文を捨てにくい場面が増えているからです。社内文書検索、会話エージェント、IDE支援、運用ログ解析、監査証跡の読解などでは、入力を短く切り詰めるだけでは性能が上がりません。長文を現実的なコストで扱える系列モデルが必要です。
用語解説
- Test-Time Training
- 推論時にもモデル内部の一部を更新して、その入力列に適応させる考え方です。TTT Layers ではこの発想を系列モデルの基本動作に組み込み、文脈適応そのものを状態更新に使います。
- 隠れ状態
- 系列モデルが過去の情報を持ち回る内部表現です。従来RNNではベクトルですが、この論文では「重みを持つ小さな学習器」に置き換えることで、より表現力の高い状態にしようとしています。
- 自己教師あり損失
- 外部ラベルなしで入力自身から学習信号を作る損失です。TTT Layers ではトークンから作った train view と label view の対応を学習し、その更新結果を次の予測に使います。
- Inner Loop / Outer Loop
- Inner Loop は各入力列に対してその場で行う小さな学習、Outer Loop は通常の事前学習です。TTT Layers を理解するには、この二重ループで「推論中の適応」と「それをうまく働かせるための事前学習」を分けて考える必要があります。
- Dual Form
- TTT の逐次更新を、行列演算中心の形に書き換えて並列化しやすくした実装上の工夫です。アイデア自体は学習器更新ですが、実際に速く動かすにはこのハードウェア寄りの再定式化が重要です。
技術の仕組み
TTT Layers の発想を一言で言うと、コンテキストを固定サイズベクトルに押し込める代わりに、入力列ごとに更新される小さな学習器へ圧縮する というものです。
基本アイデア
通常の系列モデルは、過去を何らかの状態 s_t にまとめて保持します。TTT Layers ではこの状態をベクトルではなく、重み W_t を持つモデル f として表します。各トークンを読むたびに、そのトークンから作った自己教師ありタスクで W_t を更新し、その更新後のモデルで出力を作ります。
つまり状態遷移は「次の状態ベクトルを計算する」ではなく、「その場で小さな学習ステップを1回進める」です。論文ではこの仕組みを、RNNとattentionの中間にある新しい系列レイヤとして扱っています。
状態はベクトルではなく学習器
論文では、各TTT layerの状態をモデル f のパラメータ W として持ちます。単純な版では線形モデルを使い、より表現力の高い版では2層MLPを使います。
- TTT-Linear: 状態が線形モデル
- TTT-MLP: 状態が2層MLP
ここが重要です。従来RNNでは状態容量が固定ベクトルの次元数にほぼ縛られますが、TTT Layers では状態の中に「学習規則を通じて構造を記憶する能力」を持たせています。長文で必要になるのは単なる値の保持ではなく、関係性の圧縮なので、そのために小さな学習器を状態化しているわけです。
各トークンで自己教師あり更新をかける
各トークン x_t に対して、TTT layer は train view、label view、test view の3つを作ります。論文中の実装イメージでは、theta_K x_t を train view、theta_V x_t を label view、theta_Q x_t を test view として使います。
Inner Loop では、現在の状態モデル f(.; W_t) が train view から label view を再構成できるよう、MSEベースの自己教師あり損失で W_t を更新します。更新後のモデルに test view を入力して、その時刻の出力 z_t を得ます。
この設計のポイントは、ラベル付きデータを推論時に必要としないことです。入力列の中だけで学習信号を作れるので、長い文脈に合わせて状態をその場で適応させられます。
Outer Loopで「どう適応すべきか」を学ぶ
Inner Loop だけだと、どんな自己教師ありタスクを解かせれば次トークン予測に役立つのかが不明です。そこで論文では、train view / label view / test view を作る投影行列や初期重み W_0、学習率ゲート eta(x) などを Outer Loop で学習します。
これはかなり大事な点です。TTT は単なる推論時追加学習ではなく、推論時学習がうまく機能するように事前学習された層 です。したがって、場当たり的なテスト時適応ではなく、系列モデルの正規の構成要素として設計されています。
mini-batch TTTで逐次依存を緩める
素直にやると、各トークン更新が完全に逐次依存するため、GPUやTPUで並列化しにくくなります。そこで論文は mini-batch TTT を導入し、トークンを小さな塊に分けて、その塊の中では勾配計算を並列化します。実験では b=16 を採用しています。
これは品質と速度のトレードオフです。b を小さくすると更新回数が増えて品質面で有利ですが、並列性は落ちます。逆に大きすぎるとバッチ学習に近づいて、逐次適応の良さが薄れます。論文のアブレーションでは、この mini-batch 化が性能改善にも大きく効いています。
Dual Formで行列積中心に再定式化する
TTT の実計算をそのまま書くと、小さな更新が大量に並ぶため、Tensor Core のような行列積向けハードウェアを活かしにくくなります。そこで論文は Dual Form を導入し、出力が等価になる形で計算グラフを書き換えています。
理論計算量だけを見ると得ばかりではありませんが、実装上は大きな意味があります。論文では、JAX 実装で Dual Form を使うと Primal Form より学習が 5 倍超高速になったと報告しています。つまりTTT Layersは、アルゴリズムだけでなく「どう演算器に載せるか」まで込みで成立している技術です。
バックボーンとの組み合わせ
TTT Layers は単体で完結するというより、Transformerのself-attentionを置き換える系列レイヤとして挿入できます。論文では、Transformer風バックボーンだけでなく、時間方向畳み込みを含む Mamba 風バックボーンでも評価しています。
結果として、TTT-MLP は特に長文側で伸びやすく、バックボーン選択との相互作用もあります。この点からも、TTT は単独の小技ではなく、将来的に長文向け基盤モデルの設計軸になりうると読めます。
実験と結果
論文では、単純な理論比較ではなく、短文・長文・速度の3方向から評価しています。モデルサイズは 125M、350M、760M、1.3B を使い、Transformer と Mamba を主な比較対象にしています。
何を検証したのか
主に検証しているのは次の点です。
- TTT Layers が通常のTransformerやMambaに対してどの程度競争力を持つか
- 長文になるほど性能差がどう変わるか
- 推論時に本当に線形計算量の利点を活かせるか
- 実装上の工夫が速度改善にどの程度効くか
どんなデータセットや評価指標を使ったのか
短文側では Pile、長文側では Books3 を使い、主な指標は perplexity です。長文評価では 1k から 32k までコンテキスト長を伸ばし、追加の文脈を与えたときにどれだけ予測が改善するかを見ています。
評価手順はできるだけ Mamba 論文に合わせ、学習FLOPsを揃えた比較をしています。長文Transformerについては、実務でよく使われる「短文で学習してから長文にファインチューニングする」設定もベースラインとして加えています。
アブレーションでは mini-batch TTT が大きく効いた
125M モデルのアブレーションでは、線形attention相当の出発点の perplexity が 15.91 でした。そこから、LayerNormとresidual、mini-batch TTT、学習率ゲート、Mambaバックボーンなどを追加し、最終的な TTT-Linear は 11.09 まで改善しています。
特に大きかったのは mini-batch TTT で、論文ではこの変更だけで 14.05 から 12.35 まで下がっています。これは、TTT の本質が「推論時更新」という発想だけでなく、その更新単位をどう設計するかにも強く依存していることを示しています。
長文では Mamba より文脈を使い切りやすかった
論文の中心的な結果は、TTT-Linear と TTT-MLP が、Transformerに近い形で「コンテキストを増やすほど perplexity を下げ続ける」ことです。一方で Mamba は 16k 以降で改善が頭打ちになる傾向が示されています。
Books3 の 32k 評価では、TTT-Linear と TTT-MLP が Mamba を上回りました。さらに TTT-MLP は短文ではやや不利でも、長文になるほど伸びやすい傾向が出ています。これは、より表現力の高い状態モデルが長文側で効いているという解釈と整合的です。
速度面でも理論だけでは終わっていない
速度評価では、学習時の 2k コンテキストで、v5e-256 TPU pod 上の1イテレーションが Transformer 0.30 秒、TTT-Linear 0.27 秒だったと報告されています。これは特別なシステム最適化なしでも約10%速いという結果です。
さらに、Dual Form は JAX 実装で Primal Form より 5 倍超高速でした。GPU推論カーネルの結果でも、Transformerはコンテキスト長とともにトークン当たり時間が伸びる一方、TTT系はおおむね一定に保たれるとされています。
ここから言えるのは、TTT Layers が単なる理論上の線形モデルではなく、適切に実装すれば長文で実用的なコスト優位を持ちうるということです。
何に使える？
TTT Layers はまだ研究段階ですが、長文を順番に読みながら、入力列ごとに内部状態を適応させたい用途と相性が良いです。以下は論文の結果を踏まえた具体的な使い道です。一部は推測を含みますが、その場合は明示します。
長い会話履歴を持つAIエージェント
エージェントは、過去の会話、ツール結果、計画、メモなどでコンテキストが膨らみやすいです。TTT Layers は履歴を全部保持する代わりに、その履歴から小さな状態モデルを逐次更新できます。
このため、長時間のタスクで「最近の情報だけでなく、過去の傾向や構造も反映したい」場面に向いています。論文自体はエージェント評価ではありませんが、長文で追加コンテキストを使い切りやすいという結果から、履歴肥大化への対抗策として示唆があります。
長文RAGの読解レイヤ
RAGでは検索器が文書候補を絞っても、最終的に数十ページ分を読む場面があります。TTT Layers は、入力列ごとに局所的な圧縮ルールを学びながら読むので、検索後の長文読解器として面白い候補です。
特に、マニュアル、契約書、障害報告書、監査ログのように、文書全体の構造を徐々に掴みながら読む必要があるタスクで有望です。Transformerほど重くしたくないが、単純な固定状態RNNでは弱い、という中間帯で価値があります。
コード理解やログ解析
コードやログは、遠い位置に関連情報が散らばりやすい一方で、順に読んでいく処理とも相性があります。TTT Layers の考え方は、ファイルを順にスキャンしながら内部状態を適応させる解析器と噛み合います。
たとえば、大規模ログの異常パターン抽出、長い差分のレビュー支援、巨大な設定ファイル群の依存関係理解などで、メモリ量を抑えながら文脈適応を持たせる設計が考えられます。
ストリーミング処理系
TTT Layers は系列を順次処理するので、ストリーミング入力とも相性があります。たとえば、音声認識の後段言語モデル、継続的なイベント監視、リアルタイム議事録要約などで、過去全履歴をフルattentionで持ち続けるより扱いやすい可能性があります。
これは論文の直接検証ではありませんが、線形計算量と入力列ごとの適応という性質から、将来の応用先として十分自然です。
開発や事業へのヒント
この論文から得られる最大のヒントは、長文性能を上げる方法は「コンテキスト窓を広げる」か「RAGを足す」だけではないということです。モデル自身に、推論中に少し学ぶ機構を持たせる という方向があります。
小さな適応を推論パイプラインに入れる発想
TTT Layers をそのまま再現しなくても、入力ごとに軽い適応をかける設計はアプリ層でも使えます。たとえば、案件ごとに用語分布が違う社内検索、顧客ごとに定型が違うサポート支援、プロジェクトごとにコード様式が違う開発支援では、推論前後に軽い更新や再重み付けを入れる余地があります。
つまり、「汎用モデルを固定で使う」発想から一歩進んで、セッション単位・文書単位で適応する推論 を考えるきっかけになります。
attention一辺倒ではない長文設計
長文処理というと、最近は sparse attention、KV cache 圧縮、検索併用が話題になりがちです。TTT Layers は別方向で、長文を全部保持せず、学習規則で圧縮する立場を取っています。
これはプロダクト設計でも重要です。長文支援機能を作るとき、全部を高コストで見せるのではなく、どこかで「構造として圧縮して保持する」段を入れたほうがよい場合があります。モデル内部でもアプリ内部でも同じ発想が使えます。
実装はアルゴリズムだけでは足りない
この論文は、良いアイデアを出すだけでは不十分で、mini-batch 化や Dual Form のようなハードウェア最適化まで必要だと教えてくれます。AI機能を事業化するときも、研究上の精度改善だけでなく、実行時間・メモリ・並列化の形まで含めて設計しないと実運用に乗りません。
特に、社内ツールやSaaSで長文処理コストが重いなら、「アルゴリズム改善」と「システム実装」を分けずに見る姿勢が必要です。
注目すべき今後の方向性
今後注目すべきなのは、推論中にまったく更新しない固定モデルと、フル再学習の中間にある設計です。TTT Layers はその代表例です。将来的には、セッション適応、ユーザー適応、ドメイン適応を、もっと安価に層内部へ組み込む流れが強まるかもしれません。
これは推測ですが、エージェントや長文アシスタントが一般化するほど、入力列ごとの微調整能力は価値を増していくはずです。
限界
TTT Layers にも明確な限界があります。
まず、理論上は線形計算量でも、実装は単純ではありません。逐次更新、mini-batch 化、Dual Form、専用カーネルなどが絡むため、普通のTransformerを置き換えるだけで簡単に導入できる技術ではありません。
次に、TTT-MLP は表現力が高い一方で、論文でも memory I/O の課題が残るとされています。長文で有望でも、実運用で常にTransformerより速いとはまだ言い切れません。
また、この論文の評価は主に言語 modeling の perplexity 中心です。長文QA、コード生成、ツール利用、マルチモーダル推論のような実タスクでどこまで優位性が維持されるかは、今後の追加検証が必要です。
さらに、推論時更新を含む以上、再現性やデバッグの難しさも増えます。同じモデルでも入力順やセグメント分割の影響を受けやすくなる可能性があり、プロダクション投入時には挙動監視が重要です。
よくある質問
Q. TTT Layers は Transformer の代替ですか？
A. 完全な置き換え候補ではありますが、現時点では研究段階です。少なくとも論文では、長文での文脈利用という点でTransformerに近い性質を持ちながら、計算量を線形に抑える方向を示しています。ただし、実装難度やカーネル最適化まで含める必要があるため、すぐに既存LLM全部を置き換える段階ではありません。
Q. Test-Time Training というと、毎回ファインチューニングするのと同じですか？
A. 同じではありません。TTT Layers は各入力列に対して層内部の小さな状態モデルを更新しますが、その更新規則自体は Outer Loop で事前学習されています。一般的な毎回ファインチューニングよりは、はるかに軽く、系列処理の一部として組み込まれています。
Q. 実務で取り入れるなら、どこから考えるべきですか？
A. まずは「長い入力を全部保持する必要があるか」「入力ごとに軽く適応したいか」を整理するのが先です。RAG後の読解、長い会話履歴、プロジェクト依存のコード支援などで、入力単位の適応が効くなら、この論文の考え方が参考になります。実装する場合は、モデル改善だけでなく並列化戦略まで含めて設計する必要があります。
Q. TTT-Linear と TTT-MLP はどちらが有望ですか？
A. 論文の読み方としては、短めの文脈や実装容易性では TTT-Linear が扱いやすく、長文での伸びしろでは TTT-MLP が有望です。TTT-MLP はより表現力が高い代わりに memory I/O の課題が残るので、どちらが良いかは利用場面と実装基盤次第です。
Q. この技術は RAG やエージェントと一緒に使えますか？
A. 使える可能性は高いです。論文はその組み合わせを直接検証していませんが、長文入力から追加文脈を活かしやすい性質は、RAG後の長文読解やエージェント履歴処理と相性が良いです。特に「全部を保持すると高コストだが、固定状態では足りない」場面で検討価値があります。
今日の学び
この論文は、長文コンテキストを扱うときに、Transformerは重く、従来RNNは文脈圧縮が弱いという課題を扱っています。そこで、隠れ状態を小さな学習器にし、推論中にも自己教師ありで更新する TTT Layers という仕組みで解こうとしました。
ここから得られるヒントは、長文AIの改善策はメモリを増やすことだけではなく、入力ごとにどう内部状態を適応させるか を設計することでも生まれるということです。長文RAG、エージェント、コード支援のように履歴が長くなりがちな領域では、今後も注目しておきたい発想です。