Learning to (Learn at Test Time): RNNs with Expressive Hidden States
Fuente:
arXiv
Enregistré dans:
| Auteurs principaux: | , , , , , , , , , , , |
|---|---|
| Format: | Preprint |
| Publié: |
2024
|
| Sujets: | |
| Accès en ligne: | |
| Tags: |
Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
|
| _version_ | 1866914014636802048 |
|---|---|
| author | Sun, Yu Li, Xinhao Dalal, Karan Xu, Jiarui Vikram, Arjun Zhang, Genghan Dubois, Yann Chen, Xinlei Wang, Xiaolong Koyejo, Sanmi Hashimoto, Tatsunori Guestrin, Carlos |
| author_facet | Sun, Yu Li, Xinhao Dalal, Karan Xu, Jiarui Vikram, Arjun Zhang, Genghan Dubois, Yann Chen, Xinlei Wang, Xiaolong Koyejo, Sanmi Hashimoto, Tatsunori Guestrin, Carlos |
| contents | 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 states. We present a practical framework for instantiating sequence modeling layers with linear complexity and expressive hidden states. 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. Similar to Transformer, TTT-Linear and TTT-MLP can keep reducing perplexity by conditioning on more tokens, while Mamba cannot after 16k context. TTT-MLP still faces challenges in memory I/O, but shows larger potential in long context, pointing to a promising direction for future research. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2407_04620 |
| institution | arXiv |
| publishDate | 2024 |
| record_format | arxiv |
| spellingShingle | Learning to (Learn at Test Time): RNNs with Expressive Hidden States Sun, Yu Li, Xinhao Dalal, Karan Xu, Jiarui Vikram, Arjun Zhang, Genghan Dubois, Yann Chen, Xinlei Wang, Xiaolong Koyejo, Sanmi Hashimoto, Tatsunori Guestrin, Carlos Machine Learning Artificial Intelligence Computation and Language 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 states. We present a practical framework for instantiating sequence modeling layers with linear complexity and expressive hidden states. 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. Similar to Transformer, TTT-Linear and TTT-MLP can keep reducing perplexity by conditioning on more tokens, while Mamba cannot after 16k context. TTT-MLP still faces challenges in memory I/O, but shows larger potential in long context, pointing to a promising direction for future research. |
| title | Learning to (Learn at Test Time): RNNs with Expressive Hidden States |
| topic | Machine Learning Artificial Intelligence Computation and Language |
| url | https://arxiv.org/abs/2407.04620 |