Self-attention mechanisms are good at handling extended contexts since they can capture associations across sequences. However, they are computationally expensive, i.e., they have high quadratic complexity, so as the length of the sequence increases, the amount of time and memory required also increases. Recurrent Neural Networks (RNNs), on the other hand, have linear complexity and are therefore computationally more efficient. However, RNNs perform poorly in long settings due to the constraint on their hidden state that all data must be contained in a fixed-size representation.
To overcome these limitations, a team of researchers from Stanford University, UC San Diego, UC Berkeley, and Meta AI proposed a unique class of sequence modeling layers that combines more expressive hidden states with the linear complexity of RNNs. The main concept is to adopt a self-supervised learning step as the update rule to transform the hidden states into a machine learning model. This means that the hidden states are updated by efficiently training the input sequences even during the testing phase. These levels are called training-at-test (TTT) layers.
TTT-Linear and TTT-MLP are two different varieties of TTT layers introduced. The hidden state of TTT-MLP is a two-layer multilayer perceptron (MLP), while the hidden state of TTT-Linear is a linear model. The team tested the performance of these TTT layers against a robust Transformer model and a state-of-the-art RNN, Mamba, evaluating them on models with parameters ranging from 125 million to 1.3 billion.
Evaluation shows that both TTT-Linear and TTT-MLP perform as well as or better than the baseline. Like the Transformer, the TTT layer keeps getting smaller conditional on additional tokens. Perplexity is a metric that evaluates how well the model can predict sequences. This is a big advantage that the TTT layer takes advantage of the extended context well. Meanwhile, Mamba stops improving at 16,000 tokens.
After some preliminary optimizations, TTT-Linear matches Mamba in wall-clock time, a measure of the actual time that elapses during processing, and outperforms Transformer in speed for sequences of up to 8,000 tokens. Although it has a good chance of managing long contexts, TTT-MLP still has issues with memory I/O operations.
The team summarises their main contributions as follows:
- We introduce a unique class of sequence modeling layers called training-at-test (TTT) layers, where the model updated by self-supervised learning serves as the hidden state. This view opens up new avenues in sequence modeling research by integrating the training loop into the forward pass of the layer.
- A direct instantiation of the TTT layer, called TTT-Linear, was introduced and the team showed that it performed better than both Transformers and Mamba in evaluations of model sizes from 125 million to 1.3 billion parameters, suggesting that the TTT layer has the ability to improve the performance of sequence models.
- The team also created mini-batch TTT and dual formats to make the TTT layer more hardware-efficient, which makes TTT-Linear a useful building block for large language models. These optimizations make integrating the TTT layer into real-world applications more feasible.
Please check paper. All credit for this work goes to the researchers of this project. Also, don't forget to follow us: twitter.
participate Telegram Channel and LinkedIn GroupsUp.
If you like our work, you will love our Newsletter..
Please join us 46k+ ML Subreddit
Tanya Malhotra is a final year undergraduate student from the University of Petroleum and Energy Studies, Dehradun, doing a BTech in Computer Science Engineering with specialisation in Artificial Intelligence and Machine Learning.
She is an avid fan of Data Science and has strong analytical and critical thinking skills with a keen interest in learning new skills, group leadership and managing organized work.
