Diffusion-RWKV: Escalonando Arquiteturas do Tipo RWKV para Modelos de Difusão

Diffusion-RWKV: Scaling RWKV-Like Architectures for Diffusion Models

April 6, 2024

Autores: Zhengcong Fei, Mingyuan Fan, Changqian Yu, Debang Li, Junshi Huang

cs.AI

Resumo

Os Transformers catalisaram avanços nas áreas de visão computacional e processamento de linguagem natural (PLN). No entanto, a complexidade computacional substancial impõe limitações para sua aplicação em tarefas de contexto longo, como a geração de imagens de alta resolução. Este artigo introduz uma série de arquiteturas adaptadas do modelo RWKV utilizado em PLN, com as modificações necessárias ajustadas para modelos de difusão aplicados a tarefas de geração de imagens, denominadas Diffusion-RWKV. Semelhante à difusão com Transformers, nosso modelo foi projetado para lidar eficientemente com entradas segmentadas em sequência com condições adicionais, além de escalar de forma eficaz, acomodando tanto parâmetros em grande escala quanto conjuntos de dados extensos. Sua vantagem distintiva se manifesta na complexidade reduzida de agregação espacial, tornando-o excepcionalmente hábil no processamento de imagens de alta resolução, eliminando assim a necessidade de operações de janelamento ou cache em grupo. Resultados experimentais em tarefas de geração de imagens condicionais e não condicionais demonstram que o Diffusion-RWKV alcança desempenho equivalente ou superior aos modelos de difusão baseados em CNN ou Transformers nas métricas FID e IS, enquanto reduz significativamente o uso total de FLOPs computacionais.

English

Transformers have catalyzed advancements in computer vision and natural language processing (NLP) fields. However, substantial computational complexity poses limitations for their application in long-context tasks, such as high-resolution image generation. This paper introduces a series of architectures adapted from the RWKV model used in the NLP, with requisite modifications tailored for diffusion model applied to image generation tasks, referred to as Diffusion-RWKV. Similar to the diffusion with Transformers, our model is designed to efficiently handle patchnified inputs in a sequence with extra conditions, while also scaling up effectively, accommodating both large-scale parameters and extensive datasets. Its distinctive advantage manifests in its reduced spatial aggregation complexity, rendering it exceptionally adept at processing high-resolution images, thereby eliminating the necessity for windowing or group cached operations. Experimental results on both condition and unconditional image generation tasks demonstrate that Diffison-RWKV achieves performance on par with or surpasses existing CNN or Transformer-based diffusion models in FID and IS metrics while significantly reducing total computation FLOP usage.