T-Stitch: Accelerating Sampling in Pre-Trained Diffusion Models with Trajectory Stitching

Thanks for your valuable feedback. We would like to address your concerns as below.

Q1 - Concerns on the ideal cut-off for switching models.

Please refer to our general response.

Q2 - Limitations when T-Stitch beyond pair-wise.

Applying T-Stitch beyond pairwise naturally increases the number of potential allocation strategies, which we also agree with the reviewer that it introduces additional challenges. To address this, our initial manuscript provides a practical guideline (Section A.2) for such scenarios by framing T-Stitch as a compute allocation problem, which aims to find a few configuration sets from a lookup table to satisfy a computational budget and then apply to generation tasks.

Besides, with more models becoming available, as mentioned by the reviewer, our experiments in Table 3 indicate that simply adopting a small model to speed up the largest model performs favorably compared to the searched combination of 4 denoisers by DDSM. Notably, as mentioned in Lines 446–447, we only adopt the smallest network in their model family to accelerate the largest network. This suggests that we may not need too many denoisers in the sampling steps to achieve speedup in practice. As our current study has demonstrated broad effectiveness in many practical scenarios, we assume there is a potential in future work to extend this idea more intelligently when it is beyond pair-wise.

Thanks for your constructive comments. We would like to address your additional concerns as follows.

Q1 - Concerns on the reduced time cost from T-Stitch, and comparison with progressive distillation.

Essentially, T-Stitch provides a free lunch for accelerating the large diffusion model sampling by directly adopting a small model. This speedup not only demonstrates favorable efficiency gain (“meaningful computational savings at little to no cost to the quality of generated images”, as recognized by Reviewer 4dLD), but is also broadly applicable ("Compatible with existing acceleration techniques", as recognized by Reviewer Aq9T).

Relation with step-distillation based method. In fact, T-Stitch allocates different compute budgets at different steps, which is a complementary technique with reducing sampling steps, not competing with it. In Figure 33, we have shown that T-Stitch works with step-distilled models as a complementary acceleration technique. Furthermore, we would like to note that T-Stitch is training-free, thus may not be directly comparable to training-based approaches. However, in Figure 15, we additionally provide an ablation study of the direct comparison for comprehensive reference.

Q2 - Shared latent space may limit the applicability of T-Stitch.

We have never assumed that all different models have shared latent space, as it is impractical. Also, not all models in the public model zoo can achieve the same effectiveness in T-Stitch.

In fact, we design T-Stitch by starting from the insight that different diffusion models can share similar sampling trajectories if trained on the same data distribution (Lines 95-98), thus pointing out that we can directly allocate an efficient small model at the early steps for accelerating large model sampling. And we are indeed able to build upon those small models to validate our idea.

Besides, we have demonstrated that T-Stitch is generally applicable to many scenarios, as shown in our experiments and highly recognized by Reviewer 4dLD ("various backbone architectures , with various samplers") and Reviewer Aq9T ("Broader Impact & Applications"). Therefore, this similar latent space does not hinder the broad applicability of T-Stitch in practice.

Q3 - Further explanations on Figure 3, and compared to L2 distance.

Initially, we have explained this phenomenon in Lines 94-101 with two insights, where we assume that 1) different diffusion models trained on the similar dataset would capture the same data distribution. 2) Moreover, recent works have shown that compared to the large diffusion model, small diffusion models are able to capture relatively good low frequencies at the early steps, indicating the power of small diffusion models to generate good global image structures given the same condition.

We further evidenced the second insight in Figure 17 and show that it actually happens: applying the DiT-S at the early steps minorly affects the global structures (tiger, table, dog, etc) while it gradually loses more details with a more increased fraction of DiT-S timesteps. This experiment suggests that we can exploit the advantage of small models in capturing good low-frequencies to achieve speedup.

"Does the L2-distance exhibit a similar pattern as shown in Fig. 3?"

Yes. In Section A.21 of our revised submission, we show that the L2 distances between different DiT models at the early denoising steps are much smaller than those of later steps, similar to the observation in Figure 3.

Q4 - A shared latent space is a strict condition, T-Stitch for models trained with different configurations such as noise schedules and different sampling steps.

It is quite challenging for us to find pretrained diffusion models with different noise schedules or diffusion steps in one code repository, as many of them are trained by different authors and contain very different configurations. Intuitively, models trained with different configurations (e.g. linear noise schedule and cosine noise schedule) differ more significantly at the intermediate timesteps, thus applying T-Stitch in this case possibly hurts the generation quality.

Besides, we would like to mention that our main goal is to accelerate the sampling process of a large pretrained diffusion model. In this case, our contribution is to show that a small model with a shared latent space as the large model provides a free lunch for accelerating diffusion model sampling in many practical scenarios. At this stage, our work has demonstrated broad effectiveness and has been recognized by most reviewers. Thus, we will leave those explorations for future works.

Thanks for your very positive reviews and constructive comments, we would like to address your additional questions below.

Q1 - Paper polishing.

Thanks for your great advice, we will polish our manuscript based on those comments.

Q2 - Selecting the switching point between the small and large model.

Please refer to our general response.

Q3 - Further discussion with early-exit works.

In general, we aim to explore the compute budget allocation for diffusion model sampling, which is orthogonal with individual model acceleration techniques such as early-exiting or model compression that specifically focus on one model, as discussed in Section A.8.

Compared to recent early-exit works, such as DeeDiff [A], AdaDiff [B], and DuoDiff [C], T-Stitch does not require training since we directly drop the small model at the early denoising steps. Furthermore, compared to Adaptive Score Estimation (ASE) [D] which heuristically designs block-exiting strategies based on different architectures then finetunes the target diffusion model with substantial training cost, we found our speed-quality trade-offs are clearly better under the equal experimental setting, as shown below.

The above results are adopted from Table 3 of [D]. FID-5K is evaluated based on ImageNet-256 and DDIM sampler. “Acceleration” refers to the acceleration in sampling speed. “n” in “Dn-DiT” represents the acceleration scale. Details for different settings can be found in Table 2 of ASE [D].

The above results are from our Figure 1, which is based on the same experimental setting: ImageNet-256, DDIM sampler, and FID-5K. Note that due to different implementations, our DiT-XL baseline performance can be slightly different. We have included this comparison in Section A.22 of the revised submission.

[A] Tang, Shengkun, et al. "Deediff: Dynamic uncertainty-aware early exiting for accelerating diffusion model generation." arXiv preprint arXiv:2309.17074 (2023).

[B] Zhang, Hui, et al. "Adadiff: Adaptive step selection for fast diffusion." arXiv preprint arXiv:2311.14768 (2023).

[C] Fernández, Daniel Gallo, et al. "DuoDiff: Accelerating Diffusion Models with a Dual-Backbone Approach." arXiv preprint arXiv:2410.09633 (2024).

[D] Moon, Taehong, et al. "Early Exiting for Accelerated Inference in Diffusion Models." ICML 2023 Workshop on Structured Probabilistic Inference {&} Generative Modeling. 2023.

Thanks for your very positive and comprehensive reviews. Below, we would like to address your additional questions.

Q1 - Overall concerns.

Thanks for providing those valuable comments. We would like to briefly address them below.

Limitations Analysis. We have added a brief discussion on limitations in Section A. 23 of our revised submission.

Theoretical Gaps. Please refer to our response to Reviewer 1Rg8 Q1.

Architectural Considerations. In general, T-Stitch works well as long as both small and large models share similar latent space and spatial dimension of latents, as shown in Table 8. We are also orthogonal to individual model optimization, as discussed in Section A. 8. As our current study has demonstrated broad effectiveness, we would like to leave other interesting explorations (multi-model scenarios, feature space alignment) for future works.

Practical Implementation Challenges. Overall, the implementation of T-Stitch is simple: let the small model do denoise sampling first then switch into the large model for the subsequent timesteps. With minor memory overhead (Table 4), T-Stitch is general and is “more likely to be adopted” (Reviewer 4dLD) in many practical scenarios, such as ControlNet (Figure 30) and text-to-video generation (Figure 34).

Q2 - T-Stitch for diffusion models of very different architectures.

The fundamental insight of T-Stitch is that different diffusion models trained on similar datasets can share a similar latent space. As demonstrated in Table 8 of the Appendix, our experiments show that T-Stitch performs very well when applied to very different model families, such as U-ViT and DiT.

"Are there specific architectural compatibility requirements?"

Yes. T-Stitch requires the latent noise from both models to have the same spatial dimensions to allow seamless switching during denoising sampling. This design is inspired by the observation that widely adopted models (e.g., DiTs, SD, and fine-tuned SDs) typically share the same latent shape (Lines 208–211). We leave the development of more challenging stitching strategies for future work.

Q3 - More analysis of the improved prompt alignment?

We speculate that stylized SD models, such as those trained with DreamBooth, are more prone to overfitting and catastrophic forgetting [A, B] due to being trained on very few images. On the other hand, the initial SD model, trained on large-scale text-to-image datasets, may help complement the forgotten knowledge. By adopting the small SD model during the early steps, it can provide general priors at the beginning [C], thus compensating for the missing concepts in the prompts for the overfitted stylized SD models.

In our experiments, we found this approach generally applicable to both standard SD models and fine-tuned/stylized SD models (Figures 26, 27, and 28), as well as to other diffusion model acceleration techniques such as DeepCache (Figure 21) and token merging (Figure 22).

[A] Zhang, Lvmin, Anyi Rao, and Maneesh Agrawala. "Adding conditional control to text-to-image diffusion models." ICCV. 2023.

[B] Ruiz, Nataniel, et al. "Dreambooth: Fine tuning text-to-image diffusion models for subject-driven generation." CVPR. 2023.

[C] Graikos, Alexandros, et al. "Diffusion models as plug-and-play priors." NeurIPS (2022).

Q4 - Primary failure modes of T-Stitch.

Note that our primary goal in this work is to accelerate diffusion model sampling while preserving its generation quality as much as possible. Considering the fact that T-Stitch adopts a small model to accelerate the target large model, the primary failure mode lies in the weakness of the chosen small model.

For example, if the small model is not significantly faster than the larger model while its generation quality is substantially worse, T-Stitch inevitably results in a poor speed-quality trade-off. We expect future works to train or distill an efficient small diffusion model to specialize in the early sampling steps for better speed-quality trade-offs. Besides, T-Stitch will also underperform if the chosen small model has a very different latent space compared to the large model. For these reasons, we have provided a principle of model selection in our initial submission (Lines 227-232).