HarmoniCa: Armonizzare l'Addestramento e l'Inferenza per una Migliore Cache delle Caratteristiche nell'Accelerazione del Diffusion Transformer

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

October 2, 2024

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

cs.AI

Abstract

I Transformer di Diffusione (DiTs) hanno guadagnato importanza per la loro eccezionale scalabilità e straordinarie prestazioni nelle attività generative. Tuttavia, i considerevoli costi di inferenza ostacolano la loro implementazione pratica. Il meccanismo di cache delle caratteristiche, che coinvolge l'archiviazione e il recupero di calcoli ridondanti attraverso i passaggi temporali, promette di ridurre il tempo di inferenza per passaggio nei modelli di diffusione. La maggior parte dei metodi di caching esistenti per i DiT sono progettati manualmente. Anche se l'approccio basato sull'apprendimento cerca di ottimizzare le strategie in modo adattivo, soffre di discrepanze tra addestramento e inferenza, che compromettono sia le prestazioni che il rapporto di accelerazione. Attraverso un'analisi dettagliata, individuiamo che queste discrepanze derivano principalmente da due aspetti: (1) Disprezzo del Passaggio Precedente, dove l'addestramento ignora l'effetto dell'uso della cache nei passaggi temporali precedenti, e (2) Discrepanza dell'Obiettivo, dove l'obiettivo dell'addestramento (allineare il rumore previsto in ciascun passaggio temporale) si discosta dall'obiettivo dell'inferenza (generare un'immagine di alta qualità). Per attenuare queste discrepanze, proponiamo HarmoniCa, un nuovo metodo che armonizza l'addestramento e l'inferenza con un nuovo framework di caching basato sull'apprendimento costruito su Addestramento Step-Wise per la Denoising (SDT) e Obiettivo Guidato da Proxy dell'Errore dell'Immagine (IEPO). Rispetto al paradigma di addestramento tradizionale, il nuovo SDT proposto mantiene la continuità del processo di denoising, consentendo al modello di sfruttare le informazioni dai passaggi temporali precedenti durante l'addestramento, in modo simile a come opera durante l'inferenza. Inoltre, progettiamo IEPO, che integra un efficace meccanismo di proxy per approssimare l'errore finale dell'immagine causato dal riutilizzo della caratteristica memorizzata. Pertanto, IEPO aiuta a bilanciare la qualità finale dell'immagine e l'utilizzo della cache, risolvendo il problema dell'addestramento che considera solo l'impatto dell'uso della cache sull'output previsto in ciascun passaggio temporale.

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.