Learning to (Learn at Test Time): RNNs with Expressive Hidden States

By Yu Sun, Xinhao Li, Karan Dalal, Jiarui Xu, Arjun Vikram, Ge Zhang, Yann Dubois, Xinlei Chen, Xiaolong Wang, Oluwasanmi Koyejo and 2 more. International Conference on Machine Learning, 2024

311 citations recorded by the source corpus.

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.

Subfields: Artificial Intelligence

View Paper: https://www.semanticscholar.org/paper/35ab93f41115e860bee5e202b71061addfd0fd5d

Citations in the corpus

• Parallelizing Linear Transformers with the Delta Rule over Sequence Length (2024, unknown): “Finally, we remark that delta rule is closely related to meta or online learning via gradient descent [70, 38], which has been revisited in recent works like Longhorn [56] and TTT [107].”