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

Thanks to the reviewer for the constructive comments.

Literature Review and Citations

• Missing work. We appreciate the suggestions and will include ToMe, TokenFusion, and attention-driven efficiency methods in the revision.
• Citation updates. We will replace arXiv references with conference/journal versions where available.

Experiments

• Baselines. We have already included $\Delta$-DiT in Sec. H, showing our method outperforms it in both accuracy and efficiency. For [1], it was accepted by CVPR after the ICML deadline and appeared on arXiv after Oct. 2024. According to ICML policy, it qualifies as a concurrent work. Moreover, its code is not available, making comparison infeasible within the rebuttal window.
• Performance Significance. We respectfully disagree since a larger batch or model size can show larger absolute speedup. While some latency results are based on small models (e.g., DiT-XL/2 with batch size 8) for controlled ablations or comparisons, they still show consistent speedups without sacrificing quality (e.g., Tab. 2). On larger models like PixArt-$\Sigma$ $2048\times2048$, our method yields much larger absolute gains and higher acceleration ratios (Sec. E).
• Failure Cases. Across all tested settings (i.e., models, datasets, and samplers), our method shows robust and consistent performance. We did not observe clear failure cases that warrant separate analysis. Moreover, such analysis is uncommon in related works and not part of standard practice in this field.

Methods

• Novelty Concerns. While our work builds upon the foundational concept (i.e., learn a caching strategy) introduced by Learning-to-Cache (L2C), our contributions and improvements are both significant and distinct:

  1. Identifying Key Problems: We first identify two critical discrepancies between training and inference overlooked in L2C 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 L2C.

  2. Novel Training Strategies: To address these problems, we propose two novel training techniques(i.e., SDT+IEPO). Our entire optimization processes differ fundamentally compared with L2C:

     1. For the training paradigm, L2C 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, L2C focuses on intermediate noise and cache usage. We further optimize the final image error and cache usage with our IEPO.

  3. Reduced Training Cost: For our method, no training images are needed, and training time is much shorter than L2C. This makes it applicable to large-scale models like PixArt.

  4. Substantial Performance Improvements: Our method consistently outperforms L2C. For instance, Tab. 2 shows +1.24 FID and +6.74 IS on DiT-XL/2. Fig. 7 shows +30.90 IS and −12.34 FID at $1.6\times$ speedup.

• Generalization. Our router is static once trained and does not change per sample during inference. The method is also image-free during training without tuning the weights, reducing the risk of overfitting. Moreover, we can observe strong generalization across multiple datasets and model sizes in this paper.
• Theoretical Foundation. The reviewer suggests using a teacher-forcing strategy—this is exactly what SDT employs during training. Our experiments (e.g., Fig. 5, Tab. 9) validate its effectiveness in reducing error accumulation. A more detailed theoretical analysis is planned for future work, as we mainly focus on practical speedup while maintaining the performance in this work.

Questions for Authors

• Prompts in Figs. 12, 15, and 16 are randomly generated using GPT-4.
• The router structure is identical to L2C, with all the details in Sec. 5.1 and Sec. D.

[1] You et. al. Layer- and Timestep-Adaptive Differentiable Token Compression Ratios for Efficient Diffusion Transformer. Arxiv 2024.12

Thanks for the reviewer’s constructive comments.

Claims And Evidence

• Comparison with rectified flow models. We clarify that our approach is not based on pretraining, unlike flow-based models. A direct comparison is not entirely appropriate, as our method merely trains a small tensor (i.e., the Router) and only focuses on improving inference efficiency in a post-training manner. Thus, we would like to emphasize that our methods can bring significant benefits compared with other post-training methods (e.g., quantization, pruning, and other caching methods in Sec. 5). Moreover, HarmoniCa can be effectively applied to rectified flow models. We employ the pretrained models and evaluation in LFM. As shown in the following table ( $T=100$, Euler Solver, and $BS=8$), our method achieves substantial speedup ratios without performance degradation.

  Datasets	Models	FID$\downarrow$	Latency (s)$\downarrow$	Speedup$\uparrow$
  LSUN-Bedroom $256\times256$	DiT-L/2	5.22	1.76	-
  + HarmoniCa	DiT-L/2	5.13	1.09	$1.62\times$
  CelebaHQ $256\times256$	DiT-L/2	5.42	1.76	-
  + HarmoniCa	DiT-L/2	5.41	1.07	$1.65\times$
  ImageNet $256\times256$	DiT-B/2 (CFG=1.5)	5.06	1.37	-
  + HarmoniCa	DiT-B/2 (CFG=1.5)	5.04	0.87	$1.58\times$

• Relationship between the discrepancy and the caching system. We further explain the relationship as below (also can be found in Sec. 4.1):

  1. Prior Timestep Disregard: In previous research, such as Learning-to-Cache, the training process employs a raw image with manually added noise as input $\mathbf{x}_t$(see Fig. 3), which is unaffected by any cached feature reuse. However, during accelerated inference, the input (i.e., a noisy image) is disrupted by cached feature reuse in earlier timesteps (as illustrated in Lines 187-191). Moreover, the cache context in current timestep is shaped by the caching strategy in earlier timesteps (e.g., $t +1,\ldots, T$), which is unconsidered as shown in Fig. 3. This discrepency of disregarding the caching effects of prior timestep can lead to significant error, which is caused by cache usage, accumulation (see Fig. 5). To address this, SDT (Sec. 4.2) ensures that the model learns to incorporate the prior timestep's effect on the cache, aligning the training process with the inference stage and reducing error accumulation.
  2. Objective Mismatch: Equipped with the caching strategy, the training objective in previous research typically minimizes each timestep noise prediction error caused by caching usage. However, inference aims to generate high-quality images, the same as those from non-accelerated models. This mismatch between the two processes leads to an optimization deviation (Sec. C) and image distortion (Fig. 6). By integrating IEPO (Sec. 4.2), we ensure that the training process considers the final image error, which mitigates the optimization shift and improves the quality of generated images.

Relation To Broader Scientific Literature

We appreciate the reviewer’s insight. While Learning-to-Cache introduced the concept, HarmoniCa extends this by addressing the discrepancies between training and inference, and also requires fewer training resources (Sec. 4.3). Moreover, HarmoniCa achieves much higher efficiency and performance across all metrics compared with previous methods, including caching, quantization, and pruning. Regarding the concern about the marginal improvements, we acknowledge that the gains in some small models may appear modest. However, the benefits become more pronounced in large-scale models (e.g., Figs. 3 and 12). Besides, we emphasize that HarmoniCa can be effectively applied to rectified flow models (see Claims And Evidence).

Other Strengths And Weaknesses

• Lacks a clear logical argument. Regarding the concern, we clarify our general and natural logical argument process as follows. First, we have carefully analyzed the two discrepancies between training and inference. Then, we proposed our framework to solve the problems. Both parts include detailed logical analyses and experimental validation. These processes provide the rationality to support the large number of experimental findings in Sec. 5. We will make the paper logic clearer in the revision.
• Hyperparameters limit the adaptability. We would like to clarify that we have three hyperparameters in total: interval $C$, threshold $\tau$, and regularization coefficient $\beta$. The superior performance and speedup across all experiments validate the efficacy of using fixed $C$ and $\tau$. Therefore, the only one that requires adjustment is $\beta$, which determines the trade-off between acceleration and image quality. We will explore the automatic methods to determine $\beta$ in the future.

Thanks for the reviewer’s constructive comments.

Methods And Evaluation Criteria

Contribution insufficiently novel. While our work builds upon the foundational concept (i.e., learn a caching strategy) introduced by Learning-to-Cache (L2C), our contributions and improvements are both significant and distinct:

1. Identifying Key Problems: We first identify two critical discrepancies between training and inference overlooked in L2C 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 L2C.

2. Novel Training Strategies: To address these problems, we propose two novel training techniques(i.e., SDT+IEPO). Our entire optimization processes differ fundamentally compared with L2C:

   1. For the training paradigm, L2C 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, L2C focuses on intermediate noise and cache usage. We further optimize the final image error and cache usage with our IEPO.

3. Reduced Training Cost: For our method, no training images are needed, and training time is much shorter than L2C. This makes it applicable to large-scale models like PixArt.

4. Substantial Performance Improvements: Our method consistently outperforms L2C. For instance, Tab. 2 shows +1.24 FID and +6.74 IS on DiT-XL/2. Fig. 7 shows +30.90 IS and −12.34 FID at $1.6\times$ speedup.

Experimental Designs Or Analyses

• The performance improvement of Learning-to-Cache is quite limited. First, both methods without re-training the model gain comparable performance, which is also much better than that of the non-accelerated models in Tab. 13. Thus, it is reasonable to show our superiority with experiments at more challenging settings (i.e., higher CUR). For example, in Fig. 7, HarmoniCa can achieve substantially higher IS and lower FID scores, 30.90 and 12.34, under $1.6\times$speedup. Moreover, our method consumes less training time in these experiments, since we employ the same training settings (e.g., iterations and batch size) as Learning-to-Cache (analysis can be found in Sec. 4.3).
• The results in the tables do not match those in the original papers. For the performance and speedup gaps of Learning-to-Cache between this work in Tab. 2 and the official paper, we have to clarify that we re-run the official implementation of Learning-to-Cache with the provided arguments without any modification. However, we conducted training procedures on 4 H800 GPUs and evaluation phases on 8 H800 GPUs, contrasting with the original experimental setup employing 8 A6000 GPUs for both stages.
• The results in Tab. 9 are inconsistent with Tab. 2. As mentioned in Sec. I, the implementation in the code of Learning-to-Cache is inconsistent with that in Alg. 1 in the official paper. The former only employs the caching strategy in non-adjacent steps, and the latter utilizes the caching strategy across all steps (the same as ours). The trick in the former can help mitigate the impact of accumulative error by the no-caching steps (4th row vs. 5th row in Tab. 15). To make a fair comparison, we implement Alg. 1 in the official paper and report the results in Tab. 9.

Thanks for the reviewer’s constructive comments.

Claims And Evidence

• Concerns regarding acceleration rates and resolution scaling. Thanks for the insightful question. We agree that the reason why the acceleration rates of HarmoniCa decline ($1.67\times\rightarrow1.56\times$) with the resolution increase ($256\times256\rightarrow2048\times2048$) is due to memory bottlenecks. Specifically, HarmoniCa achieves similar CUR (i.e., cache usage ratio) for these resolutions, but requires a cache size increase from 936 MB (BS=8) to 4320 MB (BS=1). Thus, we believe preloading and asynchronous saving can help resolve the problem and will explore this in the future.

Methods And Evaluation Criteria

• Sampler consistency during training and inference.  The sampler used during training is not required to match the one used at inference. We have already included the discussion in Sec. K. However, mismatched samplers can lead to a discrepancy in sampling trajectories. This does affect performance as reflected in our experiments. Therefore, we suggest that the sampler should be the same for optimal performance.

• The evaluation metrics rely on FID for text-to-image generation.  We have conducted experiments with more metrics (e.g., Image-Reward, LPIPS, and PSNR), which are provided in the Appendix (Sec. J). Here, we also consider DINO and human evaluations (e.g., HPSv2 and PickScore) as suggested. As shown in the following tables, our method with a higher speedup ratio far outperforms baselines and achieves comparable performance with non-accelerated models.

  We use the configurations and models in Tab. 2 in the manuscript to generate images for evaluation.

  Model	DINO $\uparrow$	HPSv2 $\uparrow$	PickScore $\uparrow$	Latency (s)$\downarrow$	Speedup$\uparrow$
  PixArt-$\alpha$ $256\times256$	0.3082	28.91	27.89	0.553	-
  + 15 steps	0.2582	27.98	23.02	0.418	$1.32\times$
  + FORA	0.2712	28.11	22.44	0.364	$1.52\times$
  + HarmoniCa	0.3235	28.72	26.65	0.346	$1.60\times$
  PixArt-$\alpha$ $512\times512$	0.3339	30.53	28.52	1.759	-
  + 15 steps	0.3127	29.79	22.03	1.291	$1.36\times$
  + FORA	0.3099	29.82	21.98	1.150	$1.53\times$
  + HarmoniCa	0.3289	30.28	27.47	1.072	$1.64\times$

Experimental Designs Or Analyses:

• Compute cost for PixArt experiments. Thank you for the question. We provide the computation costs of training and inference in the following table. It is worth noting that we only use naive DDP training without any advanced strategies like gradient checkpointing and FSDP. Thus, the requirements of computation can be further decreased in practice.

  The settings of training can be found in Sec. 5. We employ BS=1 here to obtain the memory consumption of inference.

  Model	Infer. Mem. (GB/GPU)$\downarrow$	Train. Mem. (GB/GPU)$\downarrow$	Train. Time (h)$\downarrow$
  PixArt-$\alpha$ $256\times256$	20.86	-	-
  + HarmoniCa	21.21	79.05	1.84
  PixArt-$\alpha$ $512\times512$	24.11	-	-
  + HarmoniCa	24.61	77.30	2.92

• Absence of standard deviation in experimental results. Thank you for the suggestion. We conducted five independent runs with different random seeds on the following setting: PixArt-$\alpha$ $256\times256$, using DPM-Solver++ (20 steps) and evaluating on 5,000 images. The results show high consistency across runs, and more importantly, our caching-accelerated models consistently outperform the non-accelerated models in all evaluation metrics, confirming our method's robustness and effectiveness.

  ”x/y’’ in the table denotes the results that come from non-accelerated models and HarmoniCa. We use the same settings in Tab. 2 in the manuscript.

  Seed	IS$\uparrow$	FID$\downarrow$	sFID$\downarrow$
  8	33.28/33.27	37.31/35.44	94.78/92.10
  16	33.61/33.62	37.43/35.46	94.74/92.11
  24	33.59/33.60	37.55/35.48	95.01/92.32
  32	33.65/33.65	37.11/35.34	94.87/92.13
  40	33.64/33.64	37.05/35.29	94.88/92.19
  deviation	0.138/0.144	0.188/0.073	0.093/0.081