HarmoniCa: Гармонизация обучения и вывода для улучшения кэша признаков в ускорении диффузионного трансформера.

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

October 2, 2024

Авторы: Yushi Huang, Zining Wang, Ruihao Gong, Jing Liu, Xinjie Zhang, Jinyang Guo, Xianglong Liu, Jun Zhang

cs.AI

Аннотация

Трансформеры диффузии (DiTs) стали известны своей выдающейся масштабируемостью и выдающейся производительностью в генеративных задачах. Однако их значительные затраты на вывод затрудняют практическое применение. Механизм кэширования признаков, который включает в себя сохранение и извлечение избыточных вычислений через временные шаги, обещает снизить время вывода на каждом шаге в моделях диффузии. Большинство существующих методов кэширования для DiT разрабатываются вручную. Хотя обучающийся подход пытается оптимизировать стратегии адаптивно, он сталкивается с расхождениями между обучением и выводом, что затрудняет как производительность, так и коэффициент ускорения. После детального анализа мы выявляем, что эти расхождения в основном происходят из двух аспектов: (1) Пренебрежение предыдущим временем, когда обучение игнорирует влияние использования кэша на более ранних временных шагах, и (2) Несоответствие целей, когда цель обучения (выравнивание предсказанного шума на каждом временном шаге) отклоняется от цели вывода (генерация изображения высокого качества). Для устранения этих расхождений мы предлагаем HarmoniCa, новый метод, который гармонизирует обучение и вывод с помощью новой обучающейся кэширующей структуры, построенной на обучении пошагового шумоподавления (SDT) и цели, направленной на ошибку изображения с прокси-управляемым объектом (IEPO). По сравнению с традиционной парадигмой обучения, вновь предложенный SDT поддерживает непрерывность процесса шумоподавления, позволяя модели использовать информацию с предыдущих временных шагов во время обучения, аналогично тому, как она работает во время вывода. Более того, мы разрабатываем IEPO, который интегрирует эффективный механизм прокси для приближения окончательной ошибки изображения, вызванной повторным использованием кэшированного признака. Таким образом, IEPO помогает балансировать качество окончательного изображения и использование кэша, разрешая проблему обучения, которое учитывает только влияние использования кэша на предсказанный вывод на каждом временном шаге.

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.