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

要約

拡散トランスフォーマー (DiT) は、生成タスクにおける卓越したスケーラビリティと並外れたパフォーマンスにより注目を集めています。
ただし、かなりの推論コストが実用的な展開の妨げとなります。
特徴キャッシュ メカニズムには、タイムステップにわたる冗長な計算の保存と取得が含まれており、拡散モデルにおけるステップごとの推論時間を短縮することが期待されています。
DiT の既存のキャッシュ方法のほとんどは手動で設計されています。
学習ベースのアプローチは戦略を適応的に最適化しようとしますが、トレーニングと推論の間に不一致が発生し、パフォーマンスと加速率の両方が妨げられます。
詳細な分析により、これらの不一致は主に 2 つの側面に起因することがわかりました。(1) 前のタイムステップの無視。トレーニングでは、より早いタイムステップでのキャッシュ使用の影響が無視されます。(2) 目的の不一致。トレーニングのターゲット (予測ノイズを調整する)。
各タイムステップ) は、推論の目的 (高品質の画像の生成) から逸脱します。
これらの矛盾を軽減するために、段階的ノイズ除去トレーニング (SDT) と画像エラー プロキシ ガイド付き目標 (IEPO) に基づいて構築された新しい学習ベースのキャッシング フレームワークを使用してトレーニングと推論を調和させる新しい方法である HarmoniCa を提案します。
従来のトレーニング パラダイムと比較して、新しく提案された SDT はノイズ除去プロセスの連続性を維持し、推論時の動作方法と同様に、モデルがトレーニング中に以前のタイムステップからの情報を活用できるようにします。
さらに、キャッシュされた特徴の再利用によって引き起こされる最終的な画像エラーを近似するための効率的なプロキシ メカニズムを統合する IEPO を設計します。
したがって、IEPO は、最終的な画質とキャッシュ使用率のバランスをとるのに役立ち、各タイムステップでの予測出力に対するキャッシュ使用量の影響のみを考慮するトレーニングの問題を解決します。

要約(オリジナル)

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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google