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

Аннотация

Диффузионные трансформеры (DiT) отличаются в выдаче изображений и видеогенерации, но сталкиваются с вычислительными проблемами из-за квадратичной сложности само-внимания. Мы предлагаем DiTFastAttn, новый метод сжатия после обучения для смягчения вычислительной узкой точки DiT. Мы выявляем три ключевые избыточности в вычислении внимания во время вывода DiT: 1. пространственная избыточность, где многие головы внимания сосредотачиваются на локальной информации; 2. временная избыточность, с высокой схожестью между выходами внимания соседних шагов; 3. условная избыточность, где условные и безусловные выводы проявляют значительную схожесть. Для решения этих избыточностей мы предлагаем три техники: 1. Внимание окна с остаточным кэшированием для сокращения пространственной избыточности; 2. Снижение временной схожести для использования сходства между шагами; 3. Устранение условной избыточности для пропуска избыточных вычислений во время условной генерации. Чтобы продемонстрировать эффективность DiTFastAttn, мы применяем его к DiT, PixArt-Sigma для задач генерации изображений и OpenSora для задач генерации видео. Результаты оценки показывают, что для генерации изображений наш метод сокращает до 88\% операций с плавающей запятой и достигает до 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.