PaperlayerView exact revision on arXiv

Source-linked AI summary

Sprint: Sparse-Dense Residual Fusion for Efficient Diffusion Transformers

Dogyun Park, Moayed Haji-Ali, Yanyu Li, Willi Menapace, Sergey Tulyakov, Hyunwoo J. Kim, Aliaksandr Siarohin, Anil Kag

arXiv:2510.21986v2cs.CV

TL;DR

DiTs’ quadratic training cost makes large-scale pretraining expensive, and naïve token dropping can harm representations and full-token generalization. SPRINT combines dense shallow features with sparse deep features and a two-stage training schedule. On ImageNet-1K, it reduces training cost by up to 9.8× while matching or surpassing strong baselines, and its PDG halves inference cost while improving sample quality.

  • Problem

    DiTs’ quadratic training cost makes large-scale pretraining expensive, while naïve token dropping can degrade representations and generalization to full-token inference.

    Source: §1 · p. 1

  • Method

    SPRINT fuses dense shallow features with sparse deep features after aggressive token dropping, then uses short full-token fine-tuning after masked pre-training.

    Source: §1 · p. 2

  • Results

    SPRINT reduces training cost by up to 9.8× while matching or surpassing strong-baseline quality, and PDG halves inference cost while improving sample quality.

    Source: §1 · p. 2, §5 · p. 10

  • Takeaways & Limitations

    SPRINT is a simple, effective, and general approach for efficient DiT training.

    Source: §5 · p. 10

    Available computational resources prevented experiments on large-scale text-to-image or video diffusion models.

    Source: H LIMITATION AND FUTURE WORK · p. 36

Abstract

from arXiv · show

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\"ive 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.

1 INTRODUCTION

DiTs’ quadratic training cost makes large-scale pretraining expensive, while naïve token dropping can harm representations and full-token inference performance. SPRINT combines dense shallow processing with sparse deep processing, then uses short full-token fine-tuning to address the train–inference gap.

1 INTRODUCTION

  • 9.8× fewer training GFLOPs than standard SiT training, with similar or better quality on ImageNet-1K 256² class-conditional generation.The figure also reports up to 5.6× and 9.8× lower training cost than vanilla models while improving generation quality.Source: Figure 1 · p. 1, §1 · p. 2
  • Some token-dropping methods add substantial parameters, while others support only moderate ratios and break down at 75%.The introduction also notes that naïve dropping can degrade representations and generalization with full-token inference inputs.Source: §1 · pp. 1–2
  • SPRINT processes all tokens in shallow layers, drops typically 75% before the middle blocks, and fuses dense shallow features with sparse deep features.Short full-token fine-tuning follows long masked pre-training to let the middle blocks adapt to dense inputs and close the train–inference gap.Source: §1 · p. 2
  • PDG uses dense shallow features for the unconditional estimate, reducing guidance-sampling cost by nearly 50% while improving generation quality.Figure 2 reports improved FDD and visual quality versus vanilla REPA with 57% of inference FLOPs.Source: §1 · p. 2, Figure 2 · p. 3
  • Dense–shallow features are more noise-invariant and semantically expressive, score higher on CKNNA than vanilla DiT, and specialize in local rather than deep-path global semantics.These representation findings are reported with Figs. 6, 3b, and 4, respectively.Source: §1 · p. 2
  • SPRINT is reported to apply across SiT and UViT architectures, SD and FLUX VAE latent spaces, and 256 and 512 resolutions, with further gains when combined with REPA (Yu et al., 2024).The authors describe it as simple, architecture-agnostic, and complementary to alignment-based methods.Source: §1 · p. 2

2 RELATED WORK

Prior work accelerates DiT training through representation alignment or token dropping. The cited token-dropping approaches have limits at aggressive drop rates, while SPRINT is described as robust at high rates and compatible with alignment methods.

2 RELATED WORK

  • REPA aligns intermediate DiT activations with DINOv2 features, while Lee et al. (2025) use contrastive loss and Wang & He (2025) use dispersive loss.HASTE (Wang et al., 2025) reports that alignment signals can conflict with diffusion objectives and destabilize training; the paper says these objectives can be combined with its scheme.Source: §2 · p. 3
  • Prior token-dropping methods include masked-token reconstruction, patch mixing, and bypassing token subsets, but are reported to degrade at 75% dropping and to be difficult to pair with alignment losses.The passage identifies MDTv2 (Gao et al., 2023), MaskDiT, MicroDiT, and TREAD among these approaches.Source: §2 · p. 3
  • SPRINT is presented as alignment-friendly and robust at high drop rates, complementing alignment objectives that can be combined with token dropping.The related-work discussion says combining these objectives can further boost performance (Tab. 2).Source: §2 · p. 3

3 SPRINT: SPARSE-DENSE RESIDUAL FUSION FOR EFFICIENT DIFFUSION TRANSFORMERS

SPRINT reduces DiT training compute by keeping dense shallow features and processing sparse deep features, then fusing them for full-token prediction. Its design also supports efficient Path-Drop Guidance and structured token sampling.

3.1 PRELIMINARIES

  • DiT attention cost grows quadratically with token count, so dropping tokens reduces compute, but standard DiTs process all tokens at every layer.The paper motivates specialization by noting that deeper-layer representations become redundant as features shift from local patterns to global semantics.Source: §3.1 · p. 4, §3.2 · p. 4

3.2 BOTTLENECK IN STANDARD DIT TRAINING

  • The encoder processes dense tokens for local evidence, deeper layers use sparse tokens for global relationships, and final layers restore all tokens for prediction.SPRINT divides the DiT into encoder, middle blocks, and decoder, with dense–sparse fusion connecting their features.Source: §3.2 · p. 4, §3.3 · p. 4

3.3 SPARSE–DENSE RESIDUAL FUSION

  • SPRINT restricts expensive middle blocks to sparse tokens while retaining dense shallow features, accelerating convergence and remaining stable at aggressive drop ratios.The decoder predicts all tokens despite masked deep-path inputs; the paper links this to stronger gradient flow, richer representations, and faster FID improvement.Source: §3.3 · pp. 4–5
  • SPRINT’s dense–shallow path preserves local textures, while its sparse–deep path captures global shapes; fusion recovers coherent semantics and fine detail.The encoder sends dense features through a residual path, while the middle blocks process a sparse subset before fusion with the dense features for full-token prediction.Source: §3.3 · pp. 4–5
  • After sparse pre-training, brief full-token fine-tuning adapts the middle blocks to dense inputs and addresses the potential train–inference gap.The paper gives 100K–200K fine-tuning iterations as an example, compared with 1M–4M pre-training iterations.Source: §3.3 · p. 5

3.4 EFFICIENT PATH-DROP GUIDANCE (PDG)

  • Path-Drop Guidance uses a full conditional pass and bypasses the middle blocks for the unconditional estimate, nearly halving guidance FLOPs and latency per step.The unconditional estimate uses the dense shallow path with [MASK] tokens in place of the middle-block output.Source: §3.4 · pp. 5–6

3.5 STRUCTURED GROUP-WISE TOKEN SUBSAMPLING

  • Structured group-wise subsampling selects k tokens from each n × n group, maintaining local coverage while avoiding fixed sampling patterns.For images, the paper uses n = 2 and k = 1, corresponding to a 75% drop ratio.Source: §3.5 · p. 6

4 EXPERIMENT

Across ImageNet experiments, SPRINT improves generation quality while reducing training and inference costs, with tests spanning architectures, resolutions, ablations, and fine-tuning. Its sparse-pretraining representations support strong results after full-token adaptation.

4.1 EXPERIMENTAL DETAILS

  • Experiments primarily use SiT-B/2 or SiT-XL/2 on ImageNet-1K at 256 × 256, with pretrained Stable Diffusion or Flux VAEs and a 75% token-drop ratio.Generation quality is assessed with FDD, FID, IS, precision, and recall; training and inference efficiency are measured with FLOPs.Source: §4.1 · pp. 6–7

4.2 SYSTEM-LEVEL COMPARISON

  • At 400K SD-VAE iterations, SPRINT improves FDD by +88.5 over SiT-XL/2 with 1.32× fewer FLOPs; by 1M iterations, it achieves a 1.95× computational speedup.Across REPA and U-ViT at 400K iterations, it improves FDD by +45.1 and +63.4, respectively; visual analysis reports coherent global structure and fine details at 100K iterations.Source: §4.2 · pp. 7–8, Table 2 · p. 7

4.3 COMPARISON WITH STATE-OF-THE-ART MODELS

  • At 400 epochs, SPRINT reaches 75.4 FDD versus 79.5 for the improved SiT baseline using 53% of its training FLOPs; PDG reaches 58.4 FDD at 57% of inference cost.With REPA, SPRINT reduces FDD from 78.8 to 75.6 using 71% of training FLOPs; with PDG at 400 epochs, it surpasses the 800-epoch official REPA model by 17.8 FDD and 0.21 FID using 27% of training FLOPs.Source: §4.3 · p. 8, Table 3 · p. 8

4.4 ANALYSIS

  • PDG outperforms both CFG baselines at nearly half the inference cost; at 512 × 512 it achieves 1.96 FID versus 2.63 for SiT using 50% of training compute.At 10 sampling steps, SPRINT improves FID from 7.37 to 6.29 and FDD from 205.2 to 174.5 versus SiT-XL/2.Source: Figure 7 · p. 9, §4.4 · pp. 9–10
  • Removing either parallel path sharply degrades performance; the default middle-block depth achieves the best FID (27.3) at the lowest reported cost (7.47G).Removing the dense shallow path raises FID from 27.5 to 85.1; configurations shifting layers out of the middle blocks raise cost and degrade FID to 29.1 or 49.2, while reducing either path to one block yields FID 61.5 or 44.4.Source: §4.4 · p. 9
  • At a 75% drop ratio, structured sampling improves FID from 30.1 to 27.5, while increasing the ratio to 87.5% raises FID to 50.2.Across ratios from 0% to 75%, FID decreases from 54.1 to 27.5.Source: §4.4 · p. 9

4.5 BENEFITS OF FINE-TUNING

  • After 20K full-token fine-tuning steps, SPRINT recovers over 94% of the FID improvement achieved at 200K, and gains beyond 50K steps are marginal.Before fine-tuning, samples are class-aligned with coherent global structure but can lack high-frequency details; 20K steps largely restore those details.Source: §4.5 · p. 10, Figure 9 · p. 10, Figure 8 · p. 10

5 CONCLUSION

SPRINT combines dense-shallow and sparse-deep features through residual fusion, enabling up to 75% token dropping while preserving representation quality. On ImageNet-1K, it reduces training cost by up to 9.8×, and PDG halves inference cost while improving sample quality.

5 CONCLUSION

  • SPRINT’s two-stage schedule pairs masked pre-training with short full-token fine-tuning to close the train–inference gap.The framework is described as architecture-agnostic and applicable across architectures, resolutions, and alignment methods.Source: §5 · p. 10

REPRODUCIBILITY

The paper provides detailed experimental configurations and notes that its training code was unavailable at submission time.

REPRODUCIBILITY

  • Appendix materials provide hyperparameters, schedules, model configurations, sampling steps, and compute resources, but training code was not released at submission time.The authors state that standardized DiT and SiT setups plus the experimental details should allow independent reproduction.Source: REPRODUCIBILITY · p. 11
  • The experiments follow the established DiT and SiT setups, which the authors identify as widely adopted in diffusion research.The passage presents these setups alongside the provided experimental details as support for independent reproduction.Source: REPRODUCIBILITY · p. 11
  • The authors state that independent reproduction should be possible using the provided details, despite training code being unavailable at submission time.Source: REPRODUCIBILITY · p. 11

A ANALYSIS DETAILS

The analysis links SPRINT’s training behavior to semantic alignment and examines its feature paths. Additional sections explain how the method is integrated with SiT, REPA, and U-ViT.

A ANALYSIS DETAILS

  • SPRINT shows stronger encoder gradient propagation than SiT as sparsity increases, alongside higher CKNNA scores and lower FID values.CKNNA compares encoder features on noisy inputs with DINOv2 features from corresponding clean images; higher scores indicate more semantically meaningful, noise-robust representations.Source: A ANALYSIS DETAILS · p. 15
  • CKNNA and FID exhibit a strong negative correlation in the reported experiments.The analysis reports FID from 10K generated images and uses PCA of encoder and middle-block features to examine their evolution across diffusion timesteps.Source: A ANALYSIS DETAILS · p. 15
  • For single-path sample analysis, one path’s representation is duplicated to replace the other before concatenation and decoding.SiT’s homogeneous transformer-block stack is decomposed into encoder, middle, and decoder sections for SPRINT.Source: A ANALYSIS DETAILS · p. 15, B SPRINT WITH DIFFERENT DIFFUSION… · p. 15

B SPRINT WITH DIFFERENT DIFFUSION TRANSFORMERS

  • For REPA integration, the DINOv2 target features receive the same token-dropping mask as sparse middle-block states; U-ViT retains its original skip connections while SPRINT adds a dense residual path and drops tokens in the middle section.The REPA projection is placed at the corresponding location in the sparse middle block.Source: B SPRINT WITH DIFFERENT DIFFUSION… · pp. 15–16

C.1 TRAINING DETAILS

SPRINT training uses a two-stage schedule, with long sparse pre-training followed by full-token fine-tuning; evaluation reports standard image-generation metrics and contrasts Path-Drop Guidance with classifier-free guidance.

C.1 TRAINING DETAILS

  • A 75%-drop pre-training stage is followed by 100K full-token fine-tuning iterations, with a single fusion projection adding approximately 0.3% of the original model size.The setup uses pre-computed Stable Diffusion and Flux VAE latents without data augmentation.Source: §C.1 · p. 16
  • During both training phases, sparse-deep features are independently dropped with 10% probability and replaced by mask tokens, alongside standard class-condition dropping.The passage also reports bf16 mixed precision and gradient clipping at 1.0 for both phases.Source: §C.1 · p. 17

C.2 EVALUATION DETAILS

  • Evaluation reports FDD, FID, IS, and Precision/Recall, typically using 50K generated samples.FDD substitutes DINOv2 features for FID's Inception features, while FID measures feature distance between real and generated image distributions.Source: §C.2 · pp. 17–18
  • Path-Drop Guidance replaces CFG's unconditional full-model branch with a weaker network, while retaining the guidance-scale formulation.The guidance scale is denoted by w.Source: §C.2 · p. 18

D COMPUTATION ANALYSIS

SPRINT reports faster training and lower memory use than SiT-XL/2, with compute measured through forward-pass FLOPs and runtime comparisons at a fixed 75% token-drop rate.

D COMPUTATION ANALYSIS

  • Training FLOPs are estimated from averaged forward-pass measurements over 100 iterations and scaled by total iterations; inference FLOPs sum forward passes across sampling timesteps.Runtime comparisons use a fixed 75% drop rate for token-dropping methods.Source: D COMPUTATION ANALYSIS · p. 18
  • 5.2 iters/sec at 256² is more than 2× SiT's 2.5 iters/sec; at 512², SPRINT reaches 2.01 iters/sec versus SiT's 0.79 iters/sec.At both resolutions it also outpaces MaskDiT, MicroDiT, and Tread.Source: D COMPUTATION ANALYSIS · p. 19
  • 19.6 GB versus SiT-XL/2's 29.6 GB at 256², and 37.9 GB versus 77.7 GB at 512², are SPRINT's reported training-memory requirements.The corresponding reductions are 33.8% and 51.2%.Source: D COMPUTATION ANALYSIS · p. 19

E.1 BASELINE DETAILS ON TABLE 1

The baseline discussion contrasts SPRINT with progressive training and token-dropping alternatives, emphasizing reported performance and efficiency differences at high drop rates.

E.1 BASELINE DETAILS ON TABLE 1

  • Progressive training uses 25.8 versus SPRINT's 31.5 GFLOPs (×10^9) per iteration, yet lags SPRINT in performance and fails to match baseline SiT results.It pretrains at 128 × 128 and fine-tunes at 256 × 256 with resized positional embeddings.Source: §E.1 · p. 19
  • MicroDiT adds six patch-mixing transformer blocks and substantially increases compute and parameter count, yet underperforms SPRINT.Its deferred masking lets local patch tokens mix before dropping.Source: §E.1 · pp. 19–20
  • At 75% dropping, Tread fails while SPRINT's dense residual path supplies complete local noise information to the decoder.Tread routes randomly dropped tokens past the middle layers, leaving those blocks to encode local noise information for velocity estimation.Source: §E.1 · p. 20

E.2 MORE DISCUSSION ON OTHER BASELINES

  • SPRINT reaches 61.8 FDD and 2.01 FID in 200 epochs, while MaskDiT needs 1600 epochs for 65.4 FDD and 2.28 FID.MaskDiT adds a reconstruction decoder, grows from 675M to 730M parameters, and is limited to moderate drop ratios such as 50%.Source: §E.2 · p. 20
  • MDT uses an encoder-decoder, side interpolator, long skip connections, and masked-token reconstruction, but does not work well at 75% dropping.The passage describes MDT's operating range as moderate ratios such as 30% and 50%.Source: §E.2 · p. 20

F.1 IMAGENET 512X512 EXPERIMENT

On ImageNet at 512², SPRINT reports lower training and inference compute than a strong baseline; at fewer sampling steps it improves reported metrics, and a 256² visual comparison reports improved quality with PDG.

F.1 IMAGENET 512X512 EXPERIMENT

  • 184.8 training TFLOPs at 400 epochs versus SiT-XL's 366.6 at 600 epochs; PDG uses 0.471 versus 0.952 inference TFLOPs while improving FDD.The passage reports comparable or better generation quality overall at 512².Source: §F.1 · p. 20

F.2 PERFORMANCE WITH FEW-STEP GENERATION

  • At 50 steps, SPRINT improves FID from 2.34 to 2.19 and IS from 244.0 to 267.7; at 10 steps, FID goes from 7.37 to 6.29 and IS from 187.3 to 211.3.The comparison is between SiT-XL/2 and SiT-XL/2 + SPRINT across NFEs.Source: §F.2 · pp. 20–21, Table 11 · p. 21

G.1 VISUAL COMPARISON ON IMAGENET 256 × 256

  • SPRINT with PDG is reported to improve visual quality over REPA using only 57% of inference FLOPs.The comparison shows SiT-XL/2 + REPA models after 1M training iterations, with SPRINT applied to one model.Source: Figure 10 · p. 22

G.2 UNSELECTED GENERATED RESULTS BY SPRINT ON IMAGENET 256 × 256

Additional feature visualizations report clearer semantic structure in SPRINT features than in the baseline. The appendix also presents unselected ImageNet generations, a 512 × 512 generation example, and future directions alongside a stated resource limitation.

G.2 UNSELECTED GENERATED RESULTS BY SPRINT ON IMAGENET 256 × 256

G.3 GENERATED RESULTS BY SPRINT ON IMAGENET 512 × 512

  • Figure 33 presents generation results for SiT-XL/2 + SPRINTPDG using path-drop guidance with w = 3.0.Source: Figure 33 · p. 34

G.4 ADDITIONAL FEATURE PCA VISUALIZATION

  • SPRINT features show clearer semantic structure than baseline features in the additional PCA examples.The visualizations compare dense-shallow and sparse-deep features from SPRINT with a standard full-token SiT-XL/2 model.Source: Figure 34 · p. 35, §G.4 · p. 35

H LIMITATION AND FUTURE WORK

  • Limited computational resources prevented experiments on large-scale text-to-image or video diffusion; integrating efficient attention and scalable training is proposed as future work.Source: H LIMITATION AND FUTURE WORK · p. 36

Related papers

Loading 2510.21986v2…