TLDR: Memory Mosaics v2 scales associative memory networks to large language model sizes (Llama-8B scale) and trains them on a trillion tokens. It introduces architectural changes (adaptive bandwidth, gated key extractor, 3-level memory) and demonstrates superior performance over transformers in new-knowledge storage and in-context learning, even with less training data. It matches transformers on training-knowledge storage, challenging the notion that simply increasing data is the only path to advanced capabilities.
In the rapidly evolving field of artificial intelligence, researchers are constantly seeking ways to enhance the capabilities of large language models. A recent paper introduces “Memory Mosaics v2,” a significant advancement in networks of associative memories, demonstrating impressive compositional and in-context learning abilities at a scale comparable to large language models like Llama-8B.
Memory Mosaics, first introduced by Zhang et al. in 2025, are essentially networks of associative memories designed to offer a more transparent way to understand how models learn and compose information. While initial versions showed promise on smaller networks and synthetic datasets, the key question was whether these favorable properties would hold when scaled up to real-world datasets and larger model sizes.
Scaling Up Memory Mosaics
This new work addresses that question directly. The researchers successfully scaled Memory Mosaics to a 10-billion parameter size, training them on a massive one trillion tokens of real-world data. This scaled-up version, dubbed “Memory Mosaics v2,” incorporates several key architectural modifications to achieve its enhanced performance. These include an adaptive bandwidth for associative memory, a gated time-variant key feature extractor, and a novel 3-level memory design.
The 3-level memory design is particularly interesting. It replaces the contextual associative memory found in earlier Memory Mosaics with two distinct associative memories: a short-term memory for recent tokens and a long-term memory for earlier tokens. This separation, inspired by observations of how attention scores behave in Memory Mosaics, allows the model to handle both immediate context and more distant information efficiently. Additionally, a persistent memory, implemented using dense neural networks, stores global knowledge learned during training.
Evaluating Performance Across Key Dimensions
To thoroughly assess Memory Mosaics v2, the researchers introduced three evaluation dimensions:
1. Persistent-knowledge storage and retrieval: This measures the model’s ability to store and recall knowledge from its training dataset. On common language benchmarks, Memory Mosaics v2 performed comparably to traditional transformer models, indicating its effectiveness in retaining learned information.
2. New-knowledge storage and retrieval: This dimension evaluates the model’s capacity to store and retrieve new information encountered during testing, which is crucial for learning new tasks. In challenging “multi-unrelated-documents question-answering” tasks, Memory Mosaics v2 significantly outperformed transformers, especially at longer context lengths (up to 32k tokens). This suggests a superior ability to absorb and recall novel information.
3. In-context Learning: This directly assesses the model’s ability to learn new tasks with minimal examples and less prior human knowledge. Memory Mosaics v2 showed a remarkable advantage over transformers in multiclass classification tasks, consistently improving performance with more demonstration examples. Transformers, in contrast, often struggled to maintain performance or even degraded with more demonstrations, particularly when dealing with anonymous labels (where semantic cues are removed).
Also Read:
- Unlocking Next-Gen AI: How 2-Simplicial Attention Boosts Language Models with Less Data
- Unlocking Smarter AI: How Large Language Models Are Learning to Reason on a Budget
The Data vs. Innovation Debate
A compelling aspect of this research is its analysis of the “risk-return trade-off” in scaling AI models. The paper compares Memory Mosaics v2 (trained on 1 trillion tokens) with transformers trained on various amounts of data, including a massive 8 trillion tokens. The findings suggest that simply increasing training data for transformers does not easily replicate the new-task learning capabilities of Memory Mosaics v2. Even with eight times more training data, transformers still lagged behind Memory Mosaics v2 in new-knowledge storage and, notably, performed worse on challenging anonymous label in-context learning tasks.
This highlights a crucial point: while more data and computation can yield improvements, innovative architectural designs, like those in Memory Mosaics v2, can offer a more efficient path to superior performance in specific, critical areas like new-task learning. The paper suggests that relying solely on the “more data, more computation” approach might face limitations, including the eventual scarcity of data itself.
For a deeper dive into the technical details and full results, you can read the full research paper here.
The development of Memory Mosaics v2 represents a promising direction for future AI research, focusing on models that are not only powerful but also more adaptable and efficient at learning novel information and tasks.
