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

Diffusion Transformers have recently demonstrated unprecedented generative capabilities for various tasks. The encouraging results, however, come with the cost of slow inference, since each denoising step requires inference on a transformer model with a large scale of parameters. In this study, we make an interesting and somehow surprising observation: the computation of a large proportion of layers in the diffusion transformer, through introducing a caching mechanism, can be readily removed even without updating the model parameters. In the case of U-ViT-H/2, for example, we may remove up to 93.68% of the computation in the cache steps (46.84% for all steps), with less than 0.01 drop in FID. To achieve this, we introduce a novel scheme, named Learning-to-Cache (L2C), that learns to conduct caching in a dynamic manner for diffusion transformers. Specifically, by leveraging the identical structure of layers in transformers and the sequential nature of diffusion, we explore redundant computations between timesteps by treating each layer as the fundamental unit for caching. To address the challenge of the exponential search space in deep models for identifying layers to cache and remove, we propose a novel differentiable optimization objective. An input-invariant yet timestep-variant router is then optimized, which can finally produce a static computation graph. Experimental results show that L2C largely outperforms samplers such as DDIM and DPM-Solver, alongside prior cache-based methods at the same inference speed. Code is available at https://github.com/horseee/learning-to-cache

Paper

Similar papers

Peer review

Reviewer NP2z6/10 · confidence 4/52024-06-14

Summary

This paper proposes a method to accelerate DiT model inference using layer caching strategy. By utilizing feature interpolation, the non-differentiable layer selection problem is transformed into a differentiable optimization problem. The routing matrix $\beta$ is learned to indicate whether the features of a certain layer at the current timestep can be reused from the cached features of the same position in the previous timestep. Extensive experimental results demonstrate the effectiveness of this method in accelerating DiT model inference and also shed light on the redundancy of layers in current DiT models.

Strengths

1. Learning-to-Cache (L2C) transforms the non-differentiable layer selection problem into a differentiable optimization problem through interpolation, which is a clever transformation and forms the basis for optimizing the subsequent routing matrix $\beta$. 2. It is meaningful to explore how to apply feature caching mechanisms to DiT models for inference acceleration. The experimental results of the paper also demonstrate the effectiveness of the method. Compared with simply reducing NFE and previous feature caching methods, L2C achieves better performance. 3. This paper is well-organized and well-written;

Weaknesses

My main concern lies in the scalability of this method: 1. L2C requires training for different DiT models and diffusion schedulers, which limits the potential applications of this method; 2. This paper reports experimental results on DiT and U-ViT series models, but does not experiment on text-to-image models based on the DiT architecture (e.g., Pixart-$\alpha$[1]). Is it because training on large-scale text-image pairs dataset is too costly? 3. The paper does not report specific training costs time; 4. Compared to the original inference process, will L2C increase additional memory due to routing matrix $\beta$ overhead and feature caching? [1] Chen, Junsong, et al. "PixArt-$\alpha $: Fast Training of Diffusion Transformer for Photorealistic Text-to-Image Synthesis." arXiv preprint arXiv:2310.00426 (2023).

Questions

see weakness

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

Limitations are discussed briefly, but could be touched upon in more detail.

Reviewer faWt6/10 · confidence 2/52024-07-09

Summary

This paper introduces L2C, a novel approach that dynamically caches computations in diffusion transformers, significantly reducing the computational load. L2C leverages the repetitive structure of transformer layers and the sequential nature of diffusion, optimizing caching decisions to produce a static computation graph. Experimental results show that L2C outperforms existing methods like DDIM and DPM-Solver, as well as prior cache-based techniques, at the same inference speed.

Strengths

- The writing is easy-to-follow - The motivation is strong. There exists many redundancy and the authors bypass it in a smart way

Weaknesses

I am not an expert in this area. There is no obvious weakness as far as I can tell

Questions

See Weaknesses

Rating

6

Confidence

2

Soundness

3

Presentation

3

Contribution

3

Limitations

See Weaknesses

Reviewer FT6Q6/10 · confidence 5/52024-07-11

Summary

The paper presents Learning-to-Cache (L2C), a method to accelerate diffusion transformers' inference by caching redundant computations across timesteps. A learnable router dynamically determines which layers can reuse calculations from previous timesteps. L2C can eliminate up to 93.68% of computations in specific steps (46.84% overall) in models like U-ViT-H/2 with minimal performance loss.

Strengths

- The proposed caching method for the diffusion process is novel and provides acceptable speed-up without requiring retraining or fine-tuning of the model weights, only adding a few learnable parameters. - The paper introduces a straightforward learnable optimization process to identify redundancies in the diffusion process. - The proposed router is time-dependent but input-invariant, enabling the formation of a static computation graph for inference. - The paper is well-written and easy to understand. The figures and result tables are easy to follow and comprehensive, offering clear visual representations and supporting the text.

Weaknesses

- The paper's contribution is incremental, primarily introducing a learnable router that determines what computations to reuse from previous timesteps. It would benefit from offering more innovations and deeper insights to significantly advance the field. - The improvement in speed-up is very similar to DeepCache, which does not require any training. There are some minor mistakes in the text, such as a typo in line 233 where "maximum" is misspelled as "mamimum."

Questions

- Step distillation techniques, which typically involve few or just one step, have successfully accelerated the diffusion process with minimal impact on output quality. Can the proposed model be integrated with distilled models to achieve additional speed-up on top of the improvements gained from distillation? - You have proposed training only the router parameters while freezing the model parameters. Is there any benefit to fine-tuning and optimizing both the model parameters and the router parameters together? - What are the implications if the diffusion main loss function is used to train the router?

Rating

6

Confidence

5

Soundness

4

Presentation

4

Contribution

3

Limitations

Since the goal is to speed up inference, it should be discussed whether this method can be combined with other acceleration techniques, such as step distillation.

Area Chair 6AMM2024-08-10

Rebuttal Discussion

Dear Reviewers, This is your AC. The authors have provided a response to the comments. Please respond to the rebuttal actively. Best, AC

Authorsrebuttal2024-08-12

Dear Reviewer FT6Q: Thank you for your valuable feedback on our work. Your constructive comments on our work are invaluable, and we genuinely hope to get feedback from you. Regarding the weaknesses you mentioned, we include the corresponding experiments as follows: 1. Experimental results when **applying Learning-to-Cache to 4-step LCM on PixArt**. 2. Generation quality compared with other cache-based methods (DeepCache, Faster Diffusion) 3. Experimental results when optimizing both the model parameters and the router parameters 4. Experimental results using the original diffusion loss. Your feedback is incredibly important to us, and we sincerely thank you for considering our rebuttal. We are more than happy to discuss them if you have any further concerns or questions. Thank you again for your time and effort to review our work and looking forward to your response. Best Regards, 
Authors of submission 1630

Reviewer FT6Q2024-08-12

Increasing the Score

I want to thank the authors for their rebuttal and answering my questions. I will raise my score to WA.

Authorsrebuttal2024-08-13

Thank you for reviewing our paper and providing valuable feedback. We're glad our rebuttal addressed your concerns, and we'll include the suggested experiments in the revised manuscript. Thanks again for your time and effort.

Authorsrebuttal2024-08-12

Dear Reviewer NP2z, We sincerely appreciate the time and effort you have dedicated to reviewing our work. We greatly appreciate your thoughtful review of our work and we are looking forward to hearing your feedback. To address your concerns regarding the scalability of our method, we have conducted additional experiments, including: 1. Experimental results on PixArt, alongside the comparison results with the few-step DPM-Solver. 2. Showing that the training cost for the router remains within an acceptable range. We are grateful for your attention to our rebuttal and are committed to addressing any additional concerns you might have. Thank you once again for your thoughtful consideration. Best regards, Authors of submission 1630

Reviewer NP2z2024-08-12

Official Comment by Reviewer NP2z

I appreciate the authors' rebuttal and extra experiments, most of my concerns have been resolved. I will keep my positive score.

Authorsrebuttal2024-08-13

Thanks so much for the quite encouraging feedback. We will continue to polish our draft, following your advice. Thank you again for your time and effort in reviewing our submission

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC