Speedrunning ImageNet Diffusion

TL;DR

SR-DiT (Speedrun Diffusion Transformer) is a community-style “kitchen sink” study that stacks every recent diffusion-training-efficiency lever onto a small SiT-B/1 backbone (140M params) and asks which combinations actually compose. Starting from REG + INVAE (the REPA-E semantic VAE), the author adds SPRINT 75%-drop token routing, RMSNorm, 2D RoPE, QK-Norm, value-residual learning, Contrastive Flow Matching, time-shifting, and balanced label sampling — reaching FID 3.49 / KDD 0.319 on ImageNet-256 at 400K iterations, without classifier-free guidance, in 80 H200-GPU-hours [§5.2, Table 1]. The headline is that a 140M model trained ~10 hours on one node matches numbers REG previously needed SiT-XL (685M, FID 3.4) for. The negative-results appendix is the other half of the contribution: SwiGLU, Muon, Prodigy, Equilibrium Matching, x-prediction, SARA loss, dispersive loss, and a time-weighted CFM variant all fail to help and are documented as such.

Key claims

Stacking eight independently-published techniques on SiT-B/1 + INVAE reaches FID 3.49 / KDD 0.319 on ImageNet-256 at 400K iterations with 140M parameters, no CFG, no PDG — comparable to REG’s reported FID 3.4 at 685M params trained much longer [§5.2 Table 1, Abstract].
The dominant gain comes from three pieces: representation alignment / entanglement (REPA → REG), token routing (SPRINT at 75% drop, +6 FID), and the semantic VAE INVAE; SPRINT alone moves FID from 10.56 to 4.58 on the REG + INVAE baseline [§5.2 Table 1, §5.3].
The smaller architectural deltas — RMSNorm, RoPE, QK-Norm, value residual — are each individually modest (≤0.5 FID) but stack consistently positively, taking FID from 4.58 → 3.64 in sequence [§5.2 Table 1].
Contrastive Flow Matching (CFM, λ=0.05) added on top of all architectural changes improves KDD from 0.353 to 0.319 with little FID change, the largest single perceptual-metric jump observed [§5.2 Table 1].
RELU² beats GELU as a standalone activation (FID 3.81 vs 4.02) but is incompatible with value-residual learning — combining them gives no additional gain over GELU + value-residual (FID 3.64) [§B.1 Table 5].
Muon optimizer fails catastrophically in this configuration (FID 48.70 at 400K, KDD 1.378) despite lower early-training loss; Prodigy underperforms Adam (FID 4.67); the author keeps Adam [§B.5 Table 8].
x-prediction with velocity loss (the JiT recipe for pixel diffusion) underperforms vanilla flow-matching in latent space with INVAE, and Equilibrium Matching could not be reproduced on this architecture at all [§B.4].
A “lopsided leaky RELU²” variant the author had seen multiple researchers try produces marginally better training loss but worse evaluation metrics across the board (FID 4.11 vs 3.81, KDD 0.366 vs 0.351) — documented explicitly to deter further reproductions [§B.1, “Lopsided Leaky RELU²”].
Time-weighted CFM (down-weighting the contrastive term at low noise) was the author’s hypothesis to improve CFM near clean images; it slightly underperforms uniform CFM (FID 3.82 vs 3.61), showing the contrastive signal is useful across all noise levels [§B.4 Table 7].
SR-DiT-B/1 transfers to ImageNet-512 at the same setup with FID 4.23 / KDD 0.306 at 400K iterations (no CFG/PDG), versus DiT-XL/2 at 20.94 and U-DiT-B at 15.39 for comparable training budgets [§5.2 Table 2].

Method

The base is SiT-B/1 (patch size 1 because INVAE has 16× spatial compression vs SD-VAE’s 8×) with the REG (Representation Entanglement for Generation) objective: a standard velocity-prediction flow-matching loss on latents, REPA-style projection-loss cosine alignment of intermediate hidden states to a frozen DINOv2 encoder’s per-token features, plus a velocity-prediction loss on a separately-diffused DINOv2 CLS embedding that is concatenated as a learned class token. INVAE replaces SD-VAE as the tokenizer.

On top of REG + INVAE, the additions are: (i) SPRINT token routing — two dense encoder blocks → routed middle blocks operating on 25% of tokens → fusion via concatenation + projection back to full-length with a learned [MASK] token → two dense decoder blocks; (ii) RMSNorm everywhere LayerNorm sat, including the QK-Norm internal normalizers and the CLS-embedding pre-concatenation norm; (iii) 2D RoPE adapted to support sparse/routed token subsets via per-token rope_ids, with the CLS token explicitly excluded from rotation; (iv) QK-Norm; (v) value-residual learning that caches the first attention block’s value vectors and mixes them into every subsequent block’s values via a learned scalar; (vi) Contrastive Flow Matching as an auxiliary loss with λ=0.05; (vii) time-shifting using the LightningDiT/RAE-style monotone transform $t \mapsto (\alpha t)/(1 + (\alpha-1)t)$ with shift derived from full latent dimension; (viii) balanced label sampling at evaluation time. Total loss is $\mathcal{L} = \mathcal{L}_{\text{vel}} + \lambda_{\text{proj}}\mathcal{L}_{\text{proj}} + \lambda_w \mathcal{L}_{\text{vel-CLS}} + \lambda_{\text{cfm}}\mathcal{L}_{\text{cfm}}$ with $\lambda_{\text{proj}}=0.5$ , $\lambda_w=0.1$ , $\lambda_{\text{cfm}}=0.05$ [§4.2 Eq. 10].

Training cost is 10 hours on 8× H200 for ImageNet-256 (80 GPU-hours), 15 hours for ImageNet-512 (120 GPU-hours).

Results

ImageNet-256, 400K iterations, no CFG, balanced sampling: SR-DiT-B/1 reaches FID 3.49, sFID 5.67, IS 221.2, Precision 0.808, Recall 0.546, KDD 0.332 with 140M params [§5.2 Table 1].
For reference, REG with SiT-XL (685M) reports FID 3.4 in its original paper at similar iteration counts; REPA with SiT-XL needs 4M steps to reach FID 5.9; the SiT-B/2 + REG baseline at 132M reaches FID 15.2 at 400K [§5.2 baselines].
The full ablation ladder (REG + INVAE → final) on the same SiT-B/1: 10.56 → 4.58 (+SPRINT) → 4.56 (+RMSNorm) → 4.09 (+RoPE) → 4.02 (+QK-Norm) → 3.64 (+ValueResidual) → 3.61 (+CFM) → 3.56 (+TimeShift) → 3.49 (+BalancedSampling) [Table 1].
ImageNet-512, 400K iterations, no CFG: FID 4.23, sFID 5.46, IS 192.3, KDD 0.306; comparison points (also no CFG, 400K) are DiT-XL/2 at FID 20.94 and U-DiT-B at 15.39 [§5.2 Table 2].
Intermediate checkpoints: SR-DiT-B/1 hits FID 9.17 at 100K and FID 4.91 at 200K iterations — already beating REG SiT-B/2 (FID 15.2 at 400K) at half the iterations and a third of the activation budget [§A.1 Table 3].

Why it’s interesting

This is the kind of paper Luma’s training-efficiency thread has been waiting for. It is the empirical composition of Diffusion training efficiency‘s key recipes — REPA-style external-DINO alignment, token routing, and semantic-VAE latents — in one carefully-ablated pipeline, run by an outside researcher with a sponsored single-node compute budget. The headline result (140M → 3.49 FID at 80 GPU-hours, no CFG) operationalizes the implicit claim of the Representation Autoencoders page: that pretrained visual structure (INVAE) plus the right intermediate-supervision recipe (REG) plus a sensible architectural prior (SPRINT + RoPE + value-residual) collapses the compute frontier for class-conditional ImageNet generation. It extends TREAD: Token Routing for Efficient Architecture-agnostic Diffusion Training in the most direct way possible — TREAD reported a 25.39× speedup at 50% token drop; SPRINT at 75% drop is the next step on that axis and the author confirms it composes with REG. It also adds two negative-result data points to the wiki’s open questions: Muon’s published wins do not transfer to small ImageNet diffusion in this exact recipe, and RELU² and value-residual saturate together rather than stacking. The absence of any internal-self-distillation alternative (SRA / Self-Flow: Self-Supervised Flow Matching for Scalable Multi-Modal Synthesis) in the SR-DiT stack is conspicuous — those papers argue against the external-DINO substrate this work doubles down on. The natural follow-up is “SR-DiT minus REG plus Self-Flow” on the same setup; that experiment is not in this paper.