DiTFastAttn: Compresión de atención para modelos de transformadores de difusión

DiTFastAttn: Attention Compression for Diffusion Transformer Models

June 12, 2024

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

cs.AI

Resumen

Los Transformadores de Difusión (DiT) sobresalen en la generación de imágenes y videos, pero enfrentan desafíos computacionales debido a la complejidad cuadrática de la autoatención. Proponemos DiTFastAttn, un novedoso método de compresión post-entrenamiento para aliviar el cuello de botella computacional de los DiT. Identificamos tres redundancias clave en el cálculo de la atención durante la inferencia de los DiT: 1. redundancia espacial, donde muchas cabezas de atención se enfocan en información local; 2. redundancia temporal, con alta similitud entre las salidas de atención de pasos vecinos; 3. redundancia condicional, donde las inferencias condicionales e incondicionales muestran una similitud significativa. Para abordar estas redundancias, proponemos tres técnicas: 1. Atención por Ventanas con Almacenamiento en Caché de Residuos para reducir la redundancia espacial; 2. Reducción de Similitud Temporal para aprovechar la similitud entre pasos; 3. Eliminación de Redundancia Condicional para omitir cálculos redundantes durante la generación condicional. Para demostrar la efectividad de DiTFastAttn, lo aplicamos a DiT, PixArt-Sigma para tareas de generación de imágenes y OpenSora para tareas de generación de videos. Los resultados de evaluación muestran que, para la generación de imágenes, nuestro método reduce hasta un 88\% de las operaciones de punto flotante (FLOPs) y logra una aceleración de hasta 1.6x en la generación de alta resolución.

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.