HarmoniCa: Het harmoniseren van training en inferentie voor een betere kenmerkcache in de versnelling van de Diffusion Transformer

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

Samenvatting

Diffusie Transformers (DiTs) hebben aan populariteit gewonnen vanwege hun uitstekende schaalbaarheid en buitengewone prestaties bij generatieve taken. Echter, hun aanzienlijke inferentiekosten belemmeren praktische implementatie. Het kenmerkcache-mechanisme, dat het opslaan en ophalen van redundante berekeningen over tijdstappen omvat, biedt hoop om de inferentietijd per stap te verminderen in diffusiemodellen. De meeste bestaande cache-methoden voor DiT zijn handmatig ontworpen. Hoewel de op leermethoden gebaseerde aanpak probeert strategieën adaptief te optimaliseren, lijdt het onder inconsistenties tussen training en inferentie, wat zowel de prestaties als de versnelling vermindert. Na gedetailleerde analyse wijzen we aan dat deze inconsistenties voornamelijk voortkomen uit twee aspecten: (1) Voorafgaande Tijdstapverwaarlozing, waarbij training het effect van cachegebruik in eerdere tijdstappen negeert, en (2) Doelstelling Mismatch, waarbij het trainingsdoel (afstemmen van voorspelde ruis in elk tijdstip) afwijkt van het doel van inferentie (het genereren van de hoogwaardige afbeelding). Om deze inconsistenties te verlichten, stellen we HarmoniCa voor, een nieuw methode die training en inferentie harmoniseert met een nieuw op leermethoden gebaseerd Caching-framework gebouwd op Stapsgewijze Denoising Training (SDT) en Begeleide Doelstelling voor Afbeeldingsfouten (IEPO). In vergelijking met het traditionele trainingsparadigma behoudt de nieuw voorgestelde SDT de continuïteit van het denoising-proces, waardoor het model informatie kan benutten van eerdere tijdstappen tijdens training, vergelijkbaar met de manier waarop het werkt tijdens inferentie. Verder ontwerpen we IEPO, dat een efficiënt proxy-mechanisme integreert om de uiteindelijke afbeeldingsfout te benaderen die wordt veroorzaakt door het hergebruik van de gecachte functie. Daarom helpt IEPO bij het balanceren van de uiteindelijke beeldkwaliteit en cachegebruik, waardoor het probleem van training wordt opgelost dat alleen de impact van cachegebruik op de voorspelde output bij elk tijdstip overweegt.

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.