HarmoniCa : Harmonisation de l'entraînement et de l'inférence pour une meilleure mise en cache des caractéristiques dans l'accélération du transformateur de diffusion

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

October 2, 2024

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

cs.AI

Résumé

Les Transformateurs de Diffusion (DiTs) ont gagné en importance pour leur extensibilité remarquable et leurs performances extraordinaires dans les tâches génératives. Cependant, leurs coûts d'inférence considérables entravent le déploiement pratique. Le mécanisme de cache de caractéristiques, qui consiste à stocker et récupérer des calculs redondants à travers les pas de temps, offre des perspectives pour réduire le temps d'inférence par étape dans les modèles de diffusion. La plupart des méthodes de mise en cache existantes pour les DiTs sont conçues manuellement. Bien que l'approche basée sur l'apprentissage tente d'optimiser les stratégies de manière adaptative, elle souffre de divergences entre l'entraînement et l'inférence, ce qui nuit à la fois aux performances et au ratio d'accélération. Après une analyse détaillée, nous identifions que ces divergences proviennent principalement de deux aspects : (1) Mépris du Pas de Temps Antérieur, où l'entraînement ignore l'effet de l'utilisation du cache aux pas de temps antérieurs, et (2) Divergence d'Objectif, où la cible d'entraînement (aligner le bruit prédit à chaque pas de temps) s'écarte de l'objectif de l'inférence (générer une image de haute qualité). Pour atténuer ces divergences, nous proposons HarmoniCa, une méthode novatrice qui harmonise l'entraînement et l'inférence avec un nouveau cadre d'apprentissage basé sur la mise en cache construit sur l'Entraînement au Débruitage par Étapes (SDT) et l'Objectif Guidé par un Proxie d'Erreur d'Image (IEPO). Comparé au paradigme d'entraînement traditionnel, le SDT nouvellement proposé maintient la continuité du processus de débruitage, permettant au modèle de tirer parti des informations des pas de temps antérieurs pendant l'entraînement, de manière similaire à son fonctionnement pendant l'inférence. De plus, nous concevons IEPO, qui intègre un mécanisme de proxie efficace pour approximer l'erreur d'image finale causée par la réutilisation de la caractéristique mise en cache. Par conséquent, IEPO aide à équilibrer la qualité finale de l'image et l'utilisation du cache, résolvant le problème de l'entraînement qui ne prend en compte que l'impact de l'utilisation du cache sur la sortie prédite à chaque pas de temps.

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.