Memory Mosaics at scale

TL;DR

This work presents Memory Mosaics v2, a scalable associative-memory-based architecture that extends memory mosaics to llama-8B with one trillion tokens to enhance new-task learning and in-context learning. By introducing adaptive bandwidth, a gated time-variant key extractor, and a 3-level memory, the approach achieves superior new-knowledge storage and in-context learning compared to transformers, often with less data and comparable or lower compute. The authors propose three evaluation dimensions—persistent-knowledge storage, new-knowledge storage, and in-context learning—to precisely diagnose capabilities, and show that Memory Mosaics v2 surpasses transformer baselines on challenging tasks like ruler QA and in-context learning benchmarks, while maintaining competitive or worse on persistent benchmarks. A risk–return analysis suggests smart architectural innovations can outperform mere data scaling, and highlights future directions for long-context efficiency and broader applicability in real-world tasks.

Abstract

Memory Mosaics [Zhang et al., 2025], networks of associative memories, have demonstrated appealing compositional and in-context learning capabilities on medium-scale networks (GPT-2 scale) and synthetic small datasets. This work shows that these favorable properties remain when we scale memory mosaics to large language model sizes (llama-8B scale) and real-world datasets. To this end, we scale memory mosaics to 10B size, we train them on one trillion tokens, we introduce a couple architectural modifications ("Memory Mosaics v2"), we assess their capabilities across three evaluation dimensions: training-knowledge storage, new-knowledge storage, and in-context learning. Throughout the evaluation, memory mosaics v2 match transformers on the learning of training knowledge (first dimension) and significantly outperforms transformers on carrying out new tasks at inference time (second and third dimensions). These improvements cannot be easily replicated by simply increasing the training data for transformers. A memory mosaics v2 trained on one trillion tokens still perform better on these tasks than a transformer trained on eight trillion tokens.

Figures (14)

• Figure 1: Average attention scores of the last token attending previous tokens. Left: Transformer with RoPE position encoding. Right: Memory Mosaics zhang-2025. The (averaged) attention scores in transformer heavily depends on token positions (curly curves), while the attention scores in memory mosaics at far tokens (e.g. position 0 to 450) are almost invariant to positions (flat curves).
• Figure 2: Left: Memory Unit. Right: Memory Mosaics v2 architecture.
• Figure 3: Semantic label in-context learning comparison between memory mosaics v2 and transformer. Memory mosaics v2 significantly outperform transformers on in-context learning with a large margin (more than 10%). Meanwhile, memory mosaics v2 benefits from more demonstration shots (x-axis), unlike transformers.
• Figure 4: Anonymous label in-context learning comparison between memory mosaics v2 and transformers. Memory mosaics v2 significantly outperform transformers on all classification tasks.
• Figure 5: Augmenting transformer with long-short term attention doesn't help in-context learning.