---
license: mit
library_name: pytorch
tags:
  - diffusion-models
  - associative-memory
  - hopfield-networks
  - memorization
  - generalization
  - cifar10
---

# dm-am-cifar10-unet96

Trained **unet96** DDPM diffusion models on **cifar10**, 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, 11.4 GiB total, spanning K = 2 to 50,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

```python
from huggingface_hub import hf_hub_download
import torch

# the model trained on K=2 samples
path = hf_hub_download("lemoncmd/dm-am-cifar10-unet96", "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

```bibtex
@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}
}
```