HarmoniCa: Armonizando el Entrenamiento y la Inferencia para un Mejor Almacenamiento en Caché de Características en la Aceleración del Transformador de Difusión

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

October 2, 2024

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

cs.AI

Resumen

Los Transformadores de Difusión (DiTs) han ganado prominencia por su escalabilidad sobresaliente y su extraordinario rendimiento en tareas generativas. Sin embargo, sus considerables costos de inferencia obstaculizan su implementación práctica. El mecanismo de caché de características, que implica almacenar y recuperar cálculos redundantes a lo largo de los pasos temporales, promete reducir el tiempo de inferencia por paso en modelos de difusión. La mayoría de los métodos de caché existentes para DiT son diseñados manualmente. Aunque el enfoque basado en aprendizaje intenta optimizar estrategias de forma adaptativa, sufre de discrepancias entre el entrenamiento y la inferencia, lo que afecta tanto al rendimiento como a la relación de aceleración. Tras un análisis detallado, identificamos que estas discrepancias provienen principalmente de dos aspectos: (1) Desconsideración del Paso Temporal Anterior, donde el entrenamiento ignora el efecto del uso de caché en pasos temporales anteriores, y (2) Desajuste de Objetivo, donde el objetivo de entrenamiento (alinear el ruido predicho en cada paso temporal) se desvía del objetivo de inferencia (generar la imagen de alta calidad). Para mitigar estas discrepancias, proponemos HarmoniCa, un método novedoso que armoniza el entrenamiento y la inferencia con un marco de caché basado en aprendizaje construido sobre el Entrenamiento de Desruido Paso a Paso (SDT, por sus siglas en inglés) y el Objetivo Guiado por un Proxy de Error de Imagen (IEPO). En comparación con el paradigma de entrenamiento tradicional, el SDT recientemente propuesto mantiene la continuidad del proceso de desruido, permitiendo que el modelo aproveche la información de pasos temporales anteriores durante el entrenamiento, de manera similar a como opera durante la inferencia. Además, diseñamos IEPO, que integra un mecanismo de proxy eficiente para aproximar el error final de la imagen causado por reutilizar la característica en caché. Por lo tanto, IEPO ayuda a equilibrar la calidad final de la imagen y la utilización de la caché, resolviendo el problema del entrenamiento que solo considera el impacto del uso de caché en la salida predicha en cada paso temporal.

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.