The memory works by tracking two kinds of "surprise" - immediate surprise (how unexpected is the current token?) and accumulated surprise (what patterns of unexpected things have we been seeing?). It uses this to decide what's worth remembering and what can be forgotten. What's clever is they formulated this as a gradient descent problem that can run efficiently in parallel despite being inherently sequential.

The really interesting part is how it integrates with the main model - they tried three approaches but the most effective was using the memory as additional context tokens alongside the input. This lets the attention mechanism figure out for itself when to use the memory versus the immediate context. And because the memory tokens are injected both at test-time and during training, the memory model and main model are trained together despite the memory model being unfrozen for each inference.

In practice, this lets them handle sequences over 2M tokens long while outperforming traditional transformers, even matching GPT-4 on some long-context reasoning tasks with far fewer parameters. It's a neat example of combining classical ideas about online learning with modern deep learning architectures.

The code isn't released yet, but the paper suggests the implementation is relatively straightforward since it builds on standard gradient descent mechanics. It'll be interesting to see if this approach influences the next generation of open source LLMs. I'm sure we will see implementations very soon even though it may take some time for open source models to be trained using this new architecture.

I'm very excited to know whether Gemini 2.0 1206 Experimental is using this new architecture. I suspect it is.