Source-linked AI summary
T-Stitch: Accelerating Sampling in Pre-Trained Diffusion Models with Trajectory Stitching
Zizheng Pan, Bohan Zhuang, De-An Huang, Weili Nie, Zhiding Yu, Chaowei Xiao, Jianfei Cai, Anima Anandkumar
TL;DR
High-quality diffusion sampling is expensive because large models typically require many costly denoising steps, and existing efficiency approaches often use one model throughout. T-Stitch uses a smaller pretrained model early in the trajectory and switches to a larger one later. Across architectures and samplers it reports speed gains with little quality loss; on DiT-XL, replacing 40% of early steps with DiT-S yields around 1.5× lossless speedup.
Problem
High-quality diffusion sampling is costly because it typically requires a large model and many denoising steps, while existing efficiency methods assume the same model is used throughout.
Source: §1 · p. 1Method
T-Stitch uses a smaller pretrained DPM in early denoising steps and switches to a larger DPM later, without retraining.
Source: §1 · pp. 1, 3Results
Takeaways & Limitations
T-Stitch complements existing fast-sampling methods, and can accelerate stylized Stable Diffusion while improving prompt alignment.
Source: §1 · p. 2T-Stitch requires a sufficiently optimized smaller model trained on the same data distribution, adds slight memory usage, and its speedup is bounded by the smaller model’s efficiency.
Source: §5 · p. 8
Abstract
from arXiv · showhide
Sampling from diffusion probabilistic models (DPMs) is often expensive for high-quality image generation and typically requires many steps with a large model. In this paper, we introduce sampling Trajectory Stitching T-Stitch, a simple yet efficient technique to improve the sampling efficiency with little or no generation degradation. Instead of solely using a large DPM for the entire sampling trajectory, T-Stitch first leverages a smaller DPM in the initial steps as a cheap drop-in replacement of the larger DPM and switches to the larger DPM at a later stage. Our key insight is that different diffusion models learn similar encodings under the same training data distribution and smaller models are capable of generating good global structures in the early steps. Extensive experiments demonstrate that T-Stitch is training-free, generally applicable for different architectures, and complements most existing fast sampling techniques with flexible speed and quality trade-offs. On DiT-XL, for example, 40% of the early timesteps can be safely replaced with a 10x faster DiT-S without performance drop on class-conditional ImageNet generation. We further show that our method can also be used as a drop-in technique to not only accelerate the popular pretrained stable diffusion (SD) models but also improve the prompt alignment of stylized SD models from the public model zoo. Code is released at https://github.com/NVlabs/T-Stitch
1. Introduction
High-quality diffusion sampling is costly because large models require many denoising steps, while existing efficiency approaches use one model throughout. T-Stitch assigns smaller models to early steps and larger models to later steps, aiming to improve speed without sacrificing quality.
1. Introduction
- 40% of early DiT-XL steps can use DiT-S with no FID degradation, yielding around 1.5× lossless speedup.Figure 1 compares FID as the fraction of early DiT-S steps increases; its setup uses DDIM with 100 timesteps and guidance scale 1.5.Source: Figure 1 · p. 2, §1 · p. 2
- The method is motivated by similar sampling trajectories across models trained on the same distribution and small models’ ability to form early global structure.The paper associates early denoising with low-frequency components and later steps with high-frequency detail.Source: §1 · p. 2
- For stylized Stable Diffusion models, pairing a smaller pretrained SD with the larger model accelerates sampling and improves prompt alignment.The paper suggests stylization fine-tuning may degrade prompt alignment; Figure 2 illustrates examples including “New York City” and “tropical beach.”Source: §1 · pp. 2–3, Figure 2 · p. 2
- T-Stitch replaces the large DPM with a smaller one for early denoising steps, then switches back to the large DPM, without retraining.The approach dynamically allocates computation across denoising steps and is presented as complementary to existing efficient sampling methods.Source: §1 · pp. 1–3
- Experiments report consistent efficiency gains across architectures and samplers, while T-Stitch complements existing fast-sampling methods.The paper also reports a training-free quality-efficiency frontier that improves on training-based model-stitching methods.Source: §1 · pp. 2–3
2. Related Works
Prior work accelerates diffusion sampling by reducing per-step cost, reducing the number of steps, or using experts and stitched models. T-Stitch instead explores speed-quality trade-offs by assigning differently sized pretrained models to different trajectory intervals.
Efficient diffusion models
- Diffusion efficiency methods reduce computation per step, reduce sampling steps, or parallelize sampling; T-Stitch complements these approaches by using a small pretrained model early and a large one later.Per-step methods include compression and lightweight architectures; step-count methods include distillation and improved differential-equation solvers.Source: §2 · p. 3
Multiple experts in diffusion models
- Prior timestep-specialized experts either add large denoisers without reducing cost or use same-sized models, whereas T-Stitch trades off small and large DPMs across steps.Examples include expert denoisers, frequency gating, and small denoisers allocated by interval.Source: §2 · p. 3
Stitchable neural networks
- SN-Net and SN-Netv2 stitch pretrained models of different scales using learned stitching layers; T-Stitch is presented as a simpler, more general alternative for diffusion speed-quality trade-offs.SN-Net jointly trains the stitches, while SN-Netv2 extends them to dense prediction tasks.Source: §2 · p. 3
3. Method
The method builds on continuous-time score-based diffusion and classifier-free guidance, then stitches pretrained denoisers of different sizes across sampling steps. Early cross-model latent similarity supports using a smaller model first, with step allocation controlling speed–quality trade-offs.
3.1. Preliminary
- Diffusion models iteratively denoise high-variance Gaussian inputs, using a learned score function to define the probability-flow ODE.The paper focuses on continuous-time score-based models and denoisers.Source: §3.1 · pp. 3–4
- Classifier-free guidance jointly uses conditional and unconditional predictions, with a guidance scale steering samples toward the conditioning signal.The approach trains both conditional and unconditional behavior in one network by replacing conditioning with a null embedding.Source: §3.1 · p. 4
3.2. Trajectory Stitching
- Early latent embeddings from different DiT models reach nearly 100% cosine similarity, motivating small-model use at the start of sampling.The comparison averages results over 32 images.Source: §3.2 · p. 4, Figure 3 · p. 5
- T-Stitch assigns the first fraction of sampling steps to a faster small denoiser and the remaining steps to a large denoiser, tuning the trade-off through r1.Increasing r1 shifts more computation to the small model; the models should be faster and sufficiently optimized, and trained on the same data or similar distributions.Source: §3.2 · p. 4
- Using additional denoisers, such as a medium model in an intermediate interval, expands the available trade-offs and can improve the Pareto frontier.A precomputed lookup table can identify configurations that meet a compute budget.Source: §3.2 · pp. 4–5
- Unlike multi-expert approaches that allocate multiple large denoisers, T-Stitch uses pretrained models of different sizes without retraining.The authors distinguish its diffusion-specific sampling strategy from speculative decoding for autoregressive models.Source: §3.2 · p. 5
4. Experiments
The experiments test T-Stitch across DiT, U-Net, and Stable Diffusion models, then examine its performance across sampling steps, samplers, and other acceleration methods. Across these settings, using smaller models early can improve speed while retaining comparable quality, with model and setup affecting the trade-off.
4.1. DiT Experiments
- Using a smaller DiT in the first 40–50% of steps brings only a minor performance drop across the XL/S, XL/B, and B/S combinations.The best and worst performance are roughly bounded by the smallest and largest models.Source: §4.1 · p. 5, Figure 5 · p. 6
- At about 10 seconds, three-model DiT stitching achieves 1.7× speedup with comparable FID (9.21 vs. 9.19) and Inception Score (243.82 vs. 245.73).The S→B→XL schedules produce a smooth Pareto frontier.Source: Figure 6 · p. 6, §4.1 · p. 6
4.2. U-Net Experiments
- With DDIM and 100 steps, the first ~50% of U-Net sampling steps can use LDM-S with comparable or better FID and Inception Scores.Time cost decreases approximately linearly as more LDM-S steps are used.Source: §4.2 · p. 6, Table 1 · p. 7
4.3. Text-to-Image Stable Diffusion
- For SD v1.4, BK-SDM Tiny can take the early 30% of steps without significant Inception Score or CLIP Score loss, while FID improves.For stylized SDs such as Inkpunk, stitching also yields natural style interpolation and can improve prompt alignment.Source: §4.3 · pp. 6–7, Table 2 · p. 7, Figure 7 · p. 7
4.4. Ablation Study
- Across step counts, DiT-S can take 40% of early steps without significant performance drop, and T-Stitch remains complementary to methods that reduce sampling steps.The authors attribute this to small models handling early low-frequency components while DiT-XL handles later high-frequency detail.Source: §4.4 · p. 7, Figure 9 · p. 8
- T-Stitch shows consistent efficiency gains with DDPM, DDIM, and DPM-Solver++ across different sampler and step settings.These samplers provide techniques orthogonal to T-Stitch’s allocation of models across denoising steps.Source: §4.4 · p. 7, Figure 8 · p. 8
- Under DDIM with 100 steps, T-Stitch achieves a clearer speed–quality Pareto advantage than model stitching across the compared metrics.Figure 10 compares the two approaches.Source: §4.4 · pp. 7–8
- T-Stitch is training-free, works with DiT and U-Net denoisers and pruned or step-distilled models, and complements caching, quantization, VAE acceleration, and token merging.These approaches can accelerate the model used for the large-model portion of the trajectory as well.Source: §4.4 · p. 8
5. Conclusion
The conclusion presents T-Stitch as a general way to improve the speed–quality trade-off of pretrained diffusion-model sampling by using smaller models early. It also identifies model requirements and extra memory use as practical constraints, and points to further work on specialized models and allocation guidance.
5. Conclusion
- T-Stitch uses pretrained smaller models at the start of denoising to improve speed and quality trade-offs over using one large diffusion model.Experiments report consistent efficiency gains across architectures, samplers, and Stable Diffusion models.Source: §5 · p. 8
- The authors identify designing specialized early- and late-stage models and developing better trade-off guidance as future work.They also call for deeper analysis of prompt alignment in stylized Stable Diffusion models.Source: §5 · p. 8
- T-Stitch requires a sufficiently optimized smaller model trained on the same data distribution, adds some memory use, and has speedup bounded by that model’s efficiency.The authors recommend it when a much faster small model is available.Source: §5 · p. 8
Societal Impact
The societal-impact passage says the use of public pretrained models avoids training costs and can lower deployment emissions, while placing privacy and misuse outside the paper’s scope.
Societal Impact
- Using public pretrained models avoids training costs while accelerating diffusion sampling for image generation.The authors connect this deployment cost reduction to lower carbon emissions.Source: Societal Impact · p. 9
- Generated images depend on user prompts and the diffusion models chosen.Source: Societal Impact · p. 9
- The paper does not address potential privacy concerns or misuse of generative models.The authors state these issues fall outside the work’s scope.Source: Societal Impact · p. 9
A. Appendix
The appendix points readers to additional deployment guidance, analyses, examples, and compatibility evaluations. It also reports supplementary evaluations of T-Stitch alongside caching and token merging.
A. Appendix
- Supplementary examples cover original, stylized, and larger Stable Diffusion models, as well as ControlNet.Source: A. Appendix · p. 11
- The supplementary material evaluates T-Stitch together with cache-based methods and reports compatibility with DeepCache for further speedup.Source: A. Appendix · p. 11
- The appendix evaluates T-Stitch with ToMe and includes generated-image examples for that combination.Source: A. Appendix · p. 11
A.1. Practical Deployment of T-Stitch
Practical T-Stitch deployment is formulated as allocating denoisers across trajectory intervals to maximize generation quality within a compute budget.
A.1. Practical Deployment of T-Stitch
- The allocation objective selects interval fractions for denoisers to maximize a generation metric under a compute-budget constraint.The denoisers have known sampling costs, ordered from smaller to larger models.Source: §A.1 · p. 12
- Known denoiser costs let the method enumerate fraction combinations and build a lookup table mapping configurations to compute budgets.Source: §A.1 · p. 12
- For a target budget, practitioners can sample configurations from the table that satisfy it and apply them to generation tasks.Source: §A.1 · p. 12
A.2. Frequency Analysis in Denoising Process
A Fourier analysis of DiT-XL latent noises indicates that early denoising steps generate low-frequency components, while later steps refine high-frequency details.
A.2. Frequency Analysis in Denoising Process
- DiT-XL log-amplitude visualizations of Fourier-transformed latent noises track frequency changes across sampling steps.Source: §A.2 · p. 12
- Low-frequency amplitudes rise rapidly from timestep 999 to 555, indicating intensive low-frequency generation early in sampling.Source: §A.2 · p. 12
- At later steps, especially t = 111 and t = 0, high-frequency amplitudes rise significantly, implying detail refinement.Source: §A.2 · p. 12
A.3. Pretrained DiTs and U-Nets
The appendix describes the pretrained DiT and U-Net model families and reports configurations for three-model stitching under DDIM and DDPM sampling.
A.3. Pretrained DiTs and U-Nets
- Tables 3 and 4 compare the pretrained DiT family and a reproduced small U-Net, with model sizes separated to enable a clear speedup.Source: §A.3 · p. 12
- Figure 12 uses DiT-S, DiT-B, and DiT-XL for three-model stitching with DDIM at 100 timesteps and guidance scales 1.5, 2.0, and 3.0.Source: Figure 12 · p. 13
- Figure 13 uses DiT-S, DiT-B, and DiT-XL for three-model stitching with DDPM at 250 timesteps and guidance scale 1.5.Source: Figure 13 · p. 13
A.4. Effect of Different Classifier-free Guidance on Three-model T-Stitch
Three-model T-Stitch maintains consistent performance across classifier-free guidance scales and interpolates a smooth Pareto frontier between DiT-S and DiT-XL.
A.4. Effect of Different Classifier-free Guidance on Three-model T-Stitch
- Three-model T-Stitch performs consistently across different classifier-free guidance scales.Source: §A.4 · p. 13
- Across guidance scales, T-Stitch interpolates a smooth Pareto frontier between DiT-S and DiT-XL.Source: §A.4 · p. 13
- The authors present this consistency as evidence of broad applicability because DPMs commonly use different guidance scales to control image generation.Source: §A.4 · p. 13
A.5. FID-50K vs. FID-5K
Across these experiments, T-Stitch offers flexible speed–quality trade-offs across model families, samplers, and combinations with other acceleration methods. The appendix also tests interval-specific finetuning, memory overhead, and robustness across practical settings.
A.5. FID-50K vs. FID-5K
- FID-5K observations are similar to FID-50K, while replacing SDv1.4 with compressed BK-SDM Small can further reduce time cost at an image-quality trade-off.The FID comparison uses DiT with DDPM, 250 steps, and guidance scale 1.5.Source: §A.5 · p. 13, Figure 14 · p. 14, §A.7 · p. 14
A.6. Compared to Directly Reducing Sampling Steps
- Against directly reducing sampling steps, T-Stitch has competitive speed–quality trade-offs and even better FIDs in the 50–100-step region.The comparison uses DDIM on DiT-XL; classifier-free guidance scale is denoted by “s” in Figure 15.Source: §A.6 · p. 14, Figure 15 · p. 14
A.10. T-Stitch with Different Pretrained Model families
- T-Stitch works across pretrained model families and U-Net-based models, while examples show that early small-model steps can preserve image quality and short-training DiT-S has only a minor effect on overall FID.The U-ViT H experiment uses DiT-S at early steps; the shorter DiT-S training schedule is 400K iterations.Source: §A.10 · p. 15, Table 5 · p. 15, §A.16 · p. 17
A.11. More Examples in Stable Diffusion
- Stable Diffusion examples show flexible speed–quality trade-offs, and stitching can improve prompt alignment and diversity for stylized models while supporting complex prompts, ControlNet, and stable repeated runs.The stylized examples use InkPunk-Diffusion and Ghibli-Diffusion; the robustness example reports eight consecutive runs.Source: §A.11 · pp. 15–16
A.12. Finetuning on Specific Trajectory Schedule
- With DiT-S/B/XL assigned 50%:30%:20% of steps, interval-specific finetuning yields an FID of 16.49 and improves FID, Precision, and Recall over the unfine-tuned schedule.The separately finetuned DiT-B and DiT-XL receive 250K additional ImageNet-1K training iterations; the generated images show improved local details.Source: §A.12 · p. 16, Table 6 · p. 16
A.13. Compared with More Stitching Baselines
- The default small-to-large schedule achieves the best FID and Inception Score among stitching baselines at similar sampling speed.The comparison uses DiT-S/XL with DDIM, 100 steps, and guidance scale 1.5.Source: §A.13 · p. 17, Table 7 · p. 17
A.14. Additional Memory and Storage Overhead of T-Stitch
- Compared with DiT-XL alone, using DiT-S for 50% of steps adds 5% parameters, 4% GPU memory, and 10% local storage while accelerating sampling by 1.76×.At early 40–50% allocations, T-Stitch has a minor effect on Precision and Recall; later allocations show clearer trade-offs consistent with FID.Source: §A.14 · p. 17, Table 8 · p. 17, §A.15 · p. 17
A.18. Compatibility with LCM
- T-Stitch combines with LCM distillation, DeepCache, and Token Merging; reported experiments show further speedups or improvements across image-quality and alignment metrics.With DeepCache, the speedup gain is slightly smaller than for T-Stitch alone because BK-SDM Tiny is only 1.5× faster than SDv1.4 under DeepCache, versus 1.7× without it.Source: §A.18 · p. 18, §A.19 · p. 19, §A.20 · p. 19