multigeom.diffusion.train

Stage 3: Train a conditional DDPM on TinyImageNet, conditioned on the synthetic responses produced in Stage 2.

Training recipe (paper appendix "Diffusion model architecture and training"):
  - DDPM linear β schedule, T = 1000, ε-prediction objective.
  - AdamW with lr 1e-4, weight decay 1e-4, cosine schedule with 500-step warmup.
  - Batch size 128 in bfloat16, gradient clipping at norm 1.0.
  - 200 epochs (the checkpoint used in all figures).
  - Conditioning vector r̃ = r + σ ξ (σ = 0.22) injected per-step.
  - Classifier-free guidance dropout at rate 0.1: with that probability,
    the conditioning vector is zeroed out so the network learns both
    p(ε | x_t, r̃) and p(ε | x_t, ∅).

Usage:
    python -m multigeom.diffusion.train \
        --config diffusion \
        --responses runs/responses/inception/V1/train_responses.pt \
        --out-dir runs/diffusion/inception/V1

Functions

cosine_with_warmup(step, warmup, total)	Multiplicative LR factor: linear warmup, then cosine decay to 0.
main([argv])	Train the diffusion model from a responses file (command-line entry point).
parse_args([argv])	Command-line arguments of python -m multigeom.diffusion.train.
train_diffusion(cfg, image_dataset, ...[, ...])	Train the conditional diffusion model of an encoder, and write final_model.pt (rewritten after every epoch), train_log.json, and last.pt (to resume an interrupted run with resume=True) to out_dir.

Classes

TinyImageNetWithResponses(*args, **kwargs)

Pair TinyImageNet images with the precomputed synthetic response vectors.