[NeurIPS'25] U-REPA: Aligning Diffusion U-Nets to ViTs

U-Nets: Great Diffusers with the help of ViT Encoders!

• 🤔 We look at Representation Alignment (REPA) from an architecture perspective.
• 🔧 We tailor REPA for U-Nets based on its difference with ViT or DiT architectures.
• 🚀 The proposed U-REPA tap the potential of U-Nets to the full: better FIDs with half REPA training!

Under the standard REPA Setting, U-Net are much performant (FID=1.41) with only half of the training budget.

Codes

Repo Structure

Our Codes are mainly based on the codebase of REPA. We sincerely thank their authors for opensourcing. Major settings, including most choices of hyperparameters and supporting packages, stick to their convention. Some key modifications are as follows:

• models_urepa.py: the code for SiT$\downarrow$.

• loss.py: the code for manifold loss.

Environment

Please run the following command to automatically install supporting packages:

pip install -r requirements.txt

For DINO-V2, it will be automatically downloaded via torch.hub.

Data Processing

We follow the same data processing procedures as REPA. Please refer to their guidance for the detailed procedures.

Training

For training a SiT$\downarrow$-XL/2, with SOTA performance, Please run the following command:

accelerate launch --machine_rank 0 --num_processes 8 --num_machines 1 train.py --data-dir=../../data/imagenet256 --resolution=256 --model-domain=urepa --model=XL --num-classes=1000 --max-train-steps=4000000 --epochs=1400 --batch-size=256 --seed=0 --path-type=linear --prediction=v --weighting=uniform --enc-type=dinov2-vit-b --proj-coeff=0.5 --encoder-depth=18 --projector-embed-dims=768 --manifold-tp=cos-l2 --manifold-weight=3

Sampling

To replicate the best FID performance with CFG, please use the SDE sampling with guidance interval.

Please run the following command:

torchrun --nnodes=1 --nproc_per_node=8 generate.py --ckpt={path-to-ckpt} --resolution=256 --model-domain=urepa --model=XL --num-classes=1000 --seed=0 --path-type=linear --prediction=v --weighting=uniform --enc-type=dinov2-vit-b --encoder-depth=18 --projector-embed-dims=768 --per-proc-batch-size=64 --num-steps=250 --cfg-scale=1.65 --mode=sde --guidance-high=0.7

Evaluation

Please refer to here for metric evaluation. Please download reference batches for ImageNet-256x256 here.

Weights

We also opensource our weights. Please feel free to download!

	Training iters	FID	Link
SiT$\downarrow$-XL/2	2M	1.41	Link

Relevant Work

U-DiT [NeurIPS'24]: Revealing the power of U-Nets

DiC [CVPR'25]: Conv3x3 U-Nets are speedy & powerful

Acknowledgement

We sincerely thank the authors of DiT SiT REPA Guided Diffusion for their code.

For certain reasons, this is a self-replicated code repo; there might be some inconsistencies with the original one with which I run all experiments (though I am trying my best to reduce them all). Please spare me if there are any problems with the code.