yingheng/LAP-checkpoints

Purify Before You Align: checkpoints

Weights for Purify Before You Align: Improving Representation Alignment for Diffusion Models (paper, code).

Class-conditional SiT-B/2 models trained on ImageNet 256×256 for 400K steps (batch 64). For each model, this repository contains the seed with the lower FID-50K; the paper reports two-seed means.

Model	File	Seed	FID-50K
No alignment	diffusion/in1k_vanilla_s1.pt	1	53.90
MAE, REPA	diffusion/in1k_mae_raw_s1.pt	1	52.41
MAE, LAP-L	diffusion/in1k_mae_res.pt	0	44.81
MAE, LAP-N	diffusion/in1k_mae_lapk5_s1.pt	1	41.53
DINOv2, REPA	diffusion/in1k_repa_s1.pt	1	41.35

FID-50K: EMA weights, Euler–Maruyama SDE with 250 steps, no guidance, ADM evaluator and ImageNet-256 reference.

Files

• diffusion/*.pt: fp32 EMA weights with the alignment projectors and the arguments needed to rebuild the model. No optimizer state, so training cannot be resumed exactly. Loads with torch.load(..., weights_only=True).
• assets/purifier*.pt: the 12 frozen LAP-N purifiers used in the paper. assets/*latents-stats.pt: channel statistics for the EQ-VAE and REPA-E VAE experiments.
• manifest.json: size and SHA256 of every file, plus seed and FID of each model.

Teacher encoders, VAE weights and images are not included. The weights are released under CC BY-NC 4.0; the code is MIT-licensed, with some files under their upstream terms.

Usage

With the code repository:

python scripts/download_checkpoints.py --output checkpoints
python sample.py --ckpt checkpoints/diffusion/in1k_mae_lapk5_s1.pt --out samples/mae_lapn

Sampling downloads the SD-VAE decoder (stabilityai/sd-vae-ft-mse); it does not need the teacher or the purifier. For LAP-N training, copy assets/*.pt into the code repository's assets/.

Limitations

Research checkpoints for studying representation alignment, trained only on ImageNet; they are not intended as general-purpose image generators. Fresh FID estimates vary slightly with sampling noise, hardware and the number of processes.

Citation

@article{lap2026,
  title   = {Purify Before You Align: Improving Representation Alignment for Diffusion Models},
  author  = {Wang, Yingheng and Li, Yaoqiang and Wu, Yaqin and Bai, Junwen and Gu, Jiatao and De Sa, Christopher and Kuleshov, Volodymyr},
  journal = {arXiv preprint arXiv:ARXIV_ID},
  year    = {2026}
}