An Evolved Universal Transformer Memory

Sakana AI's latest project tackles a problem anyone who has run a long conversation with a language model has felt: the thing chokes on its own context. Transformers store and process every past token indiscriminately, which gets expensive and sloppy as inputs stretch out. Sakana's answer, described in a new paper called An Evolved Universal Transformer Memory, is a set of small neural network classifiers called Neural Attention Memory Models, or NAMMs, that decide token by token whether to keep something in memory or let it go.

The tricky part is that keep-or-drop is a binary call, and binary calls don't play nice with the gradient-based training that powers most of deep learning. So Sakana turned to evolution instead — mutating and selecting NAMM parameters based on which versions produced the best language modeling results, no gradients required. The models lean on attention matrices, a component present in every transformer layer, which encode how much each token matters relative to the others. Because that structure is common across architectures, a single trained NAMM can be dropped into a different layer, or even a completely different model, without retraining.

And that's the part Sakana leans on hardest. They trained NAMMs on a Llama 3 8B base model and tested the combination on LongBench, InfiniteBench, and a benchmark they built themselves called ChouBun, which targets long-context Japanese text — a gap they say existing benchmarks focused on English and Chinese don't cover. Across 36 tasks spanning natural language and coding problems, NAMMs improved performance over the base Llama 3 8B and, as a side effect nobody explicitly optimized for, shrank the context size processed per layer. Two prior hand-designed memory methods, H₂O and L₂, also cut context size, but Sakana says those gains came with performance trade-offs, since those methods were built to preserve full-context performance rather than beat it.

The transfer results are where things get more interesting. NAMMs trained purely on language got dropped into a much larger Llama 70B model, into Llava Next Video for computer vision work, and into the Decision Transformer for reinforcement learning tasks — all zero-shot, no retraining. In each case the models still found useful things to forget: redundant video frames, weak action choices. Digging into what gets pruned, Sakana found early transformer layers hang onto global information like task instructions and keywords, while later layers let that go and focus on local detail. Coding tasks saw whole chunks of whitespace, comments, and boilerplate discarded, while natural language tasks lost grammatical filler mid-sentence, leaving names and core concepts intact.

Sakana frames this as an early step rather than a finished product, floating the idea of training transformers directly on top of NAMMs from the start, or alternating between learning and evolution across future model generations — a nod to how their own team likes to describe biological memory as an iterative, evolved process. The work was backed by Japan's GENIAC supercomputing grant program, and the training code plus the ChouBun benchmark are both public now.