An Evolved Universal Transformer Memory

Sakana AI

Sakana AI developed Neural Attention Memory Models (NAMMs), learnable memory systems that enable transformers to selectively retain or discard tokens based on attention patterns, improving both performance and efficiency. The NAMMs were trained on Llama 3 8B using evolutionary optimization and evaluated on three long-context benchmarks totaling 36 tasks, consistently outperforming prior hand-designed methods like H₂O and L₂. The system transfers zero-shot to other transformer architectures and modalities including video and reinforcement learning without retraining, allowing models to focus on critical information for improved performance across diverse tasks.