Note that this works within a single sequence of tokens. It might be consistent with "episodic memory" metaphor if we consider a particular transformer run as its experience.
But this might be very different from what people expect from "memory" - i.e. ability to learn vast amounts of information and retrieve it as necessary.
This is more like a refinement of transformer attention: instead of running attention over all tokens (which is very expensive as it's quadratic), it selects a subset of token spans and runs fine-grained attention only on those. So it essentially breaks transformer attention into two parts - coarse-grained (k-NN over token spans) and fine-grained (normal).
It might be a great thing for long-context situations. But it doesn't make sense when you want millions of different facts to be considered - making them into long context is rather inefficient.
It would be inefficient if you had to do it from scratch for every query, but if you can do it once as a preprocessing step and reuse the prepared context for many queries, it might start to become more efficient than a shorter context that includes only some documents but has to be reprocessed because it's different every time.
Yes, I think it might be a good solution where you have a context up to 10M of tokens and you do a lot of requests with that context. It might be relevant for agentic stuff which tends to produce long chat logs - especially with some gadgets on top, e.g. some 'episodes' might be completely removed as obsolete.
But I don't think it's a good solution for bigger amounts of data - as in that case it's more beneficial if that can be formed into independent memories.
I think something like Titans explains Gemini's excellent long context performance. That would explain why the Titan team hasn't released the training code or hyperpameters used even though they said in the paper that they would, and why soon after that it came out that DeepMind would be holding off publishing new results for 6 months to avoid giving away competitive advantages.
Interesting. Before there even was attention I was thinking that the episodic memory model offers something that could be very useful for neural nets, so its cool to see people testing that
So, infinite context length by making it compute bound instead of memory bound. Curious how much longer this takes to run and when it makes sense to use vs RAG.
Hi MacsHeadroom, first author here. Thanks for the great questions about compute/memory trade-offs.
The quick take: To give you an example of processing speed, with a 7B model on an NVIDIA V100, EM-LLM processes (or generates) about 326 tokens/sec with a 51.2K context window (which is quite competitive for these old GPUs).
More broadly, EM-LLM is designed to make ultra-long contexts (memory-prohibitive for standard O(n^2) attention) computationally tractable. The Appendix C of our paper https://openreview.net/pdf?id=BI2int5SAC details how: significantly better attention scaling, efficient O(nm) memory formation, and large KV cache management via CPU/disk offloading. While there's a slight per-chunk overhead compared to the simplest retrieval methods initially, the crucial part is our ability to handle sequences at scales infeasible for full-context models. For instance, we're successfully using 8B models with 10M token contexts on a single GPU without prohibitive delays.
Regarding RAG in particular, EM-LLM often shows significant gains on tasks needing deep understanding of a single, long, coherent context. A key reason is that EM-LLM allows each layer to retrieve and integrate relevant information from different "episodes" of the context independently, offering more nuance than a typical single RAG step, for similar overall resource use.
But this might be very different from what people expect from "memory" - i.e. ability to learn vast amounts of information and retrieve it as necessary.
This is more like a refinement of transformer attention: instead of running attention over all tokens (which is very expensive as it's quadratic), it selects a subset of token spans and runs fine-grained attention only on those. So it essentially breaks transformer attention into two parts - coarse-grained (k-NN over token spans) and fine-grained (normal).
It might be a great thing for long-context situations. But it doesn't make sense when you want millions of different facts to be considered - making them into long context is rather inefficient.