dm-am-mnist-unet128

Trained unet128 DDPM diffusion models on mnist, from the paper Memorization to Generalization: Emergence of Diffusion Models from Associative Memory.

Bao Pham, Gabriel Raya, Matteo Negri, Mohammed J. Zaki, Luca Ambrogioni, Dmitry Krotov

• Paper: https://arxiv.org/abs/2505.21777
• Code: https://github.com/Lemon-cmd/Diffusion-Models-and-Associative-Memory

What this contains

38 models, 13.9 GiB total, spanning K = 2 to 60,000.

Each file is named <K>.pt, where K is the size of the training set the model was trained on -- not a training step. Every model was trained for the same number of iterations; K is the axis the paper sweeps to move the model through its three regimes:

Regime	Roughly	Behaviour
Memorization	small K	Each training sample gets its own attractor
Spurious	intermediate K	Emergent attractors that are not training data -- the first signs of generative ability
Generalization	large K	Attractors correspond to novel, coherent samples

Sorting the files numerically walks that transition.

Checkpoint format

Each .pt is a torch.save dict:

Key	Contents
model	Model state_dict, saved from a DistributedDataParallel wrapper (keys carry a module. prefix)
ema	EMA weights, same parameters without the module. prefix
opt	Optimizer state (state, param_groups)
args	Full training config Namespace, including train_size (matches the filename)
iterations	Configured training iterations (identical across files)

Optimizer state is included, so these are resume-capable, not inference-only.

Loading

from huggingface_hub import hf_hub_download
import torch

# the model trained on K=2 samples
path = hf_hub_download("lemoncmd/dm-am-mnist-unet128", "2.pt")
ckpt = torch.load(path, map_location="cpu", weights_only=False)

ema = ckpt["ema"]                 # EMA weights, used for sampling in the paper
model = {k.removeprefix("module."): v for k, v in ckpt["model"].items()}

Unpickling ckpt["args"] needs the training repo's config classes importable (simple_parsing plus parse_utils.py from the code repo). Use weights_only=True to read only tensors.

MANIFEST.tsv lists every file with its K and byte size.

Citation

@inproceedings{Pham2025MemorizationTG,
  title  = {Memorization to Generalization: Emergence of Diffusion Models from Associative Memory},
  author = {Bao Pham and Gabriel Raya and Matteo Negri and Mohammed J. Zaki and Luca Ambrogioni and Dmitry Krotov},
  year   = {2025},
  url    = {https://arxiv.org/abs/2505.21777}
}