Diffusion-RWKV : Mise à l'échelle d'architectures de type RWKV pour les modèles de diffusion

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

April 6, 2024

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

cs.AI

Résumé

Les Transformers ont catalysé des avancées dans les domaines de la vision par ordinateur et du traitement du langage naturel (NLP). Cependant, leur complexité computationnelle substantielle pose des limites à leur application dans des tâches à contexte long, telles que la génération d'images haute résolution. Cet article présente une série d'architectures adaptées du modèle RWKV utilisé en NLP, avec les modifications nécessaires pour les modèles de diffusion appliqués à la génération d'images, appelées Diffusion-RWKV. Similairement à la diffusion avec les Transformers, notre modèle est conçu pour gérer efficacement des entrées découpées en séquences avec des conditions supplémentaires, tout en étant capable de s'adapter à des paramètres à grande échelle et à des ensembles de données étendus. Son avantage distinctif réside dans sa complexité réduite d'agrégation spatiale, le rendant particulièrement apte à traiter des images haute résolution, éliminant ainsi la nécessité d'opérations de fenêtrage ou de mise en cache par groupe. Les résultats expérimentaux sur des tâches de génération d'images conditionnelles et non conditionnelles démontrent que Diffusion-RWKV atteint des performances équivalentes ou supérieures aux modèles de diffusion basés sur les CNN ou les Transformers en termes de métriques FID et IS, tout en réduisant significativement l'utilisation totale des FLOPs de calcul.

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.