RAD / 2026

Representation-Augmented Diffusions

A new class of diffusion models that are faster to train, faster to sample, and parameter-efficient, with a smooth latent space.

70×fewer training iterations than standard diffusion

10×fewer sampling steps

4.8×fewer parameters at equal FID

Farrin Marouf Sofian^*, Kushagra Pandey^*, Stephan Mandt
University of California, Irvine · ^*equal contribution

how it works ↓

01Method

A representationat every step

Standard diffusion models only learn to denoise; whatever semantics they pick up stay buried inside the network. RAD gives the model a compact latent representation u_t at every step: it predicts u_t from the noisy state x_t, then denoises with it. These latents are learned by aligning them with an external representation model (external alignment) or with the model's own features (self-supervised, internal alignment).

x_tnoisy latent at time t u_trepresentation p_ϕrepresentation policy p_θdenoising policy qnoising process f_prefrozen encoder (DINOv2 in our runs) h_ϕprojector on block ℓ α_t=1−t, σ_t=tinterpolant βpolicy scale

1 / 7

Predict

An MLP head on denoiser block ℓ reads x_t and samples a representation.

u_t ~ p_ϕ(u_t | x_t)

Feed back

The representation is compressed into one token and added back into the denoiser at the same block.

x_t^ℓ ← x_t^ℓ + MLP(u_t)

Denoise

The remaining SiT blocks predict the velocity, now conditioned on the representation.

v_θ(x_t, u_t, t)

The model, step by step

1

Generative process Eq. 1–2

Each step samples a representation first, then a denoising step consistent with it.

Averaging over u_t makes every transition a continuous mixture of Gaussians. It is more expressive than one Gaussian, which is why RAD tolerates large steps.
2

Noising process Eq. 3, 5

While training, u_t is a noised copy of the clean image's features from a frozen encoder, on the same schedule as x_t: coarse at high noise, sharp at low noise. It sees what the generator cannot, so it acts as a teacher.
3
Variational bound Eq. 4

A conditional VAE at every step: rebuild the cleaner state from u_t, and keep the model's own u_t close to the teacher's. The KL is an information bottleneck.
Proof sketch
- Use the noising process as the variational distribution and apply Jensen's inequality.
- By Markov + Bayes, ∏ q(x_t | x_t−1) telescopes, so each step's terms group together.
- What remains (posterior entropies and the prior KL) depends only on the noising process: a constant in θ, ϕ.
4

Plug in Gaussians App. A.3.2

Same covariance on both sides, so the KL is just a distance between means: REPA's alignment loss, derived instead of assumed.

With a velocity network the denoising term becomes a weighted velocity regression. Set w_t = 1 as usual and you get the loss below.

THE LOSS a weighted variational lower bound on log p(x₀)

DenoisingThe usual velocity loss, except the denoiser also sees the representation u_t.

AlignmentPulls the predicted representation toward the target features (here a frozen encoder). It falls out of the bound as a KL term, not a heuristic add-on.

In practice the weight SNR(t)/2β², which blows up as t → 0, is replaced by its bounded form λ_t = α_t²/(α_t² + β²σ_t²) (× 0.5), with t ~ U(0,1) and β = 1.

KEY INSIGHT

In representation alignment, representations are learned but then discarded: they are never used for generation. Remove u_t from RAD's denoiser and these methods, such as REPA, become special cases.

Architecture

Training in practice

Inject early

Block ℓ = 2 for B/2 and ℓ = 8 for L/2, XL/2, inside the first third of the network. The most influential architectural choice.

B/2: block 2 → FID 8.2 · block 8 → 15.1

Use a flow policy

Replace the policy's Gaussian noise with a conditional normalizing flow (4 affine couplings, zero-initialised, so training starts Gaussian).

FID 10.3 → 8.7

Normalize · compress · modulate

Remove 60% of the token mean (γ = 0.6) and rescale to unit size; an MLP over tokens compresses N tokens to one; adaLN-style modulation, zero-initialised.

FID 8.7 → 8.2

Bounded alignment weight

Use the saturated weight λ_t × 0.5 instead of SNR(t)/2β², plus adaptive per-sample weighting with p = 1.

FID 8.2 → 7.7 (with early injection)

Tune β only at sampling

Train with β = 1; sample with β_inf ∈ [0.5, 0.9] to concentrate the policy, trading diversity for fidelity without retraining.

best β_inf: 0.7 (B/2), 0.8 (L/2, XL/2)

Small overhead

The control pathway is the same size at every scale; most of its cost is the projector, which REPA also trains but throws away.

+12% FLOPs at B/2 · +3.5% L/2 · +2.5% XL/2

Data: ImageNet, 1.28M images · SD-VAE latents 32×32×4
Teacher: DINOv2-B/14 · 256 patch tokens × 768
Backbone: SiT-B/2, L/2, XL/2 · 256 tokens
Optimizer: AdamW · lr 1e-4 · batch 256 · fp16 · EMA 0.9999
Schedule: 400K iterations · α_t = 1−t, σ_t = t · label dropout 0.1
Compute: 8 × RTX PRO 6000 · 11 / 16 / 22 h for B / L / XL

02Efficiency

Faster, smaller,fewer steps

Training convergence, XL/2

RAD (ours) REPA SiT

FID-50k vs. training iterations (log–log), no CFG.

70×

fewer iterations than SiT-XL/2: RAD-XL/2 at 100K matches SiT-XL/2 at 7M.

27×

fewer iterations than REPA-XL/2: RAD-XL/2 at 150K matches REPA-XL/2 at 4M.

Same backbone, encoder, data and schedule as REPA. The only difference is feeding the representation back.

Fewer sampling steps

RAD-XL/2 @400KREPA-XL/2 @4M

FID-50k vs. Euler ODE steps. With UniPC, RAD-XL/2 reaches 6.5 at 5 steps and 4.5 at 10.

Smaller models, same quality

RADREPASiT

FID-50k at 400K iterations, without CFG, across SiT backbones.

More linearly separable features

RAD-B/2REPA-B/2

ImageNet top-1 linear probe per block, t = 0.1. Accuracy decays more slowly once u_t is fed back.

B/2 trains 4× faster

RAD-B/2REPA-B/2REPA-XL/2

RAD-B/2 tracks the 4.8× larger REPA-XL/2 throughout training.

10× less data

RAD-B/2REPA-B/2

400K iterations on a fraction of ImageNet, best checkpoint, 50-step ODE. REPA stops improving after 10%.

Works with fast samplers

EulerUniPC-2UniPC-3

RAD-XL/2 at 400K, no guidance: FID 6.5 in 5 steps with UniPC.

FID-50k with the ADM suite · Euler–Maruyama SDE, 250 steps, unless noted · no classifier-free guidance · best β_inf per checkpoint.

03Results

Measuredgains

70×

fewer training iterations

RAD-XL/2 at 100K matches SiT-XL/2 at 7M

27×

fewer training iterations

RAD-XL/2 at 150K matches REPA-XL/2 at 4M

10×

fewer sampling steps

10 Euler steps FID 6.8 vs. REPA at 100 steps FID 7.4

4.8×

fewer parameters

RAD-B/2 FID 7.7 vs. REPA-XL/2 FID 7.9

10×

less training data

RAD-B/2 on 10% of ImageNet FID 14.9 vs. REPA-B/2 on all 26.6

+4.4

points linear probe

RAD-B/2 vs. REPA-B/2, ImageNet top-1, t = 0.1

2×

faster REG convergence

REG + RAD on SiT-B/2 at 400K FID 15.2 → 8.9

3.94

text-to-image FID, XL/2

vs. REPA 4.31 · GenEval 0.266 vs. 0.258

What matters ablations · SiT-B/2 · 400K iterations

AREPA22.9

B+ feed u_t to the denoiser10.3−12.5

C+ flow policy8.7−1.6

D+ control norm + modulation8.2−0.6

E+ early injection, p = 17.7−0.5

Feeding the representation back does most of the work: FID 22.9 → 10.3 on its own.

Encoder robust

DINOv2-B8.2

DINOv2-L8.4

SigLIP-L8.9

Injection depth earlier is better

block 28.2

block 48.5

block 610.5

block 815.1

Plugs into other methods and tasks

REG + RAD · FID-50k ↓
Backbone	Iter.	REG	+ RAD
SiT-B/2	100K	36.1	22.4
SiT-B/2	400K	15.2	8.9
SiT-L/2	100K	11.4	7.8
SiT-L/2	400K	4.6	3.9

Text-to-image · no guidance
Backbone	Method	FID ↓	GenEval ↑
B/2	REPA	10.77	0.138
B/2	RAD	8.42	0.166
XL/2	REPA	4.31	0.258
XL/2	RAD	3.94	0.266

Against the field ImageNet 256×256

		B/2		L/2		XL/2
Method	Epochs	#Params	FID↓	#Params	FID↓	#Params	FID↓
Diffusion transformers
DiT	80	130M	43.5	458M	23.3	675M	19.5
DiT	1400	–	–	–	–	675M	9.6
SiT	80	130M	33.0	458M	18.8	675M	17.2
SiT	1400	–	–	–	–	675M	8.3
Representation alignment
DiT + REPA	80	–	–	458M	15.6	675M	12.3
REPA	80	130M	24.4	458M	10.0	675M	7.9
iREPA	80	130M	21.4	458M	9.4	675M	7.5
sREPA	80	130M	20.5	458M	8.7	675M	7.2
HASTE	80	130M	19.6	458M	8.9	675M	7.3
REED	80	130M	19.5	458M	8.9	675M	7.3
TREAD	80	130M	16.0	458M	5.9	675M	4.9
REG	80	132M	15.2	460M	4.6	677M	3.4
REPA-E^†	80	–	–	–	–	675M	4.1
Representation augmentation
RAD (ours)	80	142M	7.7	472M	3.9	690M	3.7

FID-50k without CFG · Euler SDE, 250 steps · 80 epochs (400K iterations) · † REPA-E also fine-tunes the VAE end-to-end.

04Samples

What itgenerates

ImageNet class-conditional · 256×256

Curated RAD samples · 50-step ODE · CFG 3.0 · hover for the class

Text-to-image · RAD-XL/2 256×256

BLIP3o captions · frozen Qwen3-0.6B text encoder · FID-50k and GenEval measured without guidance; samples shown with CFG 3–4.

05Interpolation

A smoothlatent space

Take two real ImageNet images, read their representations with RAD's own policy head, and blend them with a weight m. RAD morphs one subject into the other while pose, viewpoint and framing stay put.

Encode

From each real image read ū = h_ϕ(x_{t_c}, t_c) at a small noise level, and invert it to noise x₁ with the ODE.

Blend

ū_m = m ū^A + (1 − m) ū^B; class embeddings blend the same way; the noises are slerped.

Sample

Generate from the blended noise with the blended representation: 50 Euler ODE steps, no guidance, no external encoder.

m = 0.50

Representation, inverted noise and class embedding are interpolated at the same m · 50 Euler ODE steps.

Compared with REPA baboon → macaque

REPA-XL/2, trained 10× longer (4M iterations), jumps between poses and scenes mid-way. RAD-XL/2 at 400K changes the subject gradually.

06BibTeX

Citethis work

@article{maroufsofian2026rad,
  title   = {Representation-Augmented Diffusions},
  author  = {Marouf Sofian, Farrin and Pandey, Kushagra and Mandt, Stephan},
  journal = {arXiv preprint},
  year    = {2026}
}