AsyncDiff: Paralelização de Modelos de Difusão por Desruído Assíncrono

AsyncDiff: Parallelizing Diffusion Models by Asynchronous Denoising

June 11, 2024

Autores: Zigeng Chen, Xinyin Ma, Gongfan Fang, Zhenxiong Tan, Xinchao Wang

cs.AI

Resumo

Os modelos de difusão têm despertado grande interesse da comunidade devido à sua notável capacidade gerativa em diversas aplicações. No entanto, sua natureza típica de desruído sequencial em múltiplos passos resulta em uma alta latência cumulativa, impedindo a possibilidade de computação paralela. Para resolver esse problema, apresentamos o AsyncDiff, um esquema de aceleração universal e plug-and-play que permite o paralelismo de modelos em múltiplos dispositivos. Nossa abordagem divide o complexo modelo de predição de ruído em vários componentes, atribuindo cada um a um dispositivo diferente. Para quebrar a cadeia de dependência entre esses componentes, o método transforma o desruído sequencial convencional em um processo assíncrono, explorando a alta similaridade entre os estados ocultos em passos consecutivos de difusão. Consequentemente, cada componente é capaz de computar em paralelo em dispositivos separados. A estratégia proposta reduz significativamente a latência de inferência, com impacto mínimo na qualidade gerativa. Especificamente, para o Stable Diffusion v2.1, o AsyncDiff alcança um aumento de velocidade de 2,7x com degradação insignificante e um aumento de 4,0x com apenas uma leve redução de 0,38 no CLIP Score, utilizando quatro GPUs NVIDIA A5000. Nossos experimentos também demonstram que o AsyncDiff pode ser facilmente aplicado a modelos de difusão de vídeo, com desempenhos promissores. O código está disponível em https://github.com/czg1225/AsyncDiff.

English

Diffusion models have garnered significant interest from the community for their great generative ability across various applications. However, their typical multi-step sequential-denoising nature gives rise to high cumulative latency, thereby precluding the possibilities of parallel computation. To address this, we introduce AsyncDiff, a universal and plug-and-play acceleration scheme that enables model parallelism across multiple devices. Our approach divides the cumbersome noise prediction model into multiple components, assigning each to a different device. To break the dependency chain between these components, it transforms the conventional sequential denoising into an asynchronous process by exploiting the high similarity between hidden states in consecutive diffusion steps. Consequently, each component is facilitated to compute in parallel on separate devices. The proposed strategy significantly reduces inference latency while minimally impacting the generative quality. Specifically, for the Stable Diffusion v2.1, AsyncDiff achieves a 2.7x speedup with negligible degradation and a 4.0x speedup with only a slight reduction of 0.38 in CLIP Score, on four NVIDIA A5000 GPUs. Our experiments also demonstrate that AsyncDiff can be readily applied to video diffusion models with encouraging performances. The code is available at https://github.com/czg1225/AsyncDiff.