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

초록

확산 트랜스포머(Diffusion Transformers, 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.