HarmoniCa: Harmonisierung von Training und Inferenz für einen besseren Merkmalscache in der Beschleunigung von Diffusionstransformatoren

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

October 2, 2024

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

cs.AI

Zusammenfassung

Diffusions-Transformer (DiTs) haben aufgrund ihrer hervorragenden Skalierbarkeit und außergewöhnlichen Leistungsfähigkeit bei generativen Aufgaben an Bedeutung gewonnen. Die beträchtlichen Inferenzkosten behindern jedoch die praktische Anwendung. Der Merkmalscache-Mechanismus, der das Speichern und Abrufen redundanter Berechnungen über Zeitschritte hinweg beinhaltet, verspricht, die Inferenzzeit pro Schritt bei Diffusionsmodellen zu reduzieren. Die meisten bestehenden Caching-Methoden für DiT sind manuell entworfen. Obwohl der lernbasierte Ansatz versucht, Strategien adaptiv zu optimieren, leidet er unter Diskrepanzen zwischen Training und Inferenz, was sowohl die Leistung als auch das Beschleunigungsverhältnis beeinträchtigt. Nach einer detaillierten Analyse identifizieren wir, dass diese Diskrepanzen hauptsächlich aus zwei Aspekten resultieren: (1) Vernachlässigung des vorherigen Zeitschritts, bei dem das Training die Auswirkung der Cache-Nutzung in früheren Zeitschritten ignoriert, und (2) Zielabweichung, bei der das Trainingsziel (Anpassung des vorhergesagten Rauschens in jedem Zeitschritt) vom Ziel der Inferenz (Erzeugung des hochwertigen Bildes) abweicht. Um diese Diskrepanzen zu mildern, schlagen wir HarmoniCa vor, eine neuartige Methode, die das Training und die Inferenz mit einem neuartigen lernbasierten Caching-Framework harmonisiert, das auf dem schrittweisen Rauschunterdrückungstraining (SDT) und dem bildfehlergeführten Ziel (IEPO) basiert. Im Vergleich zum traditionellen Trainingsparadigma erhält das neu vorgeschlagene SDT die Kontinuität des Rauschunterdrückungsprozesses aufrecht, was es dem Modell ermöglicht, Informationen aus vorherigen Zeitschritten während des Trainings zu nutzen, ähnlich wie es während der Inferenz funktioniert. Darüber hinaus entwerfen wir IEPO, das einen effizienten Proxy-Mechanismus integriert, um den endgültigen Bildfehler zu approximieren, der durch die Wiederverwendung des gecachten Merkmals verursacht wird. Daher hilft IEPO, die Balance zwischen der endgültigen Bildqualität und der Cache-Nutzung herzustellen und das Problem des Trainings zu lösen, das nur die Auswirkung der Cache-Nutzung auf die vorhergesagte Ausgabe in jedem Zeitschritt berücksichtigt.

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.