HarmoniCa: ディフュージョントランスフォーマーの高速化におけるトレーニングと推論の調和を図るための特徴キャッシュの向上

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

October 2, 2024

著者: Yushi Huang, Zining Wang, Ruihao Gong, Jing Liu, Xinjie Zhang, Jinyang Guo, Xianglong Liu, Jun Zhang

cs.AI

要旨

拡散トランスフォーマー（DiTs）は、生成タスクにおける優れたスケーラビリティと非凡なパフォーマンスによって注目を集めています。しかしながら、著しい推論コストは実用的な展開を妨げています。特徴キャッシュメカニズムは、タイムステップ間で冗長な計算を保存および取得することを含むことで、拡散モデルにおける1ステップあたりの推論時間を削減する可能性を秘めています。DiT向けのほとんどの既存キャッシング手法は手動で設計されています。学習ベースのアプローチは戦略を適応的に最適化しようとしますが、訓練と推論の間に不一致が生じ、パフォーマンスと加速比の両方に支障をきたします。詳細な分析の結果、これらの不一致は主に2つの側面から生じることが明らかになりました：（1）前のタイムステップの無視、つまり訓練が以前のタイムステップでのキャッシュ使用の影響を無視すること、および（2）目的の不一致、つまり訓練目標（各タイムステップでのノイズの予測を整合させる）が推論の目標（高品質画像の生成）から逸脱していること。これらの不一致を緩和するために、我々はHarmoniCaを提案します。これは、新しい学習ベースのキャッシングフレームワークをStep-Wise Denoising Training（SDT）およびImage Error Proxy-Guided Objective（IEPO）に基づいて構築し、訓練と推論を調和させる方法です。従来の訓練パラダイムと比較して、新しく提案されたSDTは、ノイズ除去プロセスの連続性を維持し、モデルが訓練中に以前のタイムステップからの情報を活用できるようにし、推論時と同様に動作します。さらに、我々はIEPOを設計しました。これは、再利用されたキャッシュされた特徴によって引き起こされる最終画像エラーを近似する効率的なプロキシメカニズムを統合しています。したがって、IEPOは最終画像品質とキャッシュ利用のバランスをとり、訓練が各タイムステップでの予測出力へのキャッシュ使用の影響のみを考慮する問題を解決します。

English

Diffusion Transformers (DiTs) have gained prominence for outstanding scalability and extraordinary performance in generative tasks. However, their considerable inference costs impede practical deployment. The feature cache mechanism, which involves storing and retrieving redundant computations across timesteps, holds promise for reducing per-step inference time in diffusion models. Most existing caching methods for DiT are manually designed. Although the learning-based approach attempts to optimize strategies adaptively, it suffers from discrepancies between training and inference, which hampers both the performance and acceleration ratio. Upon detailed analysis, we pinpoint that these discrepancies primarily stem from two aspects: (1) Prior Timestep Disregard, where training ignores the effect of cache usage at earlier timesteps, and (2) Objective Mismatch, where the training target (align predicted noise in each timestep) deviates from the goal of inference (generate the high-quality image). To alleviate these discrepancies, we propose HarmoniCa, a novel method that Harmonizes training and inference with a novel learning-based Caching framework built upon Step-Wise Denoising Training (SDT) and Image Error Proxy-Guided Objective (IEPO). Compared to the traditional training paradigm, the newly proposed SDT maintains the continuity of the denoising process, enabling the model to leverage information from prior timesteps during training, similar to the way it operates during inference. Furthermore, we design IEPO, which integrates an efficient proxy mechanism to approximate the final image error caused by reusing the cached feature. Therefore, IEPO helps balance final image quality and cache utilization, resolving the issue of training that only considers the impact of cache usage on the predicted output at each timestep.