title: Learning to (Learn at Test Time): RNNs with Expressive Hidden States
summary: Self-attention performs well in long context but has quadratic complexity.
  Existing RNN layers have linear complexity, but their performance in long
  context is limited by the expressive power of their hidden state. We propose a
  new class of sequence modeling layers with linear complexity and an expressive
  hidden state. The key idea is to make the hidden state a machine learning model
  itself, and the update rule a step of self-supervised learning. Since the
  hidden state is updated by training even on test sequences, our layers are
  called Test-Time Training (TTT) layers. We consider two instantiations:
  TTT-Linear and TTT-MLP, whose hidden state is a linear model and a two-layer
  MLP respectively. We evaluate our instantiations at the scale of 125M to 1.3B
  parameters, comparing with a strong Transformer and Mamba, a modern RNN. Both
  TTT-Linear and TTT-MLP match or exceed the baselines. Similar to Transformer,
  they can keep reducing perplexity by conditioning on more tokens, while Mamba
  cannot after 16k context. With preliminary systems optimization, TTT-Linear is
  already faster than Transformer at 8k context and matches Mamba in wall-clock
  time. TTT-MLP still faces challenges in memory I/O, but shows larger potential
  in long context, pointing to a promising direction for future research.
paper_authors._id: 668b5e98d358e8fd17762b9a
paper_authors.name: Yu Sun
paper_authors._id: 668b5e98d358e8fd17762b9b
paper_authors.name: Xinhao Li
paper_authors._id: 668b5e98d358e8fd17762b9c
paper_authors.name: Karan Dalal
paper_authors.status: extracted_confirmed
paper_authors.statusLastChangedAt: 2024-07-08T04:11:37.802Z
paper_authors._id: 646ebf2298e8f749fc60ce38
paper_authors.avatarUrl: /avatars/5a76c9343451c24e9697d3165cdc0af6.svg
paper_authors.fullname: Karan Dalal
paper_authors.type: user
paper_authors.user: karansdalal
paper_authors._id: 668b5e98d358e8fd17762b9d
paper_authors.name: Jiarui Xu
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2024-07-08T21:11:59.315Z
paper_authors._id: 64f55c47708362ac9b5f5c45
paper_authors.avatarUrl: /avatars/bd0ba54d6e75dc3662a3cbdec0091770.svg
paper_authors.fullname: Jiarui Xu
paper_authors.type: user
paper_authors.user: jxu605
paper_authors._id: 668b5e98d358e8fd17762b9e
paper_authors.name: Arjun Vikram
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2024-07-08T21:11:33.053Z
paper_authors._id: 63c89f809cf77a4393d2f1ce
paper_authors.avatarUrl: https://cdn-avatars.huggingface.co/v1/production/uploads/1674092377599-noauth.jpeg
paper_authors.fullname: Arjun Vikram
paper_authors.type: user
paper_authors.user: arjvik
paper_authors._id: 668b5e98d358e8fd17762b9f
paper_authors.name: Genghan Zhang
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2024-07-08T21:11:27.105Z
paper_authors._id: 65a76ff1e504d9738d636217
paper_authors.avatarUrl: /avatars/26bf5e3f19057835ee95d72c24904d77.svg
paper_authors.fullname: Genghan Zhang
paper_authors.type: user
paper_authors.user: Genghan
paper_authors._id: 668b5e98d358e8fd17762ba0
paper_authors.name: Yann Dubois
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2024-07-08T21:11:21.048Z
paper_authors._id: 641362754e5305c14f233f08
paper_authors.avatarUrl: /avatars/bce974ef60c507d22702cc7662033b28.svg
paper_authors.fullname: Yann Dubois
paper_authors.type: user
paper_authors.user: YannDubs
paper_authors._id: 668b5e98d358e8fd17762ba1
paper_authors.name: Xinlei Chen
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2024-07-08T21:11:00.420Z
paper_authors._id: 63e58e3a006a775275e59e41
paper_authors.avatarUrl: /avatars/75262a35b27a2ae1939df9118120d99e.svg
paper_authors.fullname: Xinlei Chen
paper_authors.type: user
paper_authors.user: endernewton
paper_authors._id: 668b5e98d358e8fd17762ba2
paper_authors.name: Xiaolong Wang
paper_authors._id: 668b5e98d358e8fd17762ba3
paper_authors.name: Sanmi Koyejo
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2024-07-08T21:10:33.322Z
paper_authors._id: 64931e7e2da595588288f161
paper_authors.avatarUrl: https://cdn-avatars.huggingface.co/v1/production/uploads/64931e7e2da595588288f161/4jOhJOFsU7RVFMgGk5kO7.jpeg
paper_authors.fullname: Sanmi Koyejo
paper_authors.type: user
paper_authors.user: sanmikoyejo
paper_authors._id: 668b5e98d358e8fd17762ba4
paper_authors.name: Tatsunori Hashimoto
paper_authors._id: 668b5e98d358e8fd17762ba5
paper_authors.name: Carlos Guestrin
paper_title: Learning to (Learn at Test Time): RNNs with Expressive Hidden States
paper_summary: Self-attention performs well in long context but has quadratic complexity.
  Existing RNN layers have linear complexity, but their performance in long
  context is limited by the expressive power of their hidden state. We propose a
  new class of sequence modeling layers with linear complexity and an expressive
  hidden state. The key idea is to make the hidden state a machine learning model
  itself, and the update rule a step of self-supervised learning. Since the
  hidden state is updated by training even on test sequences, our layers are
  called Test-Time Training (TTT) layers. We consider two instantiations:
  TTT-Linear and TTT-MLP, whose hidden state is a linear model and a two-layer
  MLP respectively. We evaluate our instantiations at the scale of 125M to 1.3B
  parameters, comparing with a strong Transformer and Mamba, a modern RNN. Both
  TTT-Linear and TTT-MLP match or exceed the baselines. Similar to Transformer,
  they can keep reducing perplexity by conditioning on more tokens, while Mamba
  cannot after 16k context. With preliminary systems optimization, TTT-Linear is
  already faster than Transformer at 8k context and matches Mamba in wall-clock
  time. TTT-MLP still faces challenges in memory I/O, but shows larger potential
  in long context, pointing to a promising direction for future research.
paper_ai_keywords: self-attention RNN sequence modeling hidden state machine learning model self-supervised learning Test-Time Training (TTT) TTT-Linear TTT-MLP perplexity context Transformer Mamba wall-clock time memory I/O