shumpei777/rae-cifar

RAE on CIFAR-10: pretrained weights

Weights for rae-cifar, a small-scale reproduction of Diffusion Transformers with Representation Autoencoders (RAE) on CIFAR-10. The encoders were trained from scratch on CIFAR-10 with dino-cifar.

File	Content	Size
encoder_dinov2_vit_tiny4.pt	DINOv2 ViT-Tiny/4 encoder (EMA teacher, no registers)	22 MB
normalization_stats.pt	per-channel, per-position latent mean / var (train, with flips)	0.1 MB
decoder_vitxl_gan.pt	ViT-XL decoder for the Tiny encoder, GAN fine-tuned (EMA)	1.7 GB
ditdh_s_ema.pt	DiTDH-S on the Tiny latent, 2000 epochs (EMA)	0.8 GB
encoder_dinov2_vit_s4.pt	DINOv2 ViT-S/4 encoder (EMA teacher, no registers)	85 MB
decoder_vitxl_gan_s4.pt	ViT-XL decoder for the ViT-S/4 encoder, GAN fine-tuned (EMA)	1.7 GB
disc_dino_vit_tiny4.pt	DINO (v1) ViT-Tiny/4, frozen backbone of the GAN discriminator (only needed for training)	22 MB

Results (details in the docs):

• Tiny encoder + XL decoder with GAN: test rFID 5.71, train rFID 0.89 (50k).
• ViT-S/4 encoder + XL decoder with GAN: test rFID 3.71, train rFID 0.72 (50k).
• DiT^DH-S: FID 13.59 (50k samples, no CFG, against CIFAR-10 train).

Usage

git clone https://github.com/shumpei-takezaki/rae-cifar.git && cd rae-cifar
uv sync
uv run hf download shumpei777/rae-cifar --local-dir pretrained_models
PYTHONPATH=src uv run python src/sample.py --config configs/stage2/sampling/CIFAR10/DiTDH-S_pretrained.yaml   # -> sample.png

The encoders can also be loaded directly with dino-cifar: from cifar_dino.export import load_backbone.

License

MIT. The code builds on the official RAE implementation (MIT).