Add model card

README.md

@@ -7,39 +7,49 @@ tags:
   - hopfield-networks
   - memorization
   - generalization
-  - training-checkpoints
   - cifar10
 ---
 
 # dm-am-cifar10-unet128
 
-Training-checkpoint series for a **unet128** diffusion model trained on
-**cifar10**, from the paper *Memorization to Generalization: Emergence of Diffusion Models from Associative Memory*.
+Trained **unet128** 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
 
-## Contents
+## What this contains
 
-38 checkpoints, 20.3 GiB total, spanning training steps 2 - 50,000.
-Each file is named `<step>.pt`.
+38 models, 20.3 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 with these keys:
+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` | The full training config `Namespace` used for the run |
-| `iterations` | Total configured training iterations for the run |
+| `args` | Full training config `Namespace`, including `train_size` (matches the filename) |
+| `iterations` | Configured training iterations (identical across files) |
 
-Because optimizer state is included, these are resume-capable training
-checkpoints, not inference-only weights.
+Optimizer state is included, so these are resume-capable, not inference-only.
 
 ## Loading
 
@@ -47,27 +57,27 @@ checkpoints, not inference-only weights.
 from huggingface_hub import hf_hub_download
 import torch
 
+# the model trained on K=2 samples
 path = hf_hub_download("lemoncmd/dm-am-cifar10-unet128", "2.pt")
 ckpt = torch.load(path, map_location="cpu", weights_only=False)
 
-# EMA weights are the ones used for sampling in the paper
-ema = ckpt["ema"]
-
-# Raw model weights, stripped of the DDP prefix
+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). To read only the
-tensors, use `weights_only=True`.
+(`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
-@article{pham2025memorization,
+@inproceedings{Pham2025MemorizationTG,
   title  = {Memorization to Generalization: Emergence of Diffusion Models from Associative Memory},
-  author = {Pham, Bao and Raya, Gabriel and Negri, Matteo and Zaki, Mohammed J. and Ambrogioni, Luca and Krotov, Dmitry},
-  journal = {arXiv preprint arXiv:2505.21777},
-  year   = {2025}
+  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}
 }
 ```