RiT-XL: Vanilla Diffusion Transformers Suffice in Representation Space
This repository hosts the released RiT-XL checkpoint trained for 800 epochs on ImageNet 256×256 with frozen DINOv2-Small features.
RiT (Representation Image Transformer) is a vanilla Diffusion Transformer that effectively models distributions in high-dimensional representation spaces, as presented in the paper RiT: Vanilla Diffusion Transformers Suffice in Representation Space.
Results on ImageNet 256×256
| Method | Encoder | Params | FID ↓ (CFG=1) | FID ↓ (CFG≈3.7) |
|---|---|---|---|---|
| DiT-XL | SD-VAE | 675M | 9.62 | 2.27 |
| SiT-XL | SD-VAE | 675M | 8.61 | 2.06 |
| REPA-XL | SD-VAE | 675M | 5.78 | 1.29 |
| DDT-XL | SD-VAE | 675M | 6.27 | 1.26 |
| REG-XL | SD-VAE | 675M | 1.80 | 1.36 |
| RAE-XL | DINOv2-S | 676M | 1.87 | 1.41 |
| RAE-XLDH | DINOv2-B | 839M | 1.51 | 1.16 |
| FAE-XL | FAE-DINOv2-G | 675M | 1.48 | 1.29 |
| RiT-XL (ours) | DINOv2-S | 676M | 1.45 | 1.14 |
All FIDs use 25 Heun steps with the time-shift schedule.
Sample Usage
The full training/inference code is available at lezhang7/RiT. To download the weights manually and load them in PyTorch:
import torch
from huggingface_hub import hf_hub_download
# Download the checkpoint
ckpt = hf_hub_download(repo_id="le723z/RiT", filename="checkpoint-last.pth")
# Load the state dictionary
state = torch.load(ckpt, map_location="cpu", weights_only=False)
# state['model'] / state['model_ema1'] / state['model_ema2'] are the
# trainable + two EMA-decay parameter dictionaries.
# state['model_ema1'] is the EMA decay 0.9999 (used for sampling by default).
model_weights = state['model_ema1']
To run the evaluation script (which auto-pulls this checkpoint plus the matching RAE decoder):
git clone https://github.com/lezhang7/RiT.git
cd RiT
pip install -r requirements.txt
bash scripts/eval.sh # CFG=3.7, FID ~1.14 on ImageNet 256x256
Model details
- Architecture: vanilla Diffusion Transformer — 28 layers, hidden 1152, 16 heads, SwiGLU FFN, RMSNorm, QK-norm, 2D VisionRoPE, 32 in-context class tokens, joint [CLS]-patch modeling.
- Encoder (frozen):
facebook/dinov2-with-registers-small(d=384). - Decoder (frozen): ViT-MAE-style decoder from nyu-visionx/RAE-collections, variant
decoders/dinov2/wReg_small/ViTXL_n08/model.pt. - Parameters (denoiser only): 676M.
- Training: 8×H200, batch 1536 effective, AdamW lr=5e-5, 800 epochs, x-prediction loss, dimension-aware time shift (s ≈ 4.9), CLS auxiliary loss weight λ=0.2.
- Sampling defaults: Heun, 25 steps, time-shift schedule, CFG=3.7 in interval [0.1, 0.98], coupled-noise initialization for [CLS].
Citation
@article{zhang2026rit,
title = {RiT: Vanilla Diffusion Transformers Suffice in Representation Space},
author = {Zhang, Le and Mang, Ning and Agrawal, Aishwarya},
journal = {arXiv preprint arXiv:2605.21981},
year = {2026}
}
Acknowledgments
This release reuses the frozen DINOv2 encoder + ViT decoder pairing from RAE and adopts the modernized DiT block design + in-context class tokens from JiT.