REPA Works Until It Doesn't: Early-Stopped, Holistic Alignment Supercharges Diffusion Training

TL;DR

HASTE diagnoses why REPA-style external representation alignment stops helping (and starts hurting) late in DiT training, and proposes a two-phase fix. Phase I uses a holistic alignment loss that distills both DINOv2 feature projections and attention maps into mid-depth DiT blocks; Phase II turns the alignment off at a fixed iteration (or when a gradient-cosine trigger flips negative). On ImageNet-256 with SiT-XL/2 it matches the vanilla SiT-XL/2 1400-epoch baseline FID in 50 epochs (28× speedup) and reaches REPA’s best 800-epoch FID in 500 epochs. The framing — “external supervision is a transient scaffold, not an endpoint” — directly relates to the SRA/Self-Flow internal-self-distillation thread.

Key claims

REPA’s alignment loss helps in the first ~100K-400K iterations but flips to actively harming the model later: stopping REPA at 400K iterations improves FID over the always-on baseline, while stopping early at 100K hurts [§1, Fig. 1].
The cause is a capacity mismatch: the cosine similarity between alignment and denoising gradients passes through three regimes — ignition (positive, 0-200K iters), plateau (near orthogonal, 200K-400K), conflict (negative, >400K) — and the FID divergence in Fig. 1 coincides with the negative-gradient crossover [§2.2, Fig. 3].
The mismatch is sharper at high-noise (late-diffusion) timesteps: gradient cosines for textural/high-frequency steps are near-orthogonal from the start of training, while mid-noise steps stay aligned much longer [§2.2, Fig. 4].
Feeding the teacher low-frequency only versions of each image recovers nearly identical early-FID gains as full-resolution REPA, which the paper takes as evidence that REPA’s benefit is entirely “coarse semantic scaffolding” and high-frequency teacher cues are irrelevant [§2.2, Fig. 5].
DINOv2 attention maps (relational priors / routing patterns) are complementary to its feature embeddings; distilling only attention is competitive with distilling only features, and combining the two outperforms either alone [§2.3, Table 3].
Attention should be distilled only into mid-depth DiT blocks (SiT-XL/2 blocks 4-7): shallow blocks are dominated by noise normalization and supervising them with clean-image attention destabilizes training; final blocks should remain dedicated to denoising [§2.3, Table 6].
On ImageNet-256 without CFG, SiT-XL/2 + HASTE reaches FID 5.31 at 100 epochs vs SiT+REPA’s 5.90 at 800 epochs and vanilla SiT-XL’s 8.61 at 1400 epochs [Table 1].
On ImageNet-256 without CFG, DiT-XL + HASTE reaches FID 8.39 at 50 epochs, surpassing vanilla DiT-XL trained for 1400 epochs (FID 9.62) [Table 1].
With CFG, SiT-XL/2 + HASTE matches REPA’s best (FID 1.42) at 500 epochs vs REPA’s 800 [Table 1].
The recipe transfers to text-to-image: on MS-COCO 2014 with MM-DiT, HASTE outperforms both vanilla and MM-DiT+REPA at 150K iterations under both ODE and SDE samplers (e.g. SDE w/o CFG: 6.81 vs REPA’s 7.33 vs vanilla 11.76) [Table 2].
The termination weight is robust: the attention-alignment loss weight λ_attn = 0.5 is best on SiT-L/2 at 400K iters, with larger weights (1.0-3.0) monotonically degrading [Table 5].
A fixed-iteration termination point works “nearly as well” as the gradient-angle trigger, with the gradient rule being more robust across datasets [§2.2, Eq. 2].

Method

HASTE = Holistic Alignment with Stage-wise Termination for Efficient training. During Phase I (iterations 0 ≤ t ≤ T_stop), the model trains with L = L_denoise + λ_feat · L_REPA + λ_attn · L_ATTA. L_REPA is the original REPA cosine-similarity alignment of a single DiT mid-layer feature map to a DINOv2-B patch embedding, projected via a small MLP. L_ATTA is a token-wise cross-entropy between DINOv2 attention matrices and DiT self-attention matrices, applied across a range of mid-depth student-teacher block pairs (e.g. SiT-XL/2 blocks 4-7, DINO blocks selected to match). Shallow DiT blocks (which process Gaussian-noisy latents) and the final denoising blocks are deliberately not aligned. At iteration t = T_stop, both auxiliary losses are switched off and training continues with the vanilla denoising objective. The switch point is either a fixed iteration (chosen near where the gradient-cosine analysis shows the alignment going orthogonal — ~250K for SiT-XL/2) or driven by the gradient-angle trigger directly. Authors use λ_feat following REPA defaults and λ_attn = 0.5; DINOv2-B is the teacher; SiT-XL/2, DiT-XL, and MM-DiT are the students. ImageNet experiments use SD-VAE latents with batch size 256 on A100/H100.

Results

SiT-XL/2 + HASTE, ImageNet-256, no CFG: FID 5.31 / sFID 4.72 / IS 148.5 at 100 epochs. Reaches vanilla SiT-XL/2’s 1400-epoch FID (8.61) in 50 epochs — claimed 28× acceleration over vanilla. Reaches REPA’s 800-epoch FID (5.90) by ~100 epochs [Table 1].
DiT-XL + HASTE, ImageNet-256, no CFG: FID 8.39 at 50 epochs, beating vanilla DiT-XL’s 9.62 at 1400 epochs [Table 1].
SiT-XL/2 + HASTE, ImageNet-256, with CFG: matches REPA’s 800-epoch FID 1.42 at 500 epochs; FID 1.44 at 400 epochs; FID 1.74 at 100 epochs [Table 1].
MM-DiT + HASTE on MS-COCO 2014, 150K iters: ODE no-CFG FID 9.63 vs REPA 10.40 vs vanilla 15.42; SDE no-CFG FID 6.81 vs REPA 7.33 vs vanilla 11.76 [Table 2].
Termination ablation on SiT-XL/2: holistic (REPA+ATTA) without termination achieves FID 8.1 at 100 epochs; adding termination at 50 epochs drops FID to 5.3 — termination provides a ~2.8 FID gain on top of holistic alignment alone [Table 3].
Termination across model sizes: SiT-B/2 termination at 100K improves FID 21.3 → 19.6 at 400K; SiT-L/2 termination at 250K improves 8.9 → 7.9 [Table 4].
Layer ablation on SiT-L/2: aligning DINO layers 8-11 → SiT layers 8-11 gives FID 7.9; including shallow layers 0-3 degrades to FID 8.3 [Table 6].

Why it’s interesting

This is the missing diagnostic piece in the Diffusion training efficiency external-vs-internal-supervision fork. The concept page already records that No Other Representation Component Is Needed: Diffusion Transformers Can Provide Representation Guidance by Themselves (SRA) (SRA) and Self-Flow: Self-Supervised Flow Matching for Scalable Multi-Modal Synthesis argue external DINO supervision is dispensable and doesn’t scale. HASTE arrives at a compatible conclusion from the opposite direction — keep REPA, but treat it as a transient scaffold — and provides the gradient-cosine analysis that quantifies when external alignment crosses from helpful to harmful. Three concrete implications:

Reframes Self-Flow’s scaling failure of REPA. Self-Flow reports REPA flattens past 625M params; HASTE reports REPA’s per-step gradient flips to obtuse past 400K iters on SiT-XL/2. Both observations point to the same capacity-mismatch story, just measured along different axes (model size vs training time). Termination at the gradient-cosine crossover is plausibly the missing recipe that would let REPA continue to scale — or at least delay where it stops.
The attention-distillation component is novel for diffusion. ATTA-style attention transfer is well-studied for ViTs but HASTE is the first to show DINOv2 attention maps usefully supervise diffusion transformers, and that attention and feature alignment are independently effective. This matters when reading Improved Baselines with Representation Autoencoders (RAEv2) and PixelGen: Improving Pixel Diffusion with Perceptual Supervision (P-DINO): both currently rely on feature-only DINO supervision and have not tried adding attention transfer.
The Slack note flags the practical takeaway. Terrance’s tag of @U06RMQSUUGY suggests Luma may want to consider REPA termination policy in its own training pipelines — particularly relevant if internal training infra still uses always-on REPA for video DiTs, where the per-modality story (per Self-Flow: Self-Supervised Flow Matching for Scalable Multi-Modal Synthesis) is that external alignment with V-JEPA / Depth Anything actively hurts, and where the capacity-mismatch argument here predicts the same pathology.

The paper does not address the SRA/Self-Flow alternative (internal self-distillation) head-on — it accepts REPA’s external-DINO premise and patches the termination problem. A natural unasked question: would HASTE-style early termination help SRA too, or does internal self-distillation avoid the capacity-mismatch failure mode entirely because the teacher view scales with the student?