Learning-to-Cache: Accelerating Diffusion Transformer via Layer Caching

We sincerely appreciate the valuable feedback and constructive suggestions. Thanks so much for taking time and effort to review our paper.

W1: L2C requires training for different DiT models and diffusion schedulers, which limits the potential applications of this method

Thank you for your insightful question. The training cost of L2C is low because only the router matrix is updated, leaving the model parameters unchanged. For example, in PixArt-XL-2 with a 512 resolution, only 840 parameters need to be trained, taking just 3 hours on 8 A5000 GPUs. We consider this training overhead to be relatively small; using a distillation approach to compress a small model typically requires significantly more resources, such as 4 A100 GPU days for BK-SDM [1].

[1] BK-SDM: A Lightweight, Fast, and Cheap Version of Stable Diffusion. ECCV24.

---

W2: No experiment on text-to-image models based on the DiT architecture (e.g., Pixart). Is it because training on large-scale text-image pairs datasets is too costly?

Thanks for the valuable suggestion. Here we show the experimental results on PixArt-XL-2-512x512.

• Training Cost: The training cost for this experiment is approximately 3 hours on 8 A5000 GPUs, using around 200,000 images for training. We utilized the first 200,000 samples from the SAM-LLaVA-Captions10M dataset. The router matrix converges and achieves optimal performance with 200,000 samples. We also tested with 400,000 samples but observed no performance gain, as the router does not change with additional training samples.

• Generation Quality: We test our method on the validation set of COCO2014 (30k) and COCO2017 (5k). The results, shown in the table below, indicate that our method outperforms the approach using fewer steps in DPM-Solver. We provide some qualitative examples in Figure 5 of the attached PDF.

• What the router learns: We have observed some intriguing patterns in the router, as illustrated in Figure 1 of the attached PDF. The cross-attention block displays significantly more temporal redundancy compared to other types of modules, such as self-attention and MLP. In addition, this router has the unique feature of not being cacheable at intermediate steps. We believe these specific patterns can also guide the future design of model architecture, helping to eliminate unnecessary computations.

---

W3: The paper does not report specific training costs time;

We apologize for missing this important experimental detail. Here, we report the training cost for each experiment. We use ImageNet for training, and in all cases, convergence is achieved in less than one epoch. The experiments are conducted on 8 A5000 GPUs. We would add this training cost to our next version of the manuscript. Thanks so much for pointing this out.

---

W4: Compared to the original inference process, will L2C increase additional memory due to routing matrix overhead and feature caching?

• For Routing Matrix:

The routing matrix has a very small number of parameters, calculated as $steps \times layers \times block\\_per\\_layer$. In our experiments, the largest router, used for DiT-XL/2 with 50 sampling steps, contains 1,400 parameters ($25 \times 28 \times 2$), resulting in an extra memory overhead of just 2.8KB (using FP16 for inference).

• For Feature Caching:

Yes, L2C needs extra overhead for feature caching. As the cache in the computer system and KV-Cache in LLM, L2C trades space for time by storing intermediate results in VRAM, leading to additional memory overhead. Below is the additional overhead observed in the DiT model. We believe there is still room for optimization of the memory overhead in feature caching. Thanks for raising this critical issue.

Thank you very much for your review of our manuscript. We appreciate your time and effort in evaluating our work.

It is encouraging to hear that you like our work and that no obvious weaknesses have been found. If you have any questions where you think further detail or explanation might be beneficial, we would be happy to address them.

Thank you again for your valuable time and feedback.

We extend our gratitude for your insightful feedback and suggestions

W1: The paper's contribution is incremental. It would benefit from offering more innovations and deeper insights

We greatly value your suggestion that we need to offer more profound insights in this paper. Beyond introducing a new method, we aim to share the following key insights to advance the field:

(1) Theoretical explanation of the caching-based method: In Section 3.2&3.3 and Appendix, we reveal that fast samplers and cache mechanisms represent different levels of reuse. Unlike previous heuristic-based approaches like DeepCache, our work establishes a theoretical connection and build the clear relationship between cache and fast solvers. This not only clarifies their relationship but also enhances the reliability of caching-based methods by providing theoretical support.

(2) Insights from the special patterns in the learned routers: L2C reveals intriguing patterns in the learned routers (refer to Figures 1, 2, 3 and 4 in the attached PDF). In U-ViT (Figure 2), the intermediate layers can be entirely cached, whereas the early and later layers play a crucial role. We also add one experiment on PixArt-XL-2 and for PixArt, the temporal redundancy in the cross-attention layer (Figure 1.b) is notably pronounced, with over 94% of layers (265 out of 280) being eligible for caching.These findings can be leveraged not only to accelerate model inference but also to guide the design of model architectures.

To summarize, we aim to move beyond heuristic methods in designing the cache mechanism, incorporating performance guarantees and providing more theoretical interpretability for this special cache mechanism in accelerating the diffusion models. This has not been explored before and we think is important for this area.

---

W2: The improvement in speed-up is very similar to DeepCache.

Thanks for your insightful quesion. We compare our approach with DeepCache, focusing primarily on generation quality under the same acceleration ratio. Our method improves the FID from 2.70 to 2.64, using the same model and parameters, indicating a superior caching mechanism compared to DeepCache. Additionally, DeepCache is limited to U-ViT due to it is bound to the u-shaped structure, making it inapplicable to models like DiT and PixArt. However, L2C is versatile and can be applied to all these models.

---

Q1&Limitation: Can the proposed method be combined with step distillation?

Thank you for your constructive suggestion. We built our algorithm on the distilled PixArt-XL-2-512x512, employing a 4-step LCM scheduler. The table below presents our results, which are compared against the 3-step LCM. The learned router is visualized in Figure 4 of the attached PDF. From our experiments, our method achieves approximately 1.28x acceleration, successfully caching 78 out of 168 blocks. Consistent with the results on PixArt with 20 NFEs, cross-attention remains the most redundant component in the denosing sequence.

---

Q2: Optimizing both the model parameters and the router parameters

Thanks for your great suggestion. We show the comparison results if optimizing the model parameters and the routers together:

The conclusion is that, yes, optimizing both the model and the router enhances the model's performance, as evidenced by slight improvements in FID and sFID. However, this optimization process requires additional time to train and fine-tune both components. We will include this experiment in our revised manuscript and thanks again for your inspiring suggestion.

---

Q3: what if the diffusion main loss function is used to train the router?

Thank you for your insightful question. We have adjusted the training loss accordingly, and the results are presented in the table below. We selected baselines with a very similar acceleration ratio for comparison. The results indicate that using the original diffusion loss makes it more challenging for the router to learn, resulting in slightly worse image quality compared to the distillation loss used in our submission.