Diffusion-RWKV: Escalando Arquitecturas Similares a RWKV para Modelos de Difusión

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

Resumen

Los Transformers han catalizado avances en los campos de visión por computadora y procesamiento del lenguaje natural (NLP). Sin embargo, su considerable complejidad computacional impone limitaciones para su aplicación en tareas de contexto largo, como la generación de imágenes de alta resolución. Este artículo presenta una serie de arquitecturas adaptadas del modelo RWKV utilizado en NLP, con modificaciones necesarias ajustadas para modelos de difusión aplicados a tareas de generación de imágenes, denominadas Diffusion-RWKV. Similar a los modelos de difusión con Transformers, nuestro modelo está diseñado para manejar eficientemente entradas divididas en secuencias con condiciones adicionales, al mismo tiempo que escala de manera efectiva, acomodando tanto parámetros a gran escala como conjuntos de datos extensos. Su ventaja distintiva radica en su reducida complejidad de agregación espacial, lo que lo hace excepcionalmente hábil para procesar imágenes de alta resolución, eliminando así la necesidad de operaciones de ventaneo o agrupamiento en caché. Los resultados experimentales en tareas de generación de imágenes tanto condicionadas como no condicionadas demuestran que Diffusion-RWKV logra un rendimiento igual o superior a los modelos de difusión basados en CNN o Transformers en métricas como FID e IS, mientras reduce significativamente el uso total de FLOPs computacionales.

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.