trevorbchen/diffusion_memorization

Diffusion Memorization — Model Checkpoints

Trained VAE and score-model checkpoints for the paper accompanying trevorbchen/diffusion_memorization (code) and trevorbchen/diffusion_memroization (curated results / diverse-1k subset indices).

Status: private. Accompanies a paper currently under review. Please do not share this link or make it public until the authors say otherwise.

Scope: these are the checkpoints behind the paper's headline MLP-based results. The spatial-DiT pilot experiment and its checkpoints are intentionally excluded from this release.

Contents

vae_checkpoints/
├── celeba_resnet_modernloss_d{10,15,20,...,200}/vae.pt   # final CelebA VAE, 14 widths
└── cifar10_resnet_beta005_d{10,20,...,260}/vae.pt         # final CIFAR-10 VAE, 19 widths

results/
├── celeba_diverse1k_bigmlp_sgd_lr001_m08_10k_5m/
│   └── seed{42..N}/d{10,...,200}/
│       ├── last_model.pt    # MLP score-network checkpoint (hidden=1024, depth=5)
│       ├── config.json      # exact training hyperparameters
│       ├── metrics.jsonl    # per-eval-step memorization_fraction_pixel, FID, loss
│       └── runner.log       # training log (includes any NaN/retry history)
└── cifar10_beta005_diverse1k_bigmlp_sgd_lr001_m08_10k_5m/
    └── seed{42..N}/d{20,...,260}/   (same structure)

VAEs: ResNet conv architecture, beta-VAE with free bits and capacity warmup (--arch resnet, see the code repo's standard_vae.py / train_vae_celeba_standard_tar.py for exact args). Score models: SGD (lr=0.001, momentum=0.80), trained against the diverse-1k subset in the dataset repo, 5,000,000 steps, memorization/FID evaluated every 100k steps against 10k generated samples in pixel space.

Memorization ratio: d(gen, NN1_train) / d(gen, NN2_train), memorized if < 1/3 (Bonnaire/Somepalli). See lib/metrics.py in the code repo.

Some d260 CIFAR-10 runs went non-finite under this fixed LR/momentum at high latent width — see runner.log for affected seeds; the paper's clean tables use replacement seeds for those cells rather than a different optimizer configuration (kept consistent across all d for a fair comparison).

License

MIT (matching the code repository).