DASH: Dual-Branch Score Distillation for Guidance-Calibrated Compact Diffusion Models

Parameter compression of class-conditional diffusion models exposes a structural failure mode in existing distillation objectives: supervising only the conditional branch or the composite guided prediction leaves the classifier-free guidance gap underdetermined, admitting degenerate solutions where the student's guidance gap collapses at inference.

DASH resolves this through independent supervision of both score branches, together with TIRT Transfer — a mechanism that initialises the student's per-timestep training curriculum from the teacher's converged weights as a frozen prior.

pip install -r requirements.txt

Tested on Python 3.10, PyTorch 2.0. CIFAR-10 / CIFAR-100 experiments run on a single NVIDIA T4 (16 GB). ImageNet-64 experiments run on TPU and live in a self-contained directory.

DASH/
├── models/unet.py              # ADM-style UNet (teacher 35.8M / student 6.1M)
├── diffusion/
│   ├── schedule.py             # Cosine noise schedule and forward process
│   ├── ddim.py                 # Deterministic DDIM sampler with CFG
│   └── cfg.py                  # Timestep-Adaptive Guidance (TAG) schedule
├── training/
│   ├── tirt.py                 # TIRT module and EMA helper
│   ├── losses.py               # L_im, L_un, L_an, L_DASH
│   └── checkpoint.py           # Checkpoint save / load / prune
├── configs/
│   ├── teacher_cifar10.yaml
│   ├── teacher_cifar100.yaml
│   ├── student_cifar10.yaml
│   └── student_cifar100.yaml
├── train_teacher.py            # Teacher training entry point
├── train_student.py            # DASH distillation entry point
├── evaluate.py                 # FID, IS, rho, cos(Delta), Gap MSE
└── imagenet64_TPU/             # ImageNet-64 experiments
    ├── configs/
    │   ├── teacher.yaml        # 116.5M teacher
    │   └── student.yaml        # 22.3M student, loss weights, TIRT Transfer
    ├── config.py               # YAML loader with validation
    ├── unet.py                 # ch 128/64, nrb 3/2, attention at 16 and 8
    ├── diffusion.py            # cosine schedule, DDIM, TAG
    ├── tirt.py                 # TIRT weighting (Eq. 2)
    ├── losses.py               # DASH losses, composite baseline, calibration
    ├── data.py                 # ImageNet-64 loader
    ├── train_teacher.py
    ├── train_student.py
    └── evaluate.py

Step 1 — Train teacher

python train_teacher.py \
    --config    configs/teacher_cifar10.yaml \
    --data_root /path/to/cifar \
    --out_dir   checkpoints/teacher_cifar10 \
    --download

Step 2 — Distil student

python train_student.py \
    --config       configs/student_cifar10.yaml \
    --data_root    /path/to/cifar \
    --teacher_ckpt checkpoints/teacher_cifar10/ckpt_epoch_0500.pth \
    --out_dir      checkpoints/student_cifar10

For CIFAR-100, substitute cifar10 → cifar100 throughout.

Requires torch-xla. Point --data_root at the directory holding the official downsampled ImageNet-64 files train_data_batch_1 … train_data_batch_10 and val_data (1,281,167 training images, 1000 classes).

cd imagenet64_TPU

# teacher — 130 epochs, global batch 2048, cosine LR from 2e-4, no warmup
python train_teacher.py --config configs/teacher.yaml \
    --data_root /path/to/imagenet64 --out_dir ckpt/teacher

# DASH student
python train_student.py --config configs/student.yaml \
    --data_root /path/to/imagenet64 \
    --teacher_ckpt ckpt/teacher/teacher_ep0130.pth --out_dir ckpt/dash

# evaluation
python evaluate.py --config configs/student.yaml \
    --data_root /path/to/imagenet64 \
    --ckpt ckpt/dash/dash_ep0130.pth \
    --teacher_ckpt ckpt/teacher/teacher_ep0130.pth --out_dir eval/dash

Baselines need no code changes: set train.objective in imagenet64_TPU/configs/student.yaml to composite or no_lun.

config.py asserts on load that snr_gamma = 5.0, TAG = (1.0, 4.0, 5.0), lambda_im = lambda_un = 1.0, and that the student's TIRT schedule is frozen, so a config that would silently change a reported number fails immediately.

python evaluate.py \
    --config       configs/student_cifar10.yaml \
    --ckpt         checkpoints/student_cifar10/ckpt_epoch_0300.pth \
    --teacher_ckpt checkpoints/teacher_cifar10/ckpt_epoch_0500.pth \
    --data_root    /path/to/cifar \
    --out_dir      eval_results/student_cifar10 \
    --seeds        42 123 456

Outputs: FID, IS (mean ± std over generation seeds), guidance gap ratio ρ, directional cosine cos(Δ), Gap MSE, and a results.json summary.

Sampling protocol. Teacher and student are evaluated with the identical sampler: 50-step DDIM under the TAG schedule with w_max = 4.0 (w_min = 1.0, β = 5.0), so the guidance weight runs from 3.77 at low noise to 1.23 at high noise. Reported standard deviations are over three generation seeds from a single trained checkpoint and measure sampling variance, not training variance.

All paper hyperparameters are in configs/ (CIFAR) and imagenet64_TPU/configs/ (ImageNet-64). Key values:

ImageNet-64 uses the same loss weights and TIRT settings, with a 116.5M teacher (ch=128, nrb=3) and a 22.3M student (ch=64, nrb=2; 5.2× compression), channel multipliers (1,2,3,4) and attention at 16×16 and 8×8, trained for 130 epochs at global batch 2048 with cosine decay from 2e-4 and no warmup.

λ_im = λ_un = 1.0 is not arbitrary: symmetric weights make the residual Gram matrix the identity, so the two branch constraints are equally conditioned and neither dominates the gradient.

Single-head emits no separate conditional and unconditional predictions, so ρ, cos(Δ) and MSE_Δ are undefined for it.

All results: 50-step DDIM, TAG w_max = 4.0, 50K samples.