Sprint: Sparse-Dense Residual Fusion for Efficient Diffusion Transformers

TL;DR

This work tackles the quadratic training cost of diffusion transformers by introducing SPRINT, a sparse–dense residual fusion that splits computation into a dense shallow encoder path and a sparse deep middle-path, fused through a residual mechanism. By dropping up to $75\%$ of tokens in the middle, SPRINT greatly accelerates training (up to $9.8\times$ fewer GFLOPs on ImageNet-1K $256^2$) while preserving quality, aided by two-stage training and structured group-wise subsampling. It also proposes Path-Drop Guidance (PDG) to halve inference cost with improved sample quality, and demonstrates broad applicability across architectures (SiT, REPA, U-ViT) and resolutions (256^2, 512^2). Collectively, SPRINT offers a simple, effective, and general approach for scalable diffusion-transformer training with strong empirical gains and practical impact.

Abstract

Diffusion Transformers (DiTs) deliver state-of-the-art generative performance but their quadratic training cost with sequence length makes large-scale pretraining prohibitively expensive. Token dropping can reduce training cost, yet naïve strategies degrade representations, and existing methods are either parameter-heavy or fail at high drop ratios. We present SPRINT, Sparse--Dense Residual Fusion for Efficient Diffusion Transformers, a simple method that enables aggressive token dropping (up to 75%) while preserving quality. SPRINT leverages the complementary roles of shallow and deep layers: early layers process all tokens to capture local detail, deeper layers operate on a sparse subset to cut computation, and their outputs are fused through residual connections. Training follows a two-stage schedule: long masked pre-training for efficiency followed by short full-token fine-tuning to close the train--inference gap. On ImageNet-1K 256x256, SPRINT achieves 9.8x training savings with comparable FID/FDD, and at inference, its Path-Drop Guidance (PDG) nearly halves FLOPs while improving quality. These results establish SPRINT as a simple, effective, and general solution for efficient DiT training.

Figures (33)

• Figure 1: Sparse–dense residual fusion improves the efficiency of diffusion transformer training. SPRINT decouples the computationally heavy middle blocks of DiT into a sparse–deep path and a dense–shallow residual path. Notably, on ImageNet-1K $256^2$, SPRINT achieves up to 5.6$\times$ and 9.8$\times$ lower training cost compared to vanilla models, while improving generation quality.
• Figure 2: SPRINT improves visual quality over baseline with only 57$\%$ of inference TFLOPs. We present samples from two SiT-XL/2$_\text{REPA}$ models after 1M training iterations, where SPRINT is applied to one of the models. For our approach, we further incorporate the proposed Path-Drop Guidance (PDG), yielding improved FDD scores and higher visual quality compared to vanilla REPA.
• Figure 3: Training behavior of diffusion transformers. We empirically analyze the training dynamics of SiT-B/2 and its SPRINT variants under different token-drop ratios. (a) We measure the $\ell_2$ gradient norm of $f_\theta$, showing that SPRINT enables the encoder to receive stronger gradient signals from the loss. (b) SPRINT variants achieve higher and earlier CKNNA scores than SiT, indicating SPRINT learns more semantic, noise-robust representations. (c) SPRINT converges substantially faster and to lower FID than SiT, with the gap further widening at higher drop ratios (up to 75$\%$), highlighting both the effectiveness and efficiency of our framework.
• Figure 4: Roles of dense--shallow and sparse--deep features. Dense--shallow features preserve local textures but lose global structure, while sparse--deep features capture global shapes but distort local details. Fusing both recovers high-quality outputs with coherent semantics and fine detail.
• Figure 4: Effect of dense--sparse residuals on FID.