Sakana AI teaches transformers how to forget

Humans stay sane by instantly forgetting 99% of the useless sensory noise around them. Transformers do the exact opposite: they hoard every single token, comma, and whitespace in their working memory like digital packrats, which is why extended context windows get painfully expensive and slow.

Tokyo-based research lab Sakana AI just unveiled a biologically inspired fix. Backed by Japan's METI GENIAC supercomputing initiative, they published a paper titled An Evolved Universal Transformer Memory, introducing a system called Neural Attention Memory Models (NAMMs).

Evolving a brain that knows how to dump trash

Previous efforts to trim memory contexts, like H2O or L2, relied on rigid human heuristics—for instance, keeping only the most recent tokens or heavy hitters. That usually saved memory at the direct expense of task accuracy.

Sakana AI took a different route by treating retention as a learned neural classifier. The problem is that choosing whether to keep or permanently drop a token is a binary, non-differentiable operation, meaning standard gradient descent fails completely.

So they turned to evolutionary optimization. By repeatedly mutating and selecting network parameters across rounds of trial and error, they evolved small classifiers capable of deciding exactly what to prune without needing gradients.

Technically, NAMMs turn the attention values of each token into spectrograms, compress them using an exponential moving average, and feed those compact summaries into the evolved network to score whether each token deserves to stay.

Zero-shot mastery across code, video, and robots

Because NAMMs feed strictly on attention matrices—components native to every Transformer—the learned pruning strategy exhibits wild transferability. Sakana AI evolved NAMMs entirely on a Llama 3 8B text model, yet the same weights transferred zero-shot to a massive 70B model, a video foundation model (Llava Next Video) to prune redundant frames, and a Decision Transformer in reinforcement learning to prune suboptimal actions.

Inspecting what gets pruned reveals fascinating behavior. In early layers, NAMMs preserve global context like task instructions. In deeper layers, once that context is assimilated, they dump the preamble and focus strictly on local details.

When reading code, NAMMs slice away whitespace, comments, and boilerplate text. In natural language, they strip out grammatical filler mid-sentence while keeping named entities intact. Sakana AI open-sourced the training code alongside ChouBun, a new Japanese long-context benchmark covering 36 tasks.

Why it matters

For anyone hosting long-context models, inference latency and memory footprints are the primary cost bottlenecks. NAMMs demonstrate that you don't need to rebuild foundational architectures from scratch or settle for dumb heuristic pruning just to save VRAM. The fact that Sakana AI made the code public means builders can immediately inspect and experiment with this selective pruning workflow.

Industry-wide, it challenges the brute-force arms race around massive KV caches. Proving that an attention-pruning policy learned on text can seamlessly drop redundant video frames or robotic actions without any retraining points to a genuinely universal pattern in how Transformers prioritize information.

We think Sakana AI played to its unique strengths here. Using evolutionary strategies to bypass non-differentiable bottlenecks is clever, and leveraging attention spectrograms is genuinely elegant. However, the catch looks obvious: evolutionary search is computationally voracious and relied on heavy government supercomputing support. While running the pre-evolved NAMM is cheap, evolving custom pruning models for niche architectures remains out of reach for small teams.