Paper: T-Stitch: Accelerating Sampling in Pre-Trained Diffusion Models with Trajectory Stitching (ICLR-2024)

Reviewer comment (Experiments):
This paper uses a too simple and strict design of using a weaker model in earlier steps and switching to a stronger model in later steps. However, there is not enough justification or comparison to other baseline approaches. How does it compare to more flexible baselines, e.g. interleaving strong model steps and weaker model steps throughout the whole process, or gradually reducing the probability p of using the weaker model e.g. from p=1 at t=T to and p=0 at t=0?

Author response:
**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.

| **Method**            | **FID↓**   | **Inception Score↑** | **Time Cost (s)** |
| --------------------- | --------- | ------------------- | ----------------- |
| Interleave            | 19.02     | 120.04              | 10.1              |
| Decreasing Prob       | 12.94     | 163.45              | 9.8               |
| Large-to-Small        | 27.61     | 72.60               | 10.0              |
| Small-to-Large (Ours) | **10.06** | **200.81**          | 9.9               |

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

Reviewer comment (Experiments):
Lack of comparison to other related works on multi-expert DPM approaches like [1] and [2]. Under the same inference time budget, how does the proposed approach compare to [1]? Does using a larger pre-trained model in the earlier steps and using a smaller model in later steps have better or worse generation quality compared to [1] which adopts differently designed architectures tailored toward the low-frequency features for the earlier steps or the high-frequency information for later denoising steps?

Author response:
**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.

Reviewer comment (Experiments):
In Figure 15, even the simple baseline of directly reducing the sampling steps outperforms the proposed method at the 10-50 steps range for s=1.5 and 10-20 steps range at s=2.0. This means the T-stitch approach could achieve better generation quality and latency tradeoff if combined with reducing steps and this was not investigated. More importantly, this result shows the proposed t-stitch approach does not have a strong performance even compared to this simple baseline and more comparisons to other approaches in the literature like [1] and [2] are needed.

Author response:
**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.

Reviewer comment (Reproducibility):
Two models mean more storage consume, or if they are sent into the GPU in order, they will be in and out for every batch, which is not convenient.

Author response:
**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).

Reviewer comment (Experiments):
The authors are recommended to compare the speed of their proposed method with the other accelerating methods mentioned in the second paragraph of Introduction.

Author response:
**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.

Reviewer comment (Experiments):
Memory Overhead: Adopting additional smaller models during sampling increases memory usage, which could be a concern for very large models.

Author response:
**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.

Reviewer comment (Experiments):
Finicky Tuning: Getting the right model stitching fractions to optimize the speed-quality tradeoffs may require finicky tuning based on the models and datasets. More principled guidelines could help.

Author response:
**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.

Reviewer comment (Reproducibility):
For finetuning experiments, can you elaborate on the exact finetuning procedure? Was it only on stitched intervals? How do fully finetuned models compare?

Author response:
**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.

| **Method**                     | **FID ↓** | **Inception Score ↑** |
| ------------------------------ | --------- | --------------------- |
| Pretrained                     | 16.49     | 123.11                |
| Finetuned at all timesteps     | 16.04     | 125.81                |
| Finetuned at stitched interval | 13.35     | 155.35                |

Reviewer comment (Experiments):
The prompts used for stable diffusion examples are quite simple. Have you tried more complex prompts and datasets? How robust is the method?

Author response:
**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.