Purify Before You Align: Improving Representation Alignment for Diffusion Models

Yingheng Wang; Yaoqiang Li; Yaqin Wu; Junwen Bai; Jiatao Gu; Christopher De Sa; Volodymyr Kuleshov

cs.CVOct 14, 2026

Purify Before You Align: Improving Representation Alignment for Diffusion Models

Authors: Yingheng Wang, Yaoqiang Li, Yaqin Wu, Junwen Bai, Jiatao Gu, Christopher De Sa, Volodymyr Kuleshov

Organizations: Cornell University · Independent researcher. · University of Pennsylvania

Abstract

Representation alignment improves and accelerates diffusion training by encouraging intermediate features to match a pretrained visual encoder. However, standard alignment matches pretrained representations without distinguishing what information may be more useful for generation. We introduce local appearance purification (LAP), which improves generation quality and makes encoders such as MAE more effective alignment targets by removing locally predictable content and aligning to the residual. LAP modifies only the alignment target, and requires no changes to diffusion training or inference. Through controlled interventions within a fixed encoder, we show that effective alignment depends on what information the target emphasizes, not simply how much image information it contains. Our theory gives a variational interpretation of standard alignment and, for fixed local predictors, a sum of token-wise information bounds conditioned on local appearance. To understand why purification benefits some encoders more than others, we derive an exact constant-readout test that measures how well the target can be matched without using the diffusion model's hidden state. LAP improves generation quality across five pretrained representations (MAE, AIMv2, CLIP, DINOv2, DINOv3), with the largest gain for MAE: on ImageNet-256 with SiT-B/2, it reduces FID from $53.1$ to $42.3$ . Its gains persist with stronger alignment recipes and tokenizers. Computed before diffusion training, the constant-readout score predicts the benefit of purification (Spearman $ρ=0.94$ across 13 encoders).

Figures & tables

Figure 1: Purification improves generation by changing the alignment target. (a) LAP subtracts a local prediction and aligns to the residual. (b) Both LAP-L and LAP-N improve MAE-based generation for SiT-B/2 and surpass DINOv2 REPA in the development setting.

Figure 2: One encoder, opposite alignment effects. (a) DINOv2 features split into a locally predictable component (about 2% of centered feature variance and has rank at most 17) and a residual. Each map shows the similarity of every patch to the red-boxed patch ( Appendix D ). (b) Aligning to the locally predictable component hurts generation, whereas the residual (LAP-L) keeps the benefit.

Figure 3: Better clean-image prediction generally accompanies worse generation. Dev-200/50K, FID-DINOv2 on 10K samples. The dashed line is a descriptive linear fit ( Table 32 ).

Figure 4: Our LAP-N framework. Compared with SiT training with REPA, LAP-N modifies only the alignment target by subtracting the reconstruction content predicted by a frozen local predictor $d_{\phi}$ . Training and inference are otherwise unchanged.

Teacher	$C_{0}$	Method	FID $\downarrow$	IS $\uparrow$	Prec. $\uparrow$	Rec. $\uparrow$
No alignment			53.92 $\pm 0.03$	25.34 $\pm 0.16$	0.435 $\pm 0.002$	0.602 $\pm 0.006$
MAE ViT-L	0.93	REPA	53.05 $\pm 0.90$	26.30 $\pm 0.45$	0.438 $\pm 0.004$	0.599 $\pm 0.001$
		$+$ LAP-L	45.15 $\pm 0.48$	31.44 $\pm 0.25$	0.480 $\pm 0.005$	0.617 $\pm 0.003$
		$+$ LAP-N	42.30 $\pm 1.08$	33.85 $\pm 0.44$	0.496 $\pm 0.006$	0.622 $\pm 0.010$
AIMv2 ViT-L	0.66	REPA	47.97 $\pm 0.15$	29.88 $\pm 0.01$	0.464 $\pm 0.000$	0.615 $\pm 0.007$
		$+$ LAP-L	44.31 $\pm 0.13$	32.67 $\pm 0.32$	0.484 $\pm 0.000$	0.622 $\pm 0.006$

Table 1: Purification across five teachers. SiT-B/2 on ImageNet-256, mean

\pm

sample SD over two training seeds. Teachers are ordered by raw-target

C_{0}

on dev-200. Shaded rows use LAP variants. Bold marks the best mean per teacher and metric. Individual FIDs are in Table 34 .

Teacher	Method	FID $\downarrow$	IS $\uparrow$	Prec. $\uparrow$	Rec. $\uparrow$
MAE ViT-L	iREPA	43.06 $\pm 0.62$	33.23 $\pm 0.66$	0.491 $\pm 0.004$	0.617 $\pm 0.006$
	$+$ LAP	41.32 $\pm 0.28$	34.70 $\pm 0.37$	0.502 $\pm 0.002$	0.626 $\pm 0.001$
	VA-REPA	52.82 $\pm 0.90$	26.32 $\pm 0.58$	0.437 $\pm 0.001$	0.606 $\pm 0.006$
	$+$ LAP	41.60 $\pm 0.91$	34.06 $\pm 0.53$	0.503 $\pm 0.004$	0.622 $\pm 0.004$
	REG	52.92 $\pm 0.53$	26.33 $\pm 0.41$	0.434 $\pm 0.002$	0.608 $\pm 0.002$
	$+$ LAP	41.65 $\pm 0.06$	34.31 $\pm 0.16$	0.500 $\pm 0.003$	0.626 $\pm 0.002$

Table 2: LAP on top of REPA-family recipes. Same protocol as Table 1 . Each

+

LAP row uses the variant with the best two-seed mean FID. For REG/sREPA, these are LAP-N/LAP-N for MAE and LAP-L/LAP-N for DINOv2. Bold marks the better one. Full results are in Table 25 .

Tokenizer	Method	FID $\downarrow$	IS $\uparrow$	Prec. $\uparrow$	Rec. $\uparrow$
SD-VAE	No alignment	53.92 $\pm 0.03$	25.34 $\pm 0.16$	0.435 $\pm 0.002$	0.602 $\pm 0.006$
	REPA	53.05 $\pm 0.90$	26.30 $\pm 0.45$	0.438 $\pm 0.004$	0.599 $\pm 0.001$
	$+$ LAP	42.30 $\pm 1.08$	33.85 $\pm 0.44$	0.496 $\pm 0.006$	0.622 $\pm 0.010$
EQ-VAE	No alignment	52.50 $\pm 0.04$	25.99 $\pm 0.05$	0.449 $\pm 0.003$	0.599 $\pm 0.004$
	REPA	52.41 $\pm 1.37$	26.31 $\pm 0.91$	0.452 $\pm 0.004$	0.597 $\pm 0.004$
	$+$ LAP	40.01 $\pm 0.25$	35.24 $\pm 0.36$	0.523 $\pm 0.003$	0.608 $\pm 0.005$

Table 3: LAP with other tokenizers. MAE ViT-L teacher, same protocol as Table 1 . Each

+

LAP row uses LAP-N. Bold marks the best.

Figure 5: Purification improves throughout training. Two-seed FID-10K means and ranges for SiT-B/2, using LAP-N for MAE/CLIP and LAP-L for AIMv2/DINOv3. Full trajectories and evaluation details: Section H.4 .

Dataset	Method	FID $\downarrow$	IS $\uparrow$	Prec. $\uparrow$	Rec. $\uparrow$
RESISC45	No alignment	68.06 $\pm 5.10$	4.75 $\pm 0.17$	0.421 $\pm 0.014$	0.200 $\pm 0.020$
	REPA (DINOv2)	66.34 $\pm 3.70$	4.80 $\pm 0.15$	0.426 $\pm 0.008$	0.184 $\pm 0.016$
	REPA (MAE)	64.65 $\pm 4.38$	4.94 $\pm 0.23$	0.414 $\pm 0.013$	0.205 $\pm 0.024$
	LAP (MAE)	61.47 $\pm 2.33$	4.96 $\pm 0.04$	0.433 $\pm 0.010$	0.228 $\pm 0.009$
AID	No alignment	62.41 $\pm 2.06$	6.47 $\pm 0.02$	0.365 $\pm 0.003$	0.222 $\pm 0.003$
	REPA (DINOv2)	63.68 $\pm 1.35$	6.36 $\pm 0.02$	0.376 $\pm 0.002$	0.200 $\pm 0.003$

Table 4: LAP on aerial-scene datasets. Mean

\pm

sample SD over four training seeds for RESISC45 and two for AID. FID uses 10K samples and 10K reference images. LAP uses LAP-N. Bold marks the best. Details: Appendix J .

Figure 6:

C_{0}

predicts purification gain. Raw-to-LAP-L FID-10K improvement at 100K steps. Dashed line: quadratic fit.

Appendix figures & tables41 assets

Supplementary material from the paper’s appendix.

Appendix

Teacher	Input	Norm.	$D$	Extracted features
MAE B/L/16	256	IN	768/1024	Last block, before final LN
DINOv2 S/B/L/14	224	IN	384/768/1024	Normalized patch tokens
CLIP L/14	224	CLIP	1024	Before final LN/projection
CLIP B/16	256	CLIP	768	Before final LN/projection
I-JEPA H/14	224	IN	1280	Final normalized tokens
MoCo v3 B/16	256	IN	768	Backbone patch tokens

Appendix

Table 5: Teacher preprocessing and feature extraction. Input is the actual backbone resolution after resizing the original

256\times 256

RGB image. All outputs contain 256 patch tokens. IN denotes ImageNet normalization, CLIP denotes OpenAI CLIP normalization, and half denotes mean and standard deviation

0.5

per channel. CLS and register tokens are excluded. Recorded checkpoint identifiers and normalization constants are in the reproducibility manifest.

	DINOv2 presence $R^{2}$			Velocity	Use
Model	Token (dense)	Token (paired)	Image-split	$R^{2}$	(%)
Unaligned	0.196	0.166	0.130	0.451	14.03
LAP-L (MAE)	0.268	0.226	0.172	0.114	15.38
LAP-N (MAE)	0.290	0.240	0.178	0.004	10.61
REPA	0.343	0.286	0.207	0.115	12.21
LAP-L (DINOv2)	0.342	0.289	0.213	$-0.016$	6.67

Appendix

Table 6: Image-disjoint diagnostics. ImageNet-1k/400K, mid-noise, seed 0 unless marked. Presence always probes raw DINOv2. “Token (dense)” and “Token (paired)” denote dense and matched-pool token splits, respectively. Velocity

R^{2}

and use are evaluated on disjoint images. Use is a percentage of full-layer mean-fill ablation.

Figure 7: Purification preserves readable features while reducing prediction coupling. (a) DINOv2 decodability after MAE-based alignment. (b) Absolute probe

R^{2}

(top) and normalized ablation effect (bottom) for DINOv2 targets. Negative

R^{2}

is worse than the evaluation-mean baseline. Presence always probes raw DINOv2. Whiskers give 95% image-bootstrap intervals. Seed 0, mid-noise, 400K steps, disjoint fit and evaluation images ( Section C.3 ).

Model	Noise	Presence $R^{2}$	Cosine	Velocity $R^{2}$	Use (%)
Unaligned	High	0.177	0.005	0.582	8.5
	Mid	0.196	0.008	0.451	13.9
	Low	0.196	0.006	0.333	15.0
Raw MAE	High	0.255	0.947	0.329	3.8
	Mid	0.386	0.960	0.275	17.7
	Low	0.485	0.970	0.269	22.7

Appendix

Table 7: Diagnostics across noise levels. ImageNet-1k/400K, seed 0, dense token-split protocol. Use is the rank-256 aligned-subspace effect as a percentage of the full-layer effect. Probe targets are specified in the text. A dash marks a cosine that is undefined because the projector outputs MAE-dimensional features.

Checkpoint	High	Mid	Low
REPA (400K)	13.8 [13.6, 14.0]	22.4 [21.0, 23.9]	30.1 [28.8, 31.5]
Unaligned (400K)	13.0	15.9	20.7

Appendix

Table 8: Random-subspace references. Dense token-split protocol. Mean-fill use as a percentage of full-layer ablation. REPA brackets give the observed ranges over random subspace draws. These are not training-seed intervals.

Model	Mean fill (%)	Resampling (%)
Unaligned	13.9	15.0
Raw MAE	17.7	6.5
LAP-L (MAE)	15.5	4.9
LAP-N (MAE)	11.0	1.4
REPA	10.9	2.0
LAP-L (DINOv2)	6.5	1.6

Appendix

Table 9: Intervention sensitivity. Dense token-split protocol, ImageNet-1k/400K, seed 0, mid-noise, rank 256. Each estimator uses its corresponding full-layer normalization.

Model	$r{=}32$	$r{=}128$	$r{=}256$
Unaligned	0.8 (0.2)	4.4 (0.8)	13.9 (2.8)
Raw MAE	2.0 (1.0)	8.6 (7.1)	17.7 (16.4)
REPA	0.8 (0.3)	4.8 (1.8)	10.9 (8.7)
LAP-L (DINOv2)	0.6 (0.3)	2.9 (1.6)	6.5 (6.5)
LAP-N (DINOv2)	0.2 (0.3)	2.1 (1.9)	4.5 (8.8)
LAP-L (MAE)	3.0 (0.3)	10.4 (1.7)	15.5 (5.6)

Appendix

Table 10: Rank sensitivity of mean-fill use. Dense token-split protocol, mid-noise, 400K, seed 0. Values are percentages of full-layer ablation. Parentheses show the target-probe subspace control.

		$C_{0}$ (shared direction)			$C_{0}^{\mathrm{pos}}$ (per position)
Data	Encoder	Raw	LAP-L	LAP-N	Raw	LAP-L	LAP-N
Dev-200	MAE	0.933	0.020	0.011	0.945	0.437	0.329
	CLIP	0.607	0.088	0.067	0.620	0.169	0.110
	DINOv2	0.221	0.007	0.011	0.347	0.272	0.207
	AIMv2	0.663	0.041	0.036	0.701	0.221	0.130
	DINOv3	0.525	0.018	0.017	0.587	0.310	0.168

Appendix

Table 11: Constant-readout baselines by dataset, encoder, and target. Empirical estimates from

2{,}048

images per dataset.

C_{0}

is the best expected cosine of one direction shared by all tokens ( Equation 6 );

C_{0}^{\mathrm{pos}}=\frac{1}{N}\sum_{i}\|\mathbb{E}[\hat{y}^{(i)}]\|\geq C_{0}

is the best expected cosine of a readout with one fixed direction per token position, which uses the position but not the image. LAP-N uses the purifier trained on the same dataset (

k=5

for MAE, AIMv2, and DINOv3,

k=3

for DINOv2 and CLIP); CLIP

k=5

on dev-200 gives

0.062

and

0.094

. The predictable-only DINOv2 control, the batchwise affine prediction of DINOv2 from the clean latent, has

C_{0}=0.802

and

0.733

and

C_{0}^{\mathrm{pos}}=0.804

and

0.735

on dev-200 and ImageNet-1k; LAP is not applied to it.

Encoder	Target	$C_{0}$	$C_{0}^{\mathrm{pos}}$	Seed 0	Seed 1	Mean
MAE	Raw	0.933	0.945	1641.4	1643.2	1642.3
	LAP-L	0.020	0.437	1504.3	1512.3	1508.3
	LAP-N	0.011	0.329	1486.2	1482.6	1484.4
	LAP-P	0.019	0.027	1495.9	1502.8	1499.4
	LAP-N $+$ P	0.012	0.025	1471.3	1436.4	1453.8
DINOv2	Raw	0.221	0.347	1551.2	1518.6	1534.9

Appendix

Table 12: Position-dependent purification at development scale.

C_{0}

and

C_{0}^{\mathrm{pos}}

from

2{,}048

dev-200 images ( Section C.8 ) and dev-200/50K FID-DINOv2 for training seeds 0 and 1 and their mean. The DINOv2 purifier uses

k{=}3

. Shaded rows use LAP; bold marks the lowest FID per encoder and column.

Target	Seed 0	Seed 1	Mean
Raw	53.68	52.41	53.05
LAP-L	44.81	45.48	45.15
LAP-N	43.06	41.53	42.30
LAP-P	44.02	44.62	44.32
LAP-N $+$ P	42.33	41.80	42.07

Appendix

Table 13: Position-dependent purification at full scale. MAE ViT-L teacher, ImageNet-1k/400K FID-50K for training seeds 0 and 1 and their mean; full-scale LAP-P runs were made for MAE only. Shaded rows use LAP; bold marks the lowest FID per column.

		High	Mid	Low
Base range	$R^{2}(x_{t,p}^{(i)}\!\to\!X_{p}^{(i)})$	0.037–0.051	0.472–0.501	0.926–0.934
$\Delta R^{2}$	MAE	0.655	0.266	0.013
	DINOv2	0.268	0.052	$-0.001$
	CLIP	0.114	0.026	$-0.001$
$R^{2}(h_{t}^{(i)}\!\to\!X_{p}^{(i)})$	Unaligned	0.264	0.676	0.949
$R^{2}(h_{t}^{(i)}\!\to\!X_{p}^{(i)})$	REPA	0.248	0.662	0.946

Appendix

Table 14: Clean-latent prediction. Incremental utility uses dev-200 data. Hidden-state scores use the dense token-split protocol on 400K ImageNet-1k checkpoints. The baseline row shows the range across the three encoder evaluations. High, mid, and low denote the noise bins defined above.

Figure 8: Reconstruction content is easy to recover but ineffective for alignment . (a) Dev-200/50K controls (single runs). Baseline lines and bands show two-seed means and ranges. (b) Predictable content is more decodable, while the projector matches the residual more closely (95% image-bootstrap intervals). One 400K REPA checkpoint, mid-noise, disjoint fit/evaluation images ( Appendices C and D ).

Encoder	$k$	$\rho$	Raw accuracy	Purified accuracy	$\rho$ gap
MAE ViT-L	1	0.178	0.661	0.693	$\leq 0.003$
	3	0.314	0.661	0.688	$\leq 0.003$
	5	0.352	0.661	0.687	$\leq 0.003$
DINOv2 ViT-B	3	0.125	0.798	0.794	$\leq 0.003$
CLIP ViT-L	3	0.074	0.787	0.791	$\leq 0.014$
	5	0.084	0.787	0.790	$\leq 0.014$

Appendix

Table 15: Purifier validation on dev-200. Class accuracy is held-out pooled top-1 accuracy over 200 classes. Chance accuracy is

0.005

. The last column is the training-to-validation gap in

\rho

. Bold marks the higher classification accuracy between the raw and purified targets in each row.

Dataset	Encoder	$k$	Pooled res./pred.	Token res./pred.	Both
dev-200	MAE-L	1	$0.6926/0.1492$	$0.5464/0.0603$	Pass
		3	$0.6882/0.2053$	$0.5449/0.0901$	Pass
		5	$0.6868/0.2170$	$0.5452/0.0996$	Pass
dev-200	DINOv2-B	3	$0.7942/0.2703$	$0.7986/0.1311$	Pass
		5	$0.7937/0.2795$	$0.7981/0.1533$	Pass
		7	$0.7935/0.2856$	$0.7976/0.1672$	Pass

Appendix

Table 16: Purifier feasibility checks. Each accuracy pair is residual/predicted-component top-1 accuracy on the same 4,096 held-out images. “Pass” requires both inequalities in both pooled and token probes, using the linear references and tolerances above. DINOv2

k=5,7

and CLIP

k=5

are feasibility controls; the base comparisons use fixed

k=3

checkpoints.

Figure 9: How purification changes the MAE target. (a) Cosine similarity to the fixed red-boxed query patch for the raw features, LAP-L and LAP-N, and the iREPA-style spatially normalized target (SN), on a shared scale. (b) Mean off-diagonal patch cosine over 256 held-out images. Both LAP-L and LAP-N substantially reduce MAE’s high common similarity, with a lower mean for LAP-N; per-image spatial normalization gives a near-zero value in this evaluation.

Model	$s(h_{t})$	$s(a_{t})$	$s(\text{target})$	$\cos(a_{t},\text{bias})$
No alignment	0.618	0.648	–	untrained
Raw MAE	0.716	0.948	0.876	0.737
LAP-N (MAE)	0.764	0.137	0.072	$-0.004$
REPA (DINOv2)	0.690	0.465	0.281	0.192
LAP-L (DINOv2)	0.690	0.410	0.251	$-0.007$

Appendix

Table 17: Patch similarity inside the trained models. Image-averaged mean pairwise cosine

s(\cdot)

over 256 held-out images at

t=0.5

(block 8, 400K, seed 0), and the cosine between the projector output and the projector’s bias vector.

Figure 10: The denoiser’s features under raw and purified targets. (a) Cosine to the red-boxed query patch for the block-8 hidden state $h_{t}$ and the projector output $a_{t}$ at $t=0.5$ , two held-out images, five 400K seed-0 models. (b) Mean pairwise cosine of $h_{t}$ (circles), $a_{t}$ (diamonds), and the model’s own target (bars) over 256 images. Raw MAE has the highest projector-output similarity.

Target	Seed 0	Seed 1	Mean
Raw CLIP	1548.9	1553.0	1550.9
LAP-L	1517.1	1513.9	1515.5
LAP-N, $k=3$	1537.6	1498.1	1517.9
LAP-N, $k=5$	1527.7	1516.2	1522.0

Appendix

Table 18: CLIP target transformations. Dev-200/50K FID-DINOv2. Lower is better. Means are computed before rounding the individual seed values. Bold marks the lowest FID in each column.

Target	MLP head	Conv head $+$ spatial norm
Raw MAE	1641.4 / 1643.2	1508.0 / 1496.9
LAP-N (MAE) ( $k{=}5$ )	1486.2 / 1482.6	1510.7 / 1484.4
Raw DINOv2 ( $=$ REPA)	1551.2 / 1518.6	1476.6 / 1504.3
LAP-N (DINOv2) ( $k{=}3$ )	1508.4 / 1536.2	1514.3 / 1523.7
Predictable-only (DINOv2)	1679.9 / 1641.6	1648.6 / 1653.8

Appendix

Table 19: LAP-N and iREPA-inspired operations. FID-DINOv2 for seeds 0 and 1. The right column bundles a three-layer convolutional head and spatial normalization. Bold marks the lowest FID for each encoder, head, and training seed.

Without guidance
Model	Batch	Steps	Unaligned	Raw MAE	REPA	LAP-N (MAE)	LAP-L (DINOv2)
SiT-B/2	64	400K ∗	53.92	53.05	41.75	42.30	41.47
SiT-B/2	64	800K	43.48	42.60	33.44	33.95	33.24
SiT-B/2	256	400K	36.54	34.72	27.67	28.57	27.58
SiT-L/2	64	400K	37.93	35.07	19.28	23.77	18.60
SiT-L/2	256	400K	20.14	16.05	10.01	11.97	9.78

Appendix

Table 20: Generation quality across model size, batch size, training budget, and guidance. FID-50K, two-seed means; lower is better. Within a row all methods share the configuration, and rows are not comparable to each other. The first row ( ∗ ) is the base setting of Table 1 . REPA aligns to the raw DINOv2-B target and the last column to its LAP-L target. Guidance uses CFG scale

1.5

. Bold marks the lowest FID within each row.

(a) 100K steps: FID-10K
Teacher	Target	Block 4	Block 8
MAE	Raw	87.26 / 87.09 (87.18)	87.83 / 86.54 (87.19)
MAE	LAP-N	71.73 / 71.20 ( 71.46 )	72.91 / 72.63 (72.77)
AIMv2	Raw	80.95 / 80.94 (80.95)	82.41 / 82.52 (82.47)
AIMv2	LAP-L	72.51 / 73.43 ( 72.97 )	78.98 / 77.28 (78.13)
CLIP	Raw	80.04 / 80.17 (80.11)	79.33 / 79.93 (79.63)

Appendix

Table 21: Alignment at block 4 versus block 8. (a) Inception FID-10K of the 100K-step checkpoints; (b) FID-50K at 400K steps. Two training seeds with their mean in parentheses. Target is the alignment target; all rows use REPA except the last pair, which uses REG. The DINOv2 rows and the block-4 REG rows use two training seeds separate from those of Tables 1 and 2 , so their block-8 DINOv2 values differ slightly from Table 1 ; the block-8 REG values come from the training seeds of Table 2 . Bold marks the lowest reported mean per teacher and evaluation budget.

Base recipe: REPA, SD-VAE, block 8
Teacher	Target	FID $\downarrow$	IS $\uparrow$	Prec. $\uparrow$	Rec. $\uparrow$
–	No alignment	27.70	53.78	0.603	0.530
MAE	Raw	26.37	57.46	0.606	0.532
MAE	LAP-N	17.90	79.89	0.676	0.529
AIMv2	Raw	22.32	68.58	0.638	0.528
AIMv2	LAP-L	19.68	77.35	0.654	0.545

Appendix

Table 22: Guided sampling. FID-50K, Inception Score, precision, and recall at CFG scale

1.5

for the 400K checkpoints; two-seed means (the predictable-only control has one seed). The first block is the base recipe of Table 1 ; each later block changes one component: the REG recipe ( Table 2 ), the frozen tokenizer ( Table 3 ), or the alignment depth ( Table 21 , with its training seeds). Target is the alignment target. Shaded rows use LAP; bold marks the better value within each raw/LAP pair.

Target	25K	50K	75K	100K
Unaligned	2564.6	1652.4	1454.3	1355.1
Raw self-teacher	2555.0	1622.5	1425.3	1307.5
Residual self-teacher	2537.3	1623.6	1429.5	1313.1

Appendix

Table 23: Encoder-free target comparison. Dev-200 FID-DINOv2, single runs; the unaligned value at 25K comes from a second training seed. Bold marks the lowest FID at each training checkpoint.

Target	FID $\downarrow$	IS $\uparrow$	Prec. $\uparrow$	Rec. $\uparrow$
No alignment	53.92 $\pm 0.03$	25.34 $\pm 0.16$	0.435 $\pm 0.002$	0.602 $\pm 0.006$
REPA (raw DINOv2)	41.75 $\pm 0.56$	35.50 $\pm 0.62$	0.492 $\pm 0.001$	0.630 $\pm 0.005$
REPA $+$ prediction loss ( Section 3.1 )	46.79 $\pm 0.16$	29.99 $\pm 0.27$	0.473 $\pm 0.001$	0.617 $\pm 0.007$
Predictable-only target ( Section 3.2 )	56.71 $\pm 0.42$	24.37 $\pm 0.00$	0.415 $\pm 0.001$	0.591 $\pm 0.003$

Appendix

Table 24: Standard metrics of the interventions of Section 3 . DINOv2 ViT-B teacher, protocol of Table 1 , mean

\pm

sample SD over two training seeds. The first two rows repeat Table 1 for reference. Bold marks the best mean per metric.

MAE ViT-L teacher
Recipe	Target	FID $\downarrow$	IS $\uparrow$	Prec. $\uparrow$	Rec. $\uparrow$
iREPA	Raw	43.06 $\pm 0.62$	33.23 $\pm 0.66$	0.491 $\pm 0.004$	0.617 $\pm 0.006$
	LAP-L	43.10 $\pm 0.40$	33.07 $\pm 0.62$	0.490 $\pm 0.003$	0.627 $\pm 0.005$
	LAP-N	41.32 $\pm 0.28$	34.70 $\pm 0.37$	0.502 $\pm 0.002$	0.626 $\pm 0.001$
VA-REPA	Raw	52.82 $\pm 0.90$	26.32 $\pm 0.58$	0.437 $\pm 0.001$	0.606 $\pm 0.006$
	LAP-L	43.43 $\pm 0.63$	32.65 $\pm 0.63$	0.492 $\pm 0.003$	0.620 $\pm 0.001$

Appendix

Table 25: Linear and nonlinear purification under the REPA-family recipes of Table 2 . Protocol of Table 1 , mean

\pm

sample SD over two training seeds. MAE under REG and sREPA was trained with LAP-N only. Shaded rows use LAP. Bold marks the best mean within each teacher–recipe group and metric, including ties at displayed precision.

Target	FID $\downarrow$	IS $\uparrow$	Prec. $\uparrow$	Rec. $\uparrow$
No alignment	32.30 $\pm 0.04$	40.61 $\pm 0.09$	0.584 $\pm 0.002$	0.604 $\pm 0.005$
Raw	26.10 $\pm 0.32$	52.41 $\pm 0.72$	0.611 $\pm 0.004$	0.620 $\pm 0.001$
LAP-L	25.69 $\pm 0.41$	53.11 $\pm 0.25$	0.615 $\pm 0.003$	0.614 $\pm 0.000$
LAP-N	24.94 $\pm 0.12$	55.13 $\pm 0.41$	0.614 $\pm 0.003$	0.617 $\pm 0.005$

Appendix

Table 26: Linear and nonlinear purification with the DINOv2 teacher on the frozen REPA-E VAE. Protocol of Table 1 , mean

\pm

sample SD over two training seeds. The MAE results on this tokenizer and on EQ-VAE are in Table 3 . Shaded rows use LAP; bold marks the best mean per metric.

			100K steps, FID-10K			400K steps, FID-50K
Encoder	$C_{0}$ raw	$C_{0}$ LAP-L	Raw	LAP-L	Gain	Raw	LAP-L	Gain
MAE ViT-B	0.938	0.022	86.94	75.48	11.47	53.22	43.97	9.25
MAE ViT-L †	0.933	0.020	87.19	76.53	10.66	53.05	45.15	7.90
CLIP ViT-B/16	0.704	0.246	83.46	78.92	4.53	49.05	45.70	3.35
DINOv3 ViT-H+	0.682	0.033	80.47	77.27	3.20	45.59	44.25	1.34
AIMv2 ViT-L †	0.663	0.041	82.47	78.13	4.34	47.97	44.31	3.66

Appendix

Table 27: Encoder sweep at 100K and 400K steps. Dev-200

C_{0}

of the raw and LAP-L targets; Inception FID-10K of the 100K-step checkpoints and FID-50K of the 400K-step checkpoints with the raw and the LAP-L target, two-seed means, with the gain (raw minus LAP-L) at each step count. Rows are ordered by raw

C_{0}

; † marks the five teachers of Table 1 , whose 400K values are those of Tables 1 and 34 . Bold marks the lower FID between raw and LAP-L targets within each row and step count, including ties at displayed precision. Seed-level gains are given in the text.

Figure 11:

C_{0}

predicts the purification gain at 100K and 400K steps. FID gain from LAP-L against raw-target

C_{0}

(dev-200) for 13 encoders, two-seed means with seed ranges: (a) Inception FID-10K at 100K steps, (b) FID-50K at 400K steps. Diamonds are the five teachers of Table 1 ; values in Table 27 .

Teacher	Target	FID-Inception	FID-DINOv2
–	No alignment	124.72	1652.4
MAE	Raw	122.44	1641.4
	LAP-L	107.55	1504.3
	LAP-N	106.98	1486.2
CLIP	Raw	114.64	1548.9
	LAP-L	110.72	1517.1

Appendix

Table 28: Development results in two feature spaces. Dev-200/50K, training seed 0; both metrics are computed on the same generated samples within each row. Shaded rows use LAP; bold marks the lowest FID in each feature space.

Teacher	Target	100K	200K	300K	400K
–	No alignment	88.47 $\pm 0.60$	70.79 $\pm 1.56$	62.01 $\pm 1.02$	56.61 $\pm 0.26$
MAE ViT-L	Raw	87.19 $\pm 0.91$	69.10 $\pm 1.03$	60.87 $\pm 1.34$	55.54 $\pm 1.20$
MAE ViT-L	LAP-N	72.77 $\pm 0.19$	56.55 $\pm 0.51$	49.18 $\pm 0.47$	45.05 $\pm 1.32$
AIMv2 ViT-L	Raw	82.47 $\pm 0.08$	64.33 $\pm 0.17$	56.03 $\pm 0.08$	50.73 $\pm 0.27$
AIMv2 ViT-L	LAP-L	78.13 $\pm 1.20$	60.28 $\pm 0.12$	52.41 $\pm 0.19$	47.19 $\pm 0.34$
CLIP ViT-L	Raw	79.63 $\pm 0.42$	60.94 $\pm 0.41$	52.36 $\pm 0.09$	47.43 $\pm 0.09$

Appendix

Table 29: Complete Inception FID-10K trajectories. ImageNet-1k, SiT-B/2, batch 64, no guidance. Mean

\pm

sample SD over two training seeds at every checkpoint. Each teacher shows the raw target and the LAP variant selected by the final FID-50K in Table 1 , held fixed across checkpoints, adding DINOv2 to the four teachers of Figure 5 . Shaded rows use LAP; bold marks the better value within each raw/LAP pair.

Figure 12: Samples across teachers ( Table 1 ). SiT-B/2, batch 64, SD-VAE, REPA at block 8. Each teacher shows the raw target and the LAP variant of Table 1 . Guidance scale $4.0$ ; within each comparison, a column shares class and sampling noise across rows ( Section H.5 ). Purple labels mark LAP targets.

Figure 13: Samples across REPA-family recipes ( Table 2 ). SiT-B/2, batch 64, SD-VAE, MAE ViT-L teacher with the raw target or LAP-N. Guidance scale $4.0$ ; within each comparison, a column shares class and sampling noise across rows ( Section H.5 ). Purple labels mark LAP targets.

Figure 14: Samples across frozen tokenizers ( Table 3 ). SiT-B/2, batch 64, MAE ViT-L teacher. Guidance scale $4.0$ ; within each comparison, a column shares class and sampling noise across rows ( Section H.5 ). Purple labels mark LAP targets.

Figure 15: Samples at SiT-XL/2, batch 256 ( Table 20 ). MAE ViT-L and DINOv2 ViT-B teachers with the raw target and the LAP variant of Table 20 . Guidance scale $4.0$ ; within each comparison, a column shares class and sampling noise across rows ( Section H.5 ). Purple labels mark LAP targets.

Dataset	Alignment target	$n$	FID-10K
RESISC45 (45 classes)	No alignment	4	$68.1\pm 5.1$
	DINOv2	4	$66.3\pm 3.7$
	ImageNet MAE, raw	4	$64.6\pm 4.4$
	LAP-N (ImageNet MAE)	4	$\mathbf{61.5}\pm 2.3$
AID (30 classes)	No alignment	2	$62.4\pm 2.1$
	DINOv2	2	$63.7\pm 1.3$

Appendix

Table 30: Development-scale alignment on two aerial-scene datasets. FID-10K against 10K training images of the same dataset; mean

\pm

standard deviation over

n

training seeds. Lower is better. Shaded rows are LAP targets. Bold marks the lowest mean FID within each dataset.

Dev-200, 100K steps
$\lambda_{s}$	Velocity $R^{2}$	Align. cosine	Use (%)	Proxy seed 0	Proxy seed 1
0	0.16	0.65	8.3	1157.3	1161.2
0.5	0.50	0.63	7.6	1167.2	1207.1
2	0.52	0.61	40.2	1242.0	1252.3
ImageNet-1k, 400K steps
0	0.11	0.64	11.4	825.1	821.1

Appendix

Table 31: Auxiliary prediction-loss sweep. Diagnostics are two-seed means. Dev-200 uses the token-split protocol, while full-scale diagnostics use the image-split protocol. Proxy columns give individual FID-DINOv2 scores. At

\lambda_{s}=2

, use is

37.6/42.8\%

on dev-200 and

18.9/11.5\%

at full scale (seeds 0/1). Standard FID-50K values are in Table 34 . Bold marks the lowest generation proxy per scale and training seed. Diagnostic columns are not ranked as generation performance.

Figure 16: Readability, reliance, and the auxiliary prediction loss. (a) Rank-256 mean-fill ablations of the REPA model as a percentage of the full-layer effect (ImageNet-1k/400K, seed 0, mid-noise, disjoint fit/evaluation images, 95% image-bootstrap intervals). (b) The auxiliary prediction loss on dev-200/100K: velocity sufficiency saturates by

\lambda_{s}=0.5

while FID-DINOv2 keeps worsening, and reliance on the aligned subspace (“use”) rises only at

\lambda_{s}=2

. These diagnostics use the token-split protocol. (c) The same loss at full scale, with standard FID-50K and disjoint-image diagnostics. Points are two-seed means and bars span the seed values. Values are in Table 31 .

Encoder	Utility $R^{2}$	Seed 0	Seed 1	Mean
CLIP	0.18	1548.9	1553.0	1551.0
I-JEPA	0.29	1529.4	1515.7	1522.6
DINOv2	0.33	1551.2	1518.6	1534.9
MoCo v3	0.44	1584.2	1598.0	1591.1
MAE	0.72	1641.4	1643.2	1642.3

Appendix

Table 32: Prediction utility and generation of five raw encoders. Utility is the clean-latent probe

R^{2}

of the raw target. Generation is dev-200/50K FID-DINOv2 with REPA for training seeds 0 and 1 and their mean; these are the points of Figure 3 . The unaligned model scores 1651.2. Bold marks the lowest FID in each column.

Teacher	Target	Seed 0	Seed 1	Mean
–	No alignment	1652.4	1649.9	1651.2
MAE	Raw	1641.4	1643.2	1642.3
	LAP-L	1504.3	1512.3	1508.3
	LAP-N	1486.2	1482.6	1484.4
CLIP	Raw	1548.9	1553.0	1551.0
	LAP-L	1517.1	1513.9	1515.5

Appendix

Table 33: Development-scale target ladder. Dev-200/50K FID-DINOv2 with REPA for training seeds 0 and 1 and their mean; the MAE rows are the ladder of Figure 1 (b). The centered variance explained by the DINOv2 linear split is about

2.7\%

. LAP-P is in Table 12 , the iREPA operations in Table 19 , and FID-50K for these targets in Tables 1 and 34 . Shaded rows use LAP; bold marks the lowest FID within each teacher and column.

Teacher	Target	Seed 0	Seed 1
–	No alignment	53.95	53.90
MAE ViT-L	Raw	53.68	52.41
	LAP-L	44.81	45.48
	LAP-N	43.06	41.53
AIMv2 ViT-L	Raw	48.08	47.87
	LAP-L	44.40	44.22

Appendix

Table 34: Standard FID-50K by training seed: teachers and interventions. Individual values behind Tables 1 and 24 ; ADM protocol, SDE-250, no guidance. Shaded rows use LAP; bold marks the lowest FID within each teacher and training seed.

REPA-family recipes, MAE ViT-L
Setting	Target	Seed 0	Seed 1
iREPA	Raw	43.50	42.62
	LAP-L	43.38	42.81
	LAP-N	41.12	41.52
VA-REPA	Raw	53.46	52.19
	LAP-L	42.98	43.88

Appendix

Table 35: Standard FID-50K by training seed: recipes and tokenizers. Individual values behind Tables 2 , 3 , 25 and 26 ; same protocol as Table 34 . Shaded rows use LAP; bold marks the lowest FID within each teacher–recipe or tokenizer group and training seed.

Diffusion Model Alignment

Diffusion model alignment steers generated content and internal representations toward prompts, human preferences, safety goals, and task constraints.

See all 185 papers

Jun 7, 2026

MaskAlign: Token-Subset Representation Alignment for Efficient Diffusion Training

Representation alignment with pretrained vision models has recently shown strong potential for accelerating diffusion transformer training. By aligning intermediate diffusion features with clean-image representations from self-supervised vision encoders, existing methods improve convergence and generation quality. However, such alignment also introduces a non-trivial constraint: diffusion models operate on noisy inputs whose usable information varies across timesteps, while the reference features are extracted from clean images. In this paper, we revisit this mismatch from a token-level perspective. We find that, under full-token representation alignment, tokens with large alignment-gradient norms exhibit a stable spatial preference, suggesting that the alignment objective does not affect all tokens uniformly and may encourage the model to rely on the complete set of clean-image tokens. To address this issue, we propose MaskAlign, a token-subset representation alignment method that applies alignment to randomly sampled token subsets during training. By exposing the model to different token subsets across iterations, MaskAlign reduces the dependence of representation alignment on the complete token set and encourages alignment behavior that is more stable under token-subset perturbations. To mitigate the information loss caused by directly dropping tokens, we further introduce a lightweight pre-mask token mixing block that shares information across tokens before masking.

Lianyu Pang, Tianlin Pan, Cheng Da +5

The Hong Kong University of Science and Technology · Kuaishou Technology · University of Chinese Academy of Sciences

Diffusion Model Alignment Diffusion Transformer Diffusion Model Acceleration Representation Alignment

Ask PDF

Sep 28, 2026

What Visual Generators Need from Teachers: Rethinking Representation Alignment

Representation alignment speeds up diffusion transformer training by pulling an intermediate block of the model (student) toward features of a frozen pretrained encoder (teacher). Which teacher layer to align, and for how long, is still set by convention, and each alternative costs a training run. We find that alignment helps where the student cannot linearly recover the teacher's features, not where it already resembles them. Since a deep teacher layer is largely predictable from the one below, we isolate what each layer adds, its increment, and measure how much of it an unaligned student recovers. The student fills the teacher's hierarchy from the bottom up and stalls near the top, which we call hierarchy filling: even after 400K steps it recovers almost none of the deepest. The recoverability gap is the unrecovered share of an increment, read from one unaligned checkpoint. In short runs that each align one teacher layer at one block, the gap nearly reproduces their ranking by FID improvement, and CKA, a measure of feature similarity, largely reverses it. Representation Alignment and Recoverability Estimation (RARE) picks the teacher layer with the largest gap before training. During training, it tracks each token's remaining distance to that layer, the online counterpart of the gap, weights tokens by it, and phases out the loss once the average distance stops falling. With SiT-B/2 on ImageNet

256\times256

, RARE reaches an FID of 18.02 without guidance and 4.46 with it, ahead of seven alignment baselines including REPA, iREPA and HASTE. It also trains in 14% fewer GPU-hours than iREPA. Its FID stays below iREPA's across model scales, teachers, datasets and backbones.

Yongcong Wang, Hingchin Chen, Mingyu Fan +6

Central South University · The Hong Kong University of Science and Technology · Tsinghua University +7

Diffusion Model Alignment Diffusion Transformer Representation Alignment

Ask PDF

Oct 14, 2026

Representation Conditioning for Generation

Representation-alignment methods effectively accelerate the diffusion transformer training. REPA, for example, uses a pretrained vision encoder to supervise an intermediate layer through an MLP projection head which is discarded after training. It aims to import semantics from the vision encoder to the diffusion backbone. While this objective is achieved at early layers by back-propagating the alignment loss, there is no guarantee that later layers, supervised by the velocity prediction loss alone, can well access such semantics. We propose a simple solution, representation conditioning (RECON), to improve semantics in later layers. During training, we feed the pretrained vision encoder tokens into later layers; during inference, because we no longer have the clean image, we reuse the MLP head (as a proxy to the vision encoder) and project its output tokens to the later layers. Our method adds no parameters and leads to much faster convergence: RECON is

63\times

faster than SiT and

25\times

faster than REPA. Code is available at https://github.com/saeidrazavi/RECON

Saeed Razavi, Xingjian Leng, Liang Zheng

Australian National University

Diffusion-Based Image Generation Diffusion Transformer Diffusion Model Acceleration Conditional Diffusion Models

Ask PDF

Abstract

Diffusion Model Alignment

Explore similar work

MaskAlign: Token-Subset Representation Alignment for Efficient Diffusion Training

What Visual Generators Need from Teachers: Rethinking Representation Alignment

Representation Conditioning for Generation