HarmoniCa: Harmonizing Training and Inference for Better Feature Cache in Diffusion Transformer Acceleration

Thanks for the reviewer’s constructive comments.

W1: The HarmoniCa framework’s reliance on multiple components (SDT, IEPO, and the custom cache router) could limit its accessibility and reproducibility for practitioners, especially those with limited computational resources. While the model is described in detail, implementing and integrating SDT and IEPO into an existing Diffusion Transformer may pose a significant challenge due to the additional computational and architectural overhead.

Compared to the existing learning-based approach (i.e., Learning-to-Cache), HarmoniCa requires about $\mathbf{\frac{3}{4}}$ of the training time without the need for any training images, making it a more user-friendly option. In terms of architecture, HarmoniCa does not require any modifications to the model itself. Furthermore, we only train a router which is a tensor with a shape of $T \times N$ (a few KB of memory), without any model weights. The minimal computation resource requirement for training with the gradient accumulation is under 12 GB of GPU memory for DiT-XL/2 256$\times$256, making it suitable even for an RTX 4070 Ti GPU. Moreover, users can employ gradient checkpointing to further reduce GPU memory consumption.

W2: While HarmoniCa is tested across multiple models and resolutions, it is evaluated only on ImageNet (for class-conditioned tasks) and MS-COCO (for text-to-image tasks). These datasets are popular benchmarks but may not represent the variety of distributions found in real-world data, especially in domains like video generation or multi-modal applications where diffusion models are increasingly applied.

In addition to the evaluations on ImageNet and MS-COCO, we conducted further tests using the high-quality MJHQ-30K [1] and sDCI [2] datasets with PixArt-$\alpha$ models. We added several metrics, including Image Reward [3], LPIPS [4], and PSNR. The results, summarized in the following table, demonstrate that HarmoniCa consistently outperforms DPM-Solver across all metrics on both the MJHQ and sDCI datasets. For instance, at the 512$\times$512 resolution, HarmoniCa achieves an FID of 7.01 on the MJHQ dataset, which is lower than the 7.04 FID of DPM-Solver with 20 steps, indicating better image quality. Additionally, under the same configuration, HarmoniCa achieves a PSNR of 22.09, compared to DPM-Solver's 21.41 with 15 steps, reflecting better numerical similarity. Moving forward, we will explore the application of our method to video generation and multi-modal diffusion models in our future research. We have also included the results in Sec. Q.

W3: The Image Error Proxy-Guided Objective (IEPO) is a central contribution but relies on approximating final image quality based on intermediate denoising steps. The validity of this proxy could vary significantly depending on the task and model architecture, yet the paper provides limited empirical analysis of IEPO’s sensitivity to different proxy metrics.

Our superior performance and efficiency across class-conditional and text-to-image (T2I) tasks, evaluated on 8 models and 4 samplers with resolutions ranging from 256 $\times$ 256 to 2048 $\times$ 2048, demonstrate the robustness and effectiveness of the current selected IEPO. Additionally, the strong results obtained across 5 different proxy metrics (as detailed in Sec. 5.3 and in response to Q1 from reviewer MqmJ) demonstrate that IEPO performs well with commonly used metrics such as MSE, MS-SSIM, and LPIPS for image reconstruction and quality assessment.

W4: While the paper includes several ablation studies, it does not provide a fine-grained analysis of the Step-Wise Denoising Training (SDT) component and the specific elements of the cache router. These elements are integral to HarmoniCa's efficacy, but without a breakdown of their contributions, it remains unclear how each sub-component (e.g., caching frequency, router sensitivity) directly impacts performance and efficiency.

We have provided an analysis of the Step-Wise Denoising Training (SDT) in Lines 295-297 included in our initial submission, where we analyze that SDT can reduce error at each timestep by incorporating information from prior timesteps during training. Additionally, our router is represented simply as a tensor with the shape $T \times N$, and does not involve any complex structures or design choices such as caching frequency or router sensitivity. Furthermore, we have presented a thorough ablation study in Sec. 5.3 in our initial submission. For instance, both SDT and IEPO achieve improvements exceeding 10 FID points over the baseline.

W5: The paper provides comparisons with Learning-to-Cache and FORA but lacks discussion of other potential caching and acceleration techniques, such as model pruning or quantization, commonly explored in efficiency research for generative models. While these techniques are mentioned, a more in-depth comparison or consideration of combining methods could contextualize HarmoniCa’s strengths and limitations relative to other optimizations.

Please refer to the answers to W1 and W2 from reviewer IBAN, where we compare our method with other caching approaches (DeepCache [5] and Faster Diffusion [6]), pruning methods (Diff-pruning [7]), and quantization techniques (PTQ4DiT [8] and EfficientDM [9]). All results validate the significant superiority of our method in both performance and efficiency. Furthermore, in Sec. I included in our initial submission, our method outperforms $\Delta$-DiT [10]. In Sec. F, the results demonstrate that combining quantization with our HarmoniCa method can achieve further acceleration without substantial performance degradation.

W6: Although HarmoniCa shows improvements in image quality (e.g., FID, IS), the paper could benefit from qualitative analyses of generated images. This would help illustrate how the framework’s focus on harmonizing training and inference affects visual output, potentially highlighting differences in quality between HarmoniCa and baseline methods.

We add the qualitative comparison and analyses in Sec. S. The more accurate details and objective-level traits help confirm HarmoniCa’s superiority over other baselines.

Q1: Could the authors elaborate on the criteria for selecting the error proxy in the Image Error Proxy-Guided Objective (IEPO)? Specifically, how was the proxy choice validated across different image generation tasks, and did the authors observe any trade-offs when experimenting with alternative metrics?

Please refer to the answer to the W3.

Q2: Could the authors provide additional insights into the computational resources required to implement HarmoniCa, especially the hardware and memory demands for training and inference?

Unlike methods such as pruning and quantization, HarmoniCa does not require any specialized hardware or custom CUDA kernel designs for training and inference. Taking DiT-XL/2 256$\times$256 as an example, the minimal memory demands are under 12 GB and approximately 4 GB of GPU memory for training and inference, respectively.

Q3: Does HarmoniCa generalize well beyond the specific tasks demonstrated, such as high-resolution image generation? Have the authors considered exploring HarmoniCa’s effectiveness for video generation or other data-intensive applications where DiTs are applied?

Yes, HarmoniCa generalizes well beyond the specific tasks demonstrated, including high-resolution image generation. For instance, in Sec. 5.2 and Sec. E included in our initial submission, we have shown that it performs effectively on 1024$\times$1024 PixArt-$\alpha$ and $\Sigma$, as well as 2048$\times$2048 PixArt-$\Sigma$. Specifically, it achieves over a 1.5 speedup ratio and a lower FID score for 1024$\times$1024 PixArt-$\alpha$ compared to the non-accelerated model. Additionally, we are actively exploring the application of our novel strategy to the video generation domain and other data-intensive tasks where DiTs are used.

Q3: Can the router trained with one diffusion sampler be applied to a different diffusion sampler? For example, the training uses a DDIM sampler, but the sampling uses a DPM solver.

As shown in the table below, the router trained with one diffusion sampler can indeed be applied to a different sampler, such as DPM-Solver++$\rightarrow$SA-Solver (5th row) and IDDPM$\rightarrow$DPM-Solver++ (9th row). However, the performance of these trials is much worse than the standard HarmoniCa. We believe this is due to the discrepancies in sampling trajectories and noise scheduling between the two samplers, which need to be accounted for during router training. In other words, the sampler used for training should match the one used during inference to improve the performance. We have also included the results and relevant discussions in Sec. N.

Q4: Implementation details:

1. Is  $r_{t,i}$ implemented as an output of sigmoid function following Learning-to-Cache?
2. Is  $r_{t,i}$ updated using the vanilla gradient descent or other optimizers?
3. Does Harmonica apply the stop gradient to gen_proxy? In other words,  $\lambda^{(t)}$ is an adaptive weight of $\mathcal{L}\_{MSE}$.

1. Yes (as mentioned in Sec. D).
2. We employ AdamW optimizer with the default setting in PyTorch to update the value (as mentioned in Sec. 5.1).
3. Yes, we directly employ with torch.no_grad(): to wrap gen_proxy.

[1] Ning M, Sangineto E, Porrello A, et al. Input perturbation reduces exposure bias in diffusion models[J]. arXiv preprint arXiv:2301.11706, 2023.

[2] Li M, Qu T, Yao R, et al. Alleviating exposure bias in diffusion models through sampling with shifted time steps[J]. arXiv preprint arXiv:2305.15583, 2023.

Thanks for the reviewer’s constructive comments.

W1: Harmonica introduces additional hyperparameters, such as the coefficient  $\beta$ and $C$update interval , which require more computation to tune. While this additional offline cost does not affect the inference speedup, the paper should still provide the cost spent on tuning the hyperparameters to help readers better understand the full implementation costs.

We adopt a grid search approach by sampling several values for $C$ and benchmarking feature cache performance (as shown in Fig. 7 in the initial submission). Based on this, we fix $C$ as 500 for all experiments. The improved performance and acceleration ratios validate the robustness of this fixed value which can be directly applied to other settings, making the overhead a one-time cost (about $5\times8$ H800-GPU-hours). For $\beta$, initially, we tested multiple values and found that they exhibited robustness across different configurations for a model series. Therefore, once an effective range is established, we can get a satisfactory $\beta$ with just a few trials (i.e., a few iterations). On average, determining $\beta$ takes just over 10 minutes per setting, making the process efficient and practical. In future work, we plan to explore tuning-based approaches to automate the determination of these hyperparameters.

W2: Overall, the writing could be largely improved by breaking the long sentences into shorter ones. Here are some typos I caught :

1. Line 127, "much lower training costs" -> "much lower training cost".
2. Line 213, "Harmonize" -> "harmonize".
3. Line 261, "The student" -> "the student".
4. The second DiT block of the bottom row has the same format as Teacher DiT in Figure 4 (b).

We sincerely thank the reviewer for providing valuable suggestions and pointing out these. We have fixed these typos and broken the long sentences in the paper.

W3: Terminology: "discrepancy" could be replaced with "mismatch" (as used in line 237) for better accuracy as discrepancy often refers to a lack of similarity between quantities.

We sincerely appreciate the reviewer’s thoughtful suggestion and have carefully reviewed this point. After thorough consideration, the term “discrepancy” has been retained, as it is widely used in prior studies (e.g., [1, 2]) to describe the inconsistency between training and inference. Retaining this term ensures alignment with established terminology in the field and avoids potential overlap with “Objective Mismatch.” To further address the reviewer’s concern, the definition of “discrepancy” has been clarified in Lines 52-53 of the revised manuscript. The authors hope this revision adequately addresses the reviewer’s feedback.

Q1: The authors explore different standard metrics for $\lambda^{(t)}$, but none are designed for natural images. Would perceptual metrics like LPIPS be more effective in measuring image similarity?

Thanks for the insightful suggestion. Under the same speedup ratio, we test $\text{MS-SSIM}$ and $\text{LPIPS}$ as suggested. In the table below, these metrics exhibit comparable performance compared with$\|\|x_0-x_0^{(t)}\|\|_F^2$. For instance, $\text{LPIPS}(x_0, x_0^{(t)})$ slightly outperforms in FID and sFID, while $\|\|x_0-x_0^{(t)}\|\|_F^2$ marginally excels in IS. We have included the results in Sec. M and will explore more about the influence of metric selection in the future.

Q2: Why is  $\lambda^{(t)}$ designed to be multiplied with $\mathcal{L}_{MSE}$? Are there any other alternative design choices there?

Since our goal is to use image error to guide cache usage (i.e., tend to reuse fewer features if the final error is large), $\lambda^{(t)}$, which reflects the final image error, serves as a multiplier to control the weight of $\mathcal{L}\_{MSE}$ in the loss function. A larger $\lambda^{(t)}$ (indicating a larger error) increases the emphasis on reducing $\mathcal{L}\_{MSE}$, encouraging less feature reuse, and vice versa. These explanations can be found in Lines 310-319. Using $\lambda^{(t)}$ as a simple multiplier is effective, but considering it as an exponent of $\mathcal{L}\_{MSE}$ or exploring other algebraic forms offers valuable, albeit more complex, alternatives. We plan to investigate these designs in our future research.

Thanks for the reviewer’s constructive comments.

W1: This work appears to be an incremental improvement on Learn-to-Cache, which initially made caching mechanisms learnable. The main contribution here is refining the learning mechanism of Learn-to-Cache. However, in terms of both efficiency and generation metrics, HarmoniCa shows only limited improvements.

While our work builds upon the foundational concept (i.e., learn a caching strategy) introduced by Learn-to-Cache, our contributions and improvements are both significant and distinct in the following ways:

1. Identifying Key Problems: We first identify two critical discrepancies between training and inference overlooked in Learn-to-Cache as follows：

   1. Prior timestep disregard: Ignore the error caused by reusing cached features at prior timesteps during optimization.
   2. Objective Mismatch: Focus on optimizing the intermediate output during the denoising process.

   They both significantly affect the overall performance of Learning-to-Cache.

2. Novel Training Strategies: To address these problems, we propose two novel training techniques (as recognized by reviewers zBAN and UrSw) (i.e., SDT+IEPO). Our entire optimization processes differ fundamentally compared with Learning-to-Cache:

   1. For the training paradigm, Learning-to-Cache uses random timestep sampling similar to DDPM, focusing solely on the current sampled timestep. In contrast, we implement SDT for iterative denoising from random noise, which accounts for all previous timesteps, mirroring the inference process.
   2. For the learning objective, Learning-to-Cache focuses on intermediate noise and cache usage. We further optimize the final image error and cache usage with our IEPO.

   The proposed framework achieves significantly improved performance and more efficient caching (as highlighted by reviewers zBAN, MqmJ, and UrSw).

3. Reduced Training Cost: Our method significantly reduces training costs. Compared to Learning-to-Cache, it requires significantly less training time (i.e., $\mathbf{\frac{3}{4}}$ of the training time) without any training images, making it more user-friendly and easier to implement. This is especially advantageous for advanced generative models (e.g., PixArt-$\alpha/\Sigma$), where applying Learn-to-Cache is less feasible (e.g., more training time and large pre-train datasets are required).

4. Substantial Performance Improvements:

   1. HarmoniCa consistently and pronouncedly outperforms Learn-to-Cache across all settings, particularly when fewer steps are used for sampling. For instance, in Tab. 2, HarmoniCa achieves a 1.24 improvement in FID and a 6.74 improvement in IS compared to Learn-to-Cache with the 10-step DiT-XL/2, under a higher speedup ratio.
   2. We also add IS and FID results across different speedup ratios in Sec. O, showing that our method yields significantly better results as the speedup ratio increases, e.g., 30.90 higher IS and 12.34 lower FID scores at about $\times$1.6 speedup. HarmoniCa handles the trend towards fewer timesteps and higher speedup ratios far more effectively than Learn-to-Cache.

Dear Reviewer,

Sorry for bothering you again, but the discussion period is coming to an end in two days. Could you please let us know if our responses have alleviated your concerns? If there are any further comments, we will do our best to respond.

Best regards,

Authors of Paper #1567

Dear Reviewer,

We hope this message finds you well. I apologize for the repeated follow-ups, but as the deadline to submit the revised PDF is tomorrow, we want to kindly remind you that we have provided additional clarifications and revisions in response to your comments. If there are any remaining concerns or further points you would like us to address, we would be very grateful for your feedback and would be happy to make any necessary adjustments.

We truly appreciate your time and thoughtful consideration in reviewing our work.

Thank you so much for your understanding and support.

Best regards,

Authors of Paper #1567

Dear Reviewer,

We hope this message finds you well. Apologies for following up again, but as the discussion period is coming to a close, we wanted to kindly check if our responses have sufficiently addressed your concerns. We have carefully revised the manuscript based on your feedback and have provided further clarifications in the comment section.

If you have any additional points or questions, we would be happy to address them. We truly appreciate your time and thoughtful consideration, and we are grateful for the opportunity to improve the manuscript with your feedback.

If you're celebrating, we would like to wish you a joyful and restful Thanksgiving!

Thank you once again for your continued support.

Best regards,

Authors of Paper #1567

Dear Reviewer Ukn2,

We hope this message finds you well. We apologize for reaching out repeatedly, but with less than two days remaining in the discussion period, we wanted to check whether our updates have addressed your concerns kindly.

If you have any further questions or require additional clarification, we are more than happy to provide it. Your feedback has been invaluable, and we greatly appreciate your time and effort in reviewing our work.

Thank you again for your continued support. We look forward to your feedback.

Best regards,

Authors of Paper #1567

Dear Reviewer Ukn2,

We apologize for reaching out repeatedly again. We would like to sincerely thank you for your thoughtful review and valuable feedback on our paper. We wanted to kindly check if our responses have sufficiently addressed your concerns. We have carefully revised the manuscript based on your feedback and have provided further clarifications in the comment section.

If you find that our revisions have satisfactorily addressed your concerns, we would be grateful if you could consider reflecting this in your final assessment.

Thank you again for your valuable contributions to our research.

Warm regards,

Authors of Paper #1567

Dear Reviewer,

We hope this message finds you well. We apologize for the repeated follow-ups and want to kindly remind you that we have provided additional clarifications and revisions in response to your comments. If there are any remaining concerns or further points you would like us to address, we would be very grateful for your feedback and would be happy to make any necessary adjustments.

We truly appreciate your time and thoughtful consideration in reviewing our work.

Thank you so much for your understanding and support.

Best regards,

Authors of Paper #1567

Thanks for the reviewer’s constructive comments.

W1: This paper lacks experimental comparisons with other cache-based methods, such as DeepCache and Faster Diffusion.

To highlight HarmoniCa's advantages, we compare it with DeepCache and Faster Diffusion. Notably, DeepCache is constrained by the U-shaped structure, making it unsuitable for DiTs. Since our approach is not constrained by any specific structure, we use the U-Vit[1] model with the ImageNet 256$\times$256 dataset on a single A6000 GPU. We have also included the results and experimental details in Sec. K. As shown in the table, HarmoniCa achieves a minimal FID increase of less than 0.05, while providing a 1.65$\times$ speedup, outperforming both Faster Diffusion and DeepCache.

W2: A performance comparison with other related methods under the same acceleration ratio, such as pruning or quantization, should be provided for further validation.

As shown in the table below, we compare our HarmoniCa with advanced quantization (w8a8) and pruning methods using DiT-XL/2 on ImageNet 256$\times$256. We have also included the results and experimental details in Sec. L. Our method significantly outperforms other methods, demonstrating the substantial benefit of feature cache for accelerating DiT models. It is important to note that the speedup ratio for quantization is partially determined by hardware support, which we do not rely on. Additionally, our method is orthogonal to these approaches, meaning it can be combined with them for further acceleration, e.g., 1.80$\times$speedup on H800 (results of EfficientDM + HarmoniCa have been presented in Sec. F included in our initial submission).

Q1: Can the authors provide more comprehensive comparison results with other cache-based models?

Please refer to the answer to the W1.

Q2: What are the comparative results of other related methods under the same acceleration ratio?

Please refer to the answer to the W2.

[1] Bao F, Nie S, Xue K, et al. All are worth words: A vit backbone for diffusion models[C]//Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2023: 22669-22679.

[2] He Y, Liu J, Wu W, et al. Efficientdm: Efficient quantization-aware fine-tuning of low-bit diffusion models[J]. arXiv preprint arXiv:2310.03270, 2023.

[3] Wu J, Wang H, Shang Y, et al. PTQ4DiT: Post-training Quantization for Diffusion Transformers[J]. arXiv preprint arXiv:2405.16005, 2024.

[4] Fang, Gongfan et al. “Structural Pruning for Diffusion Models.” ArXiv abs/2305.10924 (2023): n. pag.