Transcription
Google Research has just released a new paper that could be the successor to the groundbreaking "Attention is All You Need" paper, which initiated this entire new wave of AI. This paper introduces an approach called Titans, which offers a way to provide models with memory that resembles human memory, particularly long-term memory, during inference time through an intriguing surprise mechanism.
This is quite complex, but I will try to simplify it for you. The paper is titled "Titans: Learning to Memorize at Test Time," and it comes from Google Research.
The abstract of the paper outlines the issue with Transformers: the context window is limited due to a significant penalty when increasing the context window based on the Transformers architecture. But what if there were no such limitation? What if you could have an infinite number of tokens in a context window while still performing well? That is what Titans aims to achieve.
The paper describes that this more accurate modeling of dependencies comes with a quadratic cost, which limits the model to a fixed-length context. When considering the context window of major models, like GPT-4 with 128k tokens or even the largest context window with Gemini at 2 million tokens, it’s substantial. However, as our needs grow, we require ways to exceed that context window.
The paper explains how Titan-based models can learn to have long-term memory and determine what to focus on and where to direct their attention. Experimental results in language modeling, common sense reasoning, genomics, and time series tasks show that Titans are more effective than Transformers. They can also effectively scale to a context window size larger than 2 million tokens with higher accuracy in needle-in-a-haystack tasks compared to baselines. If this is accurate, it represents a significant advancement.
In the introduction, the authors discuss Transformers and the importance of the attention mechanism. Transformers, which are purely attention-based architectures, have been established as state-of-the-art models in sequence modeling. Most language models you use today, from OpenAI to Anthropic to LLaMA and DeepSeek, are based on the Transformers architecture and the attention mechanism. They are state-of-the-art due to their in-context learning and ability to learn at scale.
However, a major drawback of the Transformers architecture is its quadratic time and memory complexity concerning context length. It struggles with longer context lengths; the more information you provide during inference time or in the prompt, the worse it performs.
This is relative, as 2 million tokens is quite a lot—you can load entire videos with that. But as our need to input more information into these models grows, this problem becomes increasingly significant. In complex real-world tasks like video understanding and long-term time series forecasting, the context window can become extremely large, making the applicability of Transformers challenging in these downstream tasks.
That is the problem they are trying to solve with Titans. Titans aim to model their architecture more closely to how the human brain works. The human brain has multiple types of memory: short-term memory, long-term memory, and meta-memory. Understanding how these different memory types work together and independently is crucial to how the human brain functions, and current Transformers models do not achieve this.
Titans aim to provide multiple types of memory and allow these different types to work together. The authors argue that in an effective learning paradigm, similar to the human brain, there are distinct yet interconnected memory modules, each responsible for a component crucial to the learning process.
Now, they take a high-level view of memory and its importance. Memory is a fundamental mental process and an inseparable component of human learning. Learning and memory are two different things but are interconnected. Without memory, humans and animals would be limited to basic reflexes and stereotyped behaviors.
Most existing architectures consider memory as a neural update caused by an input—seeing, hearing, or smelling something forms a memory. Learning is defined separately as a process for acquiring effective and useful memory given an objective.
In this paper, they aim to answer a few questions:
1. What constitutes a good structure for memory?
2. What is a proper memory update mechanism?
3. What is a good memory retrieval process?
4. How to design an efficient architecture that incorporates different interconnected memory modules (short-term, long-term, meta-memory)?
5. Is a deep memory module needed to effectively store long-term memories?
One of the drawbacks of Transformers is that they struggle with very long-term memory. What if you could integrate that memory directly into the model?
The important part of this paper is that they aim to answer the above five questions by designing a long-term neural memory module that can efficiently and effectively learn to memorize at test time. I want to emphasize "test time," meaning not during pre-training, but when the model is actually running and responding to prompts. This is when they want to provide new memory to the model, allowing it to learn how to memorize and store data into its parameters at test time.
If this sounds familiar, it should. I covered a paper called "Test Time Training," which allowed a model to learn based on its prompt, updating its parameters during prompting and inference time.
The most interesting part, in my opinion, is that they designed this memory module so that an event that violates expectations—essentially, a surprise—is more memorable.
Think about it: when you're doing something boring or repetitive, those actions become so ingrained that you zone out. For example, when driving, you might arrive at a destination without remembering the drive itself. This is different from when something surprises you, like getting cut off in traffic or witnessing an accident. You will remember those surprising moments because they are not typical experiences.
This surprise mechanism is baked into the architecture of how memory works, which is fascinating. When the model is surprised, it recognizes that it needs to memorize that event. They measure the surprise of an input and present a decaying mechanism that considers the proportion of memory size and the amount of data surprise, resulting in better memory management.
Essentially, when something surprises you, it has a high memory factor initially, but over time, it learns to give that less attention or priority. Think about a surprising event in your life: at the moment, it is vivid in your memory, but as time passes, you start to forget details, and that memory becomes less significant.
This decay mechanism generalizes the forgetting mechanism in modern recurrent models.
So, what is the Titans architecture? It incorporates this type of learning and memory into an AI model. An important remaining question is how to effectively and efficiently integrate memory into a deep learning architecture.
They present Titans as a family of deep models consisting of three hyper heads: core, long-term, and persistent memory.
Think of core memory as short-term memory, responsible for the main flow of data. Long-term memory stores and remembers memories over time, while persistent memory consists of learnable but data-independent parameters that encode knowledge about a task.
They provide three different variants of the Titan architecture, each with different trade-offs. They incorporate memory as context, a layer, and a gated branch.
The authors observe that their Titan architecture outperforms all modern recurrent models and their hybrid variants across a comprehensive set of benchmarks. Titans can scale to a context window size larger than 2 million tokens, which is the current state-of-the-art limit.
The next section discusses learning to memorize at test time. Remember, "test time" refers to inference time when you prompt a model, and it generates a response. This process needs to happen quickly.
They present a neural long-term memory module, which is a meta-model that learns to memorize at test time. Long-term memory encodes the abstraction of past history into its parameters.
It’s like long-term memory in humans; we don’t remember every detail of every memory. Instead, we have an abstraction—a rough picture of what happened over time. However, memorization has often been seen as undesirable in neural networks because it limits the model's generalization.
Knowing what to memorize is crucial for these models, as it also raises privacy concerns. If a model memorizes everything, it may disclose private information and perform poorly at test time.
Now, my favorite part: the surprise metric. An event that violates expectations is called surprise. The larger the gradient with respect to the input, the more different the input data is from past data.
However, the surprise metric can lead to missing important information that follows a surprising moment. If too much attention is given to a surprising moment, subsequent events may be overlooked.
To improve the surprise metric, they break it into past surprise and momentary surprise. Past surprise measures the surprise of very recent events, while momentary surprise measures the surprise of incoming data.
They also need a way for the model to forget; it can’t just memorize everything. A forgetting mechanism is essential when dealing with large sequences of millions of tokens.
To this end, they use an adaptive forgetting mechanism that allows the memory to forget information that is no longer needed. This mechanism considers surprise and available memory to decide what to forget.
In the next section, they discuss different ways to incorporate memory, which come with trade-offs.
First, memory as context: think of it like having a personal assistant in a meeting who takes detailed notes of past discussions. Long-term memory whispers relevant information when needed, helping you combine past knowledge with current discussions to make decisions.
Next, memory as gate: imagine having two advisors in your head—one focused on the present and another drawing from years of experience. A gatekeeper decides how much to listen to each advisor.
Finally, memory as a layer processes information through different layers, with each layer representing a type of memory. For example, the first filter might process everything through long-term memory, while the second filter looks at immediate context through attention. Each layer refines the information before passing it to the next.
As I mentioned, there are trade-offs. Memory as context is best for tasks requiring detailed historical context. Memory as gate is more flexible, switching between short and long-term focus, while memory as a layer is the most efficient but slightly less powerful.
The performance across various benchmarks shows that Titan models consistently outperform other architectures.
In the needle-in-a-haystack test, which assesses the ability to remember and retrieve information from long context windows, Titans maintain consistent performance as context length increases, while other models drop off significantly.
In conclusion, they present a neural long-term memory that learns to memorize at test time. It is a recurrent model that adaptively memorizes tokens that are more surprising and close to surprising tokens.
Their experimental evaluation across diverse tasks validates that Titans are more effective than Transformers and modern linear recurrent models, especially for long contexts.
That’s it! This paper proposes a fascinating way to enhance models' long-term memory using a surprise mechanism. Congratulations to the authors of this paper! I will link the paper in the description below. If you enjoyed this video, please consider liking and subscribing, and I’ll see you in the next one.