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.

This paper presents a learning system architecture, Memory Mosaics, in which multiple associative memories work in concert to carry out a prediction task of interest. Such systems are closely related to memory networks [Weston et al., 2014, Sukhbaatar et al., 2015] and resemble transformers [Vaswani et al., 2017] despite significant differences. Like transformers, Memory Mosaics possesses some of the disentanglement and compositional capabilities that have long eluded machine learning systems [Lake and Baroni, 2018]. Unlike transformers whose internal mechanism are hard to decipher [Olsson et al., 2022, Bietti et al., 2024], Memory Mosaics achieve these capabilities in comparatively transparent ways. The three main contributions of this work are (a) recognizing and exploiting a similarity between smoothing associative memories and self-attention, (b) identifying and illustrating the predictive disentanglement principle which explains how training decomposes the overall task in interesting ways, and (c) showing that this comparatively transparent architecture matches the performance of decoding transformers on a language modeling task. Section 2 describes the basic architecture and outlines its consequences. Section 3 illustrates the predictive disentanglement principle. Section 4 extends these ideas to fully formed memory mosaics.