arXiv:2410.01723cs.CV2024-10ICML被引 17

提出HarmoniCa框架,让扩散模型推理缓存更高效且不损失图像质量。

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

  • 引入分步去噪训练,使前后步骤可复用计算结果。
  • 用图像误差代理目标平衡画质与缓存效率,实测延迟降40%以上。
  • 无需生成图像即可训练,节省25%训练时间,适合部署优化场景。

扩散Transformer(DiTs)在生成任务中表现优异,但推理成本高制约实际部署。特征缓存通过存储重复计算可提升效率,现有基于学习的缓存方法忽视前一时刻影响,且训练与推理目标不一致(预测噪声 vs 高质量图像),导致性能与效率双重受损。为此,本文提出HarmoniCa框架,首次实现训练与推理的协同优化。核心包含:分步去噪训练(SDT)确保去噪过程连续性,使前期步骤可被利用;图像误差代理引导目标(IEPO),通过高效代理逼近图像误差,在保证画质前提下优化缓存利用率。在8个模型、4种采样器及256×256至2K分辨率下广泛测试,结果表明该框架显著提升性能与速度。例如,在PixArt-α上实现超过40%的延迟降低(即理论加速2.07倍),并相较之前方法减少25%训练时间。代码已开源。

原文摘要 · Abstract (English)

Diffusion Transformers (DiTs) excel in generative tasks but face practical deployment challenges due to high inference costs. Feature caching, which stores and retrieves redundant computations, offers the potential for acceleration. Existing learning-based caching, though adaptive, overlooks the impact of the prior timestep. It also suffers from misaligned objectives--aligned predicted noise vs. high-quality images--between training and inference. These two discrepancies compromise both performance and efficiency. To this end, we harmonize training and inference with a novel learning-based caching framework dubbed HarmoniCa. It first incorporates Step-Wise Denoising Training (SDT) to ensure the continuity of the denoising process, where prior steps can be leveraged. In addition, an Image Error Proxy-Guided Objective (IEPO) is applied to balance image quality against cache utilization through an efficient proxy to approximate the image error. Extensive experiments across $8$ models, $4$ samplers, and resolutions from $256\times256$ to $2K$ demonstrate superior performance and speedup of our framework. For instance, it achieves over $40\%$ latency reduction (i.e., $2.07\times$ theoretical speedup) and improved performance on PixArt-$α$. Remarkably, our image-free approach reduces training time by $25\%$ compared with the previous method. Our code is available at https://github.com/ModelTC/HarmoniCa.

扩散模型缓存优化推理加速

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。