Research map / Architectures and efficiency / Attention and transformers

An Evolved Universal Transformer Memory

ICLR 2025 · Architectures and efficiency › Attention and transformers · approach: Attention head and block redesign · 1 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Evolves a tiny (~4k parameter) learned KV-cache eviction network (NAMM) with CMA-ES on top of a frozen Llama 3 8B, improving long-context performance while shrinking the cache, and transfers zero-shot to other transformers and modalities.

Problem

Hand-designed KV cache eviction heuristics reduce memory cost of transformers but typically trade off performance; the paper asks whether a learned memory management policy can improve both performance and efficiency.

Approach

Neural Attention Memory Models (NAMMs) score each KV cache token from a spectrogram (STFT with Hann window, EMA-reduced) of its column in the attention matrix, using a small backward-masked attention network (BAM) and evolved with CMA-ES; because inputs depend only on attention values, one shared network applies to any attention layer and transfers zero-shot across architectures and modalities.

Key results

• On LongBench, NAMM reaches 29.33 vs 28.86 for the full-context Llama 3 8B base (1.11 normalized) with about 25% of the cache (8155 tokens vs 32768).
• On InfiniteBench, NAMM raises overall average from 1.05 to 11.00 while retaining about 40% of the cache.
• Zero-shot transfer to LLaVA-NeXT-Video 7B yields 44.38 vs 43.84 base on video benchmarks with 72% cache size; on D4RL with Decision Transformer it gives 31.73 vs 29.04 (1.09 normalized).

Measured effects reported in the paper

Surprising findings

• Evolved memory policy, trained only on language tasks, improves over full-context performance on video and offline RL despite never seeing those modalities.
• The learned policy forgets redundant video frames rather than the language tokens of the final prompt.

Limitations

• Evolution is computationally expensive relative to gradient-based training, and results rely on a small set of three training tasks.
• On Llama 3 70B zero-shot transfer, NAMM does not improve overall full-context performance (34.70 vs 35.22).

Experimental setup

Datasets: LongBench InfiniteBench ChouBun LongVideoBench MLVU D4RL

Benchmarks: LongBench InfiniteBench ChouBun LongVideoBench MLVU D4RL

Metrics: normalized performance relative to full-cache base average cache size

Baselines: H2O L2 FastGen full-context base model NAMM (MLP)

Builds on: CMA-ES NTK-aware positional interpolation H2O FastGen

Closest papers by meaning

• Rethinking Attention: Exploring Shallow Feed-Forward Neural Networks as an Alternative to Attention Layers in Transformers (Student Abstract) (AAAI 2024 · similarity 0.89)
• Understanding Factual Recall in Transformers via Associative Memories (ICLR 2025 · similarity 0.89)
• Beyond Attention: Breaking the Limits of Transformer Context Length with Recurrent Memory (AAAI 2024 · similarity 0.89)
• Transformers Learn Faster with Semantic Focus (NeurIPS 2025 · similarity 0.88)
• The underlying structures of self-attention: symmetry, directionality, and emergent dynamics in Transformer training (ICML 2025 · similarity 0.88)
• What Can Transformer Learn with Varying Depth? Case Studies on Sequence Learning Tasks (ICML 2024 · similarity 0.88)