DiTFastAttn: 확산 트랜스포머 모델을 위한 어텐션 압축 기술

DiTFastAttn: Attention Compression for Diffusion Transformer Models

June 12, 2024

저자: Zhihang Yuan, Pu Lu, Hanling Zhang, Xuefei Ning, Linfeng Zhang, Tianchen Zhao, Shengen Yan, Guohao Dai, Yu Wang

cs.AI

초록

Diffusion Transformers(DiT)는 이미지 및 비디오 생성에서 뛰어난 성능을 보이지만, self-attention의 이차 복잡도로 인해 계산적 문제에 직면합니다. 우리는 DiT의 계산 병목 현상을 완화하기 위한 새로운 사후 훈련 압축 방법인 DiTFastAttn을 제안합니다. 우리는 DiT 추론 과정에서 주의 계산의 세 가지 주요 중복성을 식별했습니다: 1. 많은 주의 헤드가 지역 정보에 집중하는 공간적 중복성, 2. 인접한 단계의 주의 출력 간 높은 유사성을 보이는 시간적 중복성, 3. 조건부 및 무조건부 추론 간 상당한 유사성을 보이는 조건적 중복성. 이러한 중복성을 해결하기 위해 우리는 세 가지 기술을 제안합니다: 1. 공간적 중복성을 줄이기 위한 Window Attention with Residual Caching, 2. 단계 간 유사성을 활용하는 Temporal Similarity Reduction, 3. 조건부 생성 중 중복 계산을 건너뛰는 Conditional Redundancy Elimination. DiTFastAttn의 효과를 입증하기 위해, 우리는 이를 이미지 생성 작업을 위한 DiT와 PixArt-Sigma, 그리고 비디오 생성 작업을 위한 OpenSora에 적용했습니다. 평가 결과, 이미지 생성에서 우리의 방법은 최대 88%의 FLOPs를 줄이고 고해상도 생성에서 최대 1.6배의 속도 향상을 달성했습니다.

English

Diffusion Transformers (DiT) excel at image and video generation but face computational challenges due to self-attention's quadratic complexity. We propose DiTFastAttn, a novel post-training compression method to alleviate DiT's computational bottleneck. We identify three key redundancies in the attention computation during DiT inference: 1. spatial redundancy, where many attention heads focus on local information; 2. temporal redundancy, with high similarity between neighboring steps' attention outputs; 3. conditional redundancy, where conditional and unconditional inferences exhibit significant similarity. To tackle these redundancies, we propose three techniques: 1. Window Attention with Residual Caching to reduce spatial redundancy; 2. Temporal Similarity Reduction to exploit the similarity between steps; 3. Conditional Redundancy Elimination to skip redundant computations during conditional generation. To demonstrate the effectiveness of DiTFastAttn, we apply it to DiT, PixArt-Sigma for image generation tasks, and OpenSora for video generation tasks. Evaluation results show that for image generation, our method reduces up to 88\% of the FLOPs and achieves up to 1.6x speedup at high resolution generation.