Memory Mosaics

Memory Mosaics are networks of associative memories working in concert to achieve a prediction task of interest. The goal of Memory Mosaics is to study how to build AI for the open-world. AI for the open world requires a machine to learn a wide range of new tasks, using fewer examples and less task-specific prior knowledge from human designers.

Latest News

• [2026.10] Release source code of Memory Mosaics+ and Two-stage Memory Mosaics+ as in Open-world AI: a tale of three requirements, supporting consistent in-context learning (for the first time).
• [2026.10] "Open-world AI: a tale of three requirements" frames open-world AI and derives three requirements: rich representation, predictive disentanglement, and inference-time learning.
• [2026.10] Release source code of Memory Mosaics v2 as in Memory Mosaics at scale, featuring FSDP2, FlexAttention block-sparse masks, activation checkpointing, etc.
• [2025.09] "Memory Mosaics at scale" scales Memory Mosaics to 10B parameters and 1T training tokens, accepted by NeurIPS 2025 Oral (top 0.36%).
• [2025.01] "Memory Mosaics" is accepted by ICLR 2025.

Memory Mosaics family

architecture	paper	code
Vanilla Memory Mosaics	Memory Mosaics (ICLR 2025)	vanilla_memory_mosaics/
Memory Mosaics v2	Memory Mosaics at scale (NeurIPS 2025)	memory_mosaics_v2.yaml
Memory Mosaics	Open-world AI: a tale of three requirements (arXiv 2026)	memory_mosaics.yaml
Memory Mosaics+	Open-world AI: a tale of three requirements (arXiv 2026)	memory_mosaics+.yaml
Two-stage Memory Mosaics+	Open-world AI: a tale of three requirements (arXiv 2026)	two-stage_memory_mosaics+.yaml

Layout

configs/          one YAML per model, by family: mosaics/ transformer/ hybrid/ rnn/; debug.yaml
model/            common/ (shared layers), mosaics/, transformer/, rnn/ -- one folder per family
training/         training loop, FSDP/DDP, checkpointing
dataloader/       streaming, sharding, packing and resuming training data
inference/        KV-cache generation
cluster_corpus/   builds the corpus: normalize → embed → cluster → organize
tokenizer/        the Mistral-7B tokenizer (HuggingFace), the SFT chat format
launchers/        SLURM launcher
figure/           the architecture figure
vanilla_memory_mosaics/   the ICLR 2025 code (separate)

Setup

pip install -r requirements.txt
huggingface-cli login     # the Mistral tokenizer (mistralai/Mistral-7B-v0.1) is gated

1. Data

Two ways to build training sequences, chosen by data.loader in a config:

• pretrain: documents packed in random order, the usual pretraining. Used by Memory Mosaics v2 (memory_mosaics_v2.yaml).
• cluster: documents are grouped into semantic clusters (k-means over document embeddings), and each training sequence is packed from the documents of one cluster, i.e. related tasks. Used by the Open-world AI configs.

cluster_corpus/prepare.sh builds the data for both; see cluster_corpus/README.md.

bash cluster_corpus/prepare.sh tiny   --go     # 50k Cosmopedia docs, k=20: minutes, smoke test only
bash cluster_corpus/prepare.sh simple --go     # full Cosmopedia, k=1200 (cluster-scale embedding job)

Output goes under data_out/:

• normalized/<source>/: the documents as JSONL (stage 1), read by the pretrain loader;
• clustered_{tiny_k20,simple_k1200}/: per-cluster JSONL plus metadata_train.json / metadata_test.json (held-out test clusters), read by the cluster loader.

The configs point at data_out/normalized and data_out/clustered_simple_k1200.

2. Training

Smoke test on one GPU:

torchrun --nproc_per_node 1 -m training.train config=configs/debug.yaml

A full run, single node:

torchrun --nproc_per_node 8 -m training.train config=configs/mosaics/two-stage_memory_mosaics+.yaml

Multi-node on SLURM. The launcher snapshots the code into the run directory, writes the resolved config, and requeues the job on failure (it resumes from the latest checkpoint):

python -m launchers.slurm run name=two_stage_mmplus config=configs/mosaics/two-stage_memory_mosaics+.yaml \
    nodes=32 partition=<partition> account=<account>

Runs go to dumps/<name>/ (set DUMP_ROOT for the launcher): checkpoints/step_N/, metrics.jsonl, and base_config.yaml, the exact config that ran.

3. Generation

inference/generator.py generates with a cache for every model in the repository: prefill the prompt, then one token per step (python -m inference.chat --model_path <checkpoint> is an interactive front end).

Under attn_bias_type: doc_causal a prompt may hold several documents separated by EOS, and generation may run across EOS into new documents: p ositions restart, in-document memories read only the current document, cluster memories read the earlier ones, and recurrent states reset -- exactly as in training.

Baselines

The baselines of Open-world AI train with the same code and data as the Mosaics models, so a comparison isolates the memory:

baseline	config	what it is
Transformer	transformer.yaml	softmax attention with RoPE in every layer
Hybrid SWA	hybrid_swa.yaml	layers alternate: sliding-window attention (2k, RoPE) / full attention (no position encoding)
Hybrid DeltaNet	hybrid_deltanet.yaml	layers alternate: gated DeltaNet / full attention (no position encoding)
DeltaNet	deltanet.yaml	a recurrent model: every layer a gated DeltaNet
Transformer+, two-stage Transformer+	in configs/mosaics/	Mosaics-like transformers. See Open-world AI: a tale of three requirements.

Reference

@article{zhang2026openworld,
  title={Open-world AI: a Tale of Three Requirements},
  author={Zhang, Jianyu and Bottou, Léon},
  journal={arXiv preprint arXiv:2610.20782},
  year={2026},
  url={https://arxiv.org/pdf/2610.20782}
}
@article{zhang2026memory,
  title={Memory Mosaics at scale},
  author={Zhang, Jianyu and Bottou, L{\'e}on},
  journal={Advances in Neural Information Processing Systems},
  volume={38},
  pages={38929--38956},
  year={2026}
}
@inproceedings{zhang2025memory,
  title={Memory mosaics},
  author={Zhang, Jianyu and Nolte, Niklas and Sadhukhan, Ranajoy and Chen, Beidi and Bottou, L{\'e}on},
  booktitle={International Conference on Learning Representations},
  volume={2025},
  pages={36412--36433},
  year={2025}
}