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. Design is too simple, compared to other baselines.

We agree T-Stitch is a simple technique. However, it is also a generalizable and effective approach as highly recognized by other reviewers, i.e., “the insight …is clever….a simple yet novel idea…broadly applicable…highly practical” (Reviewer 9mJB ), “built upon the dynamics of diffusion models…clearly distinguishing it from model-wise stitching..” (Reviewer tbyv).

Furthermore, in our initial experiments, we have explored various allocation strategies, such as interleaving and large-to-small, but found that our default design can achieve the best speed-quality trade-off. In the following, we demonstrate the detailed comparison with the mentioned baselines:

1. Interleaving. During denoising sampling, we interleave the small and large models along the trajectory, ie., DiT-S -> DiT-XL -> … -> DiT-S -> DiT-XL. Eventually, DiT-S takes 50% steps and DiT-XL takes another 50% steps.
2. Decreasing Prob. Linearly decreasing the probability of using DiT-S from 1 to 0 during the denoising sampling steps.
3. Large-to-Small. Adopting the large model at the early 50% steps and the small model at the last 50% steps.
4. Small-to-Large (Ours). The default strategy of T-Stitch by adopting DiT-S at the early 50% steps and using DiT-XL at the last 50% steps.

All experiments are based on DiT-S and DiT-XL, with DDIM 100 steps and a guidance scale of 1.5. FID is calculated based on 5K images. As shown in the table below, under the similar time cost, our design achieves a superior advantage in FID and Inception Score compared to the other baselines. Nevertheless, we agree they are meaningful baselines, and we have included these comparisons in Section A.13 of the revised Appendix.

The time cost is measured by generating 8 images on one RTX 3090 in seconds (s).

Q2. Compared with MEME [1] and eDiff-I [2].

It is important to note that we are significantly different from MEME and eDiff-I:

1. T-Stitch accelerates off-the-shelf pretrained large diffusion models in a training-free manner, which is a general technique and enables “broad applicability”, as recognized by Reviewer 9mJB.
2. Existing multi-experts DPMs design new models and train them from scratch using expensive computing resources to either reduce the time cost per step with multiple smaller DPMs (MEME) or achieve better performance with multiple large DPMs without considering efficiency (eDiff-I).

Essentially, as already discussed in the Remark of Section 3, we target the research problem of computing budget allocation across different steps while benefiting from training-free, while being directly applicable to SDs, Stylized SDs, SDXL (Figure 25) and ControlNet (Figure 26). Therefore, we are not directly comparable. Furthermore, at this time, MEME is an Arxiv preprint paper that does not release code/weights, thus might be inappropriate for us to compare the speed under the same hardware.

Q3. Comparing T-Stitch with reducing sampling steps, combining T-Stitch with fewer steps was not investigated.

In fact, T-Stitch allocates different compute budgets at different steps, which is a complementary technique with reducing sampling steps, not competing with it. Figure 15 additionally provides an ablation study of the direct comparison for comprehensive reference.

More importantly, we think we have already provided the experiments of combining T-Stitch with fewer steps in Figure 8, Figure 9 and Figure 17. For example, under 10/20/50 steps, T-Stitch can maintain low FID (Figure 8 left) and high image quality (Figure 17) by replacing the early 40% steps with a small DiT-S to accelerate DiT-XL sampling, which clearly demonstrated that T-Stitch can achieve better speed-quality trade-offs if combined with reducing steps.

Dear Reviewer UhCD,

We sincerely thank you again for your great efforts in reviewing our paper. We have provided responses to address your major concerns about the comparison with other trajectory stitching baselines, more discussion with related multi-expert DPMs [1] and [2], as well as the compatibility with reducing the sampling steps. Please don’t hesitate to let us know if you have any further questions.

Best,

Authors of Submission 901

Dear Reviewer UhCD,

Thank you for raising your additional concern during the discussion. We have provided follow-up response to AC which briefly summarizes the relation between T-Stitch and existing acceleration methods. Please feel free to ask any further questions and we sincerely appreciate the opportunity to discuss with you.

Best regards,

Authors of Submission 901

Thanks for your efforts for reviewing our paper. We would like to address your concerns as below.

Q1. Two models mean more storage consumption, and it is not convenient on GPUs.

It is worth noting that compared to the large model, a small model only slightly increases the memory (x1.04) and local disk storage (x1.1), as shown in our general response Part-2.

Besides, T-Stitch is actually very efficient on GPUs since both models are pre-loaded into GPU memory before inference. During the denoising sampling, model switching is equivalent to choosing a different computational graph immediately, thus resulting in a minor overhead.

For example, as shown in our response to Reviewer UhCD Q1, based on DiT-S/XL, even interleaving DiT-S and DiT-XL during the sampling trajectory takes the similar time cost (10.1s) as our default strategy (9.9s) of adopting DiT-S for the early 50% steps then switch into DiT-XL for the last 50% steps, which is still clearly faster than only using DiT-XL for sampling (16.5s).

Q2. Compared to other acceleration techniques

In fact, we are not competing with single models/samplers/reducing steps. Instead, T-Stitch efficiently adopts pretrained small DPMs as cheap drop-in-replacements for large DPMs during sampling. Thus, we are complementary to other sampling acceleration techniques, which has been extensively discussed in the related works, as also recognized by both Reviewer 9mJB (“...complements existing techniques.”) and Reviewer tbyv (“...complementary to advanced diffusion samplers based on better ODE discretization.”). More importantly, we have already demonstrated this compatibility with comprehensive experiments on different number of sampling steps, samplers and architectures in Figure 8, Figure 9 and Table 1, respectively.

Dear Reviewer omxW,

We sincerely thank you again for your considerable efforts in reviewing our paper. We have provided responses to address your major concerns about the storage overhead, the efficiency on GPUs, as well as the compatibility with other acceleration techniques. Please feel free to reach out if you have any further questions.

Best,

Authors of Submission 901

Thanks for your very positive comments! In addition to the initial discussion in Section 4.3, we would like to provide further analysis as below,

1. Dreambooth finetuning is easy to overfit. Stylized SD models are usually finetuned on a very limited number of images, which is easy to result in overfitting and catastrophic forgetting [A, B]. For example, the Ghibili/Inkpunk Diffusion models in Figure 2 are finetuned with DreamBooth. However, according to the official documentation on Huggingface, “DreamBooth finetuning is very sensitive to hyperparameters and easy to overfit ” [C], which may inevitably result in a loss in terms of the pretrained general knowledge [D].
2. General SD may complement the knowledge. The small SD (BK-SDM Tiny) is a pruned and distilled version of the original SD v1.4, which could preserve the majority of the generality. Therefore, by adopting the small SD at the early steps, it may provide some general priors at the beginning [E], thus complementing the missing concepts in the prompts for the overfitted stylized SD models.

We agree that more in-depth analysis can be helpful to understand the prompt alignment for stylized SD models, but as our main aim is to achieve acceleration, we will leave such exploration for future work.

[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] https://huggingface.co/docs/diffusers/training/dreambooth#finetuning

[D] Lin, Yong, et al. "Speciality vs generality: An empirical study on catastrophic forgetting in fine-tuning foundation models." arXiv (2023).

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

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

Q1. Memory Overhead of T-Stitch.

We agree T-Stitch slightly increases memory usage by adopting a smaller model, as already stated in the limitations. However, in practice, this overhead can be very minor since the main bottleneck still lies in the large models. Please refer to our general response Part-2 for more details.

Q2. Finicky tuning of getting the right T-Stitch fractions to optimize the speed-quality trade-offs.

In practice, T-Stitch naturally interpolates a smooth speed-quality curve between a small and large DPM under different fractions of the small DPM. As this curve generally exists under different samplers (Figure 9), numbers of steps (Figure 8), and architectures (Tables 1,2,5), the speed under a given fraction can be roughly estimated based on the almost linearly interpolated time cost curve between the two DPMs (Figure 8 Right), while the more accurate quality measurement can be obtained by querying a pre-computed look-up table, as discussed in Section A.1 of the initial submission. We also agree with the reviewer that more principled guidelines would be helpful and will leave it for future work.

Q3. Diversity measurement of T-Stitch.

Following common practice [A], we adopt Precision to measure fidelity, and Recall to measure diversity or distribution coverage. In the table below, we report the results based on DiT-S/XL, 100 DDIM steps, and a guidance scale of 1.5. As it shows, T-Stitch maintains high Precision and Recall at the early 40-50% steps, which is consistent with FID evaluations. We agree diversity measurement is also important thus we have included this result in Section A.15 of the revised Appendix.

[A] Dhariwal, Prafulla, and Alexander Nichol. "Diffusion models beat gans on image synthesis." NeurIPS (2021): 8780-8794.

Q4. Theoretical analysis and justification on why stitching trajectories preserves sample quality.

Thanks for pointing this out. We agree that more in-depth theoretical proof is missing at this stage. In addition to our initial justifications in the first paragraph of Section 3.2, we would like to provide further analysis as follows,

Theoretically, under the same dataset, score-based models trained with the same denoising score matching loss aim to learn a consistent score function [B], which guides the probability flow at the reverse time to produce the overall ODE trajectory. If they learn the same ground-truth score, they will have the same reverse trajectory. This is evidenced by our Figure 3 where different DiTs learn similar intermediate latent embeddings, thus indicating a similar trajectory. Essentially, T-Stitch forwards the latent codes from the last timestep of a small model into the current timestep for the large model, which provides roughly similar scores, thus the stitched trajectory may interpolate a similar ODE trajectory and produce similar images, as shown in Figure 17.

Due to the current scope, we will leave more comprehensive analysis for future work.

[B] Song, Yang, et al. "Score-based generative modeling through stochastic differential equations." ICLR (2021).

Q5: Finetuning experiments

In Section A.12, we separately finetune the pretrained DiT-B and DiT-XL at their stitched intervals. Concretely,

1. During finetuning DiT-B, at each training iteration, we limit the timesteps to be only sampled from the intermediate 30% steps.
2. Similarly, during finetuning DiT-XL, we only sample the timesteps from the last 20% steps.
3. After both models are finetuned with 250K iterations, we directly apply T-Stitch based on the denoising intervals of DiT-S/B/XL with 50% : 30% : 20%.

To compare with finetuning at all intervals, we finetuned the pretrained DiT-B/XL for all timesteps with the same additional 250K iterations. The table below evaluates the same T-Stitch allocation (50% : 30% : 20%) and demonstrates that this strategy does not yield superior performance, suggesting that finetuning at stitched intervals is more effective.

Q6: More complex prompts and datasets for stable diffusion examples.

Thanks for the suggestion! We have included more examples by using more complex prompts in the Appendix Figure 24 and Figure 25, where T-Stitch performs favorably with long and complex prompts, and can be easily applied into SDXL and ControlNet (Figure 26). Moreover, we write a for-loop script for generating image samples with 8 consecutive runs under the same prompts but different latent noises. As shown in Figure 27, the output images by adopting different fractions of the small SD show great image quality for all runs. Thus it indicates T-Stitch is robust in practice.

Dear Reviewer 9mJB,

We sincerely thank you again for your great efforts in reviewing our paper. We have provided responses to address your major concerns about the memory overhead, tuning guidelines and additional analysis on T-Stitch. Beside that, we have also enriched the section for the finetuning experiment and shown more stable diffusion examples under complex prompts. Please don’t hesitate to let us know if you have any further questions.

Best,

Authors of Submission 901

Dear AC,

Thanks for your valuable time to join the discussion and read our rebuttal. We appreciate the opportunity to address your concerns as follows:

1. Why introduce model stitching as a notable baseline?

Our insight is that “using the same model throughout is a suboptimal strategy for efficiency”, thus we aim to explore more flexible trade-offs between exclusively adopting a large DPM or a small DPM for sampling. To this end, model stitching [A, B] is an effective framework that obtains numerous architectures with different speed and quality trade-offs given a pretrained model family, which provides an initial appropriate baseline that aligns with our objective. However, most existing acceleration techniques focus on individual model optimization by distillation, pruning or redesign which still suffer from fixed model complexity, making it difficult for us to study the problem of compute budget allocation during the sampling process. Model stitching, therefore, is instrumental in showcasing the superiority of T-Stitch in this case.

2. Why has T-stitch not been directly compared to other established learning/finetuning-based acceleration methods?

T-Stitch is a training-free technique for obtaining flexible trade-offs between a pretrained small and a pretrained large DPM, which might be inappropriate to directly compare with other training-based methods. Essentially, our acceleration is driven by the small DPM, where the faster and better it is, the better speed-quality trade-off we can achieve, thus it can be orthogonal to the choice of the large DPMs (eg, DiT-S/XL in Figure 5 and DiT-S/U-ViT-H in Table 5), the number of sampling steps (Figure 8), and samplers (Figure 9).

On the other hand, T-Stitch can further speed up an already accelerated DPM via established training-based methods, such as distillation. For example, as shown in Figure 28 and Figure 29 (new results), given a distilled SDXL from LCM [C] (a concurrent work), T-Stitch can achieve further speedup under 2~4 steps with high image quality by adopting a relatively smaller SD. For example, under 4 steps, applying T-Stitch (75%) reduces the sampling time of LCM distilled SDXL by 15% without sacrificing visual quality. These results demonstrate that T-stitch is generally applicable across various scenarios. Hybrid modeling such as Diffusion-GAN [D] aims to improve GAN and also needs training, which is beyond our scope of diffusion model acceleration at this stage.

We thank the AC again for the efforts in reviewing our paper, and sincerely welcome further discussions.

Best regards,

Authors of Submission 901

[A] X, Yang, et al. "Deep model reassembly." NeurIPS 2022.

[B] Z, Pan, et al. "Stitchable Neural Networks." CVPR 2023.

[C] S, Luo, et al. "Latent consistency models: Synthesizing high-resolution images with few-step inference." arXiv preprint arXiv:2310.04378, 2023.

[D] Z, Wang, et al. "Diffusion-gan: Training gans with diffusion." ICLR 2023.

Dear AC,

Thanks for your further comments. Beside that aforementioned clarifications, we would like to provide a brief summary on the relation between T-Stitch and existing DPM acceleration techniques and hopefully provide a clear big picture.

Previous fast sampling methods for diffusion models, including training-free and training-based ones, are mainly focusing on reducing sampling steps while our approach mainly targets for reducing the overall sampling time under the same sampling steps. This major difference makes T-Stitch completely complementary to previous methods. More discussions are available in the introduction and related work sections. Because our method does not need training at all, it serves as a simple drop-in accelerator that is compatible with existing training-free and training-based methods. And we have shown the compatibility with acceleration techniques as follows:

• Directly reducing sampling steps (Figure 8, Figure 17)
• Advanced samplers (Figure 9)
• Model compression (Figure 16)
• Distillation for sampling with fewer steps (Figure 28, Figure 29).

We have also demonstrated that T-Stitch is broadly applicable to various scenarios,

• Transformer (Figure 5) and U-Net (Table 1).
• Different pretrained model families, i.e., DiT + U-ViT (Table 5)
• Adopting more models for more flexible trade-offs (Figure 6)
• Pretrained SDs (Table 2 and Figure 21)
• Stylized SDs and potential better prompt alignment (Figure 22, Figure 23)
• Complex Prompts and SDXL (Figure 24, Figure 25)
• ControlNet (Figure 26)

Overall, with comprehensive experiments ("strong experimental results and ablation analysis" from Reviewer 9mJB and "good amount of experiments" from Reviewer UhCD), we have shown T-Stitch offers all its benefits without the need for training, serving as a general tool that can be readily applied to existing DPM acceleration techniques for achieving further speedup, while maintaining comparable performance.

We sincerely thank the AC and reviewers again for the valuable time in engaging the discussion.

Best regards,

Authors of Submission 901