HarmoniCa: Harmonizando Treinamento e Inferência para um Melhor Cache de Recursos em Aceleração de Transformadores de Difusão

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

Resumo

Os Transformadores de Difusão (DiTs) ganharam destaque pela escalabilidade excepcional e desempenho extraordinário em tarefas generativas. No entanto, seus consideráveis custos de inferência impedem a implantação prática. O mecanismo de cache de características, que envolve armazenar e recuperar cálculos redundantes ao longo dos passos de tempo, promete reduzir o tempo de inferência por passo nos modelos de difusão. A maioria dos métodos de cache existentes para DiT são projetados manualmente. Embora a abordagem baseada em aprendizado tente otimizar estratégias de forma adaptativa, ela sofre de discrepâncias entre treinamento e inferência, o que prejudica tanto o desempenho quanto a taxa de aceleração. Após uma análise detalhada, identificamos que essas discrepâncias derivam principalmente de dois aspectos: (1) Desconsideração do Passo de Tempo Anterior, onde o treinamento ignora o efeito do uso de cache nos passos de tempo anteriores, e (2) Desajuste de Objetivo, onde o alvo de treinamento (alinhar o ruído previsto em cada passo de tempo) se desvia do objetivo da inferência (gerar a imagem de alta qualidade). Para mitigar essas discrepâncias, propomos o HarmoniCa, um método inovador que harmoniza treinamento e inferência com um novo framework de Cache baseado em aprendizado construído sobre Treinamento de Desnublamento Passo a Passo (SDT) e Objetivo Guiado por Proxy de Erro de Imagem (IEPO). Comparado ao paradigma de treinamento tradicional, o recém-proposto SDT mantém a continuidade do processo de desnublamento, permitindo que o modelo aproveite informações de passos de tempo anteriores durante o treinamento, de forma semelhante ao que ocorre durante a inferência. Além disso, projetamos o IEPO, que integra um mecanismo de proxy eficiente para aproximar o erro final da imagem causado pelo reuso da característica em cache. Portanto, o IEPO ajuda a equilibrar a qualidade final da imagem e a utilização do cache, resolvendo o problema do treinamento que considera apenas o impacto do uso de cache na saída prevista em cada passo de tempo.

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.