Diffusion-RWKV: Scalabilità delle Architetture Simili a RWKV per Modelli di Diffusione

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

April 6, 2024

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

cs.AI

Abstract

I Transformer hanno catalizzato progressi nei campi della visione artificiale e dell'elaborazione del linguaggio naturale (NLP). Tuttavia, la notevole complessità computazionale pone limitazioni alla loro applicazione in compiti a contesto lungo, come la generazione di immagini ad alta risoluzione. Questo articolo introduce una serie di architetture adattate dal modello RWKV utilizzato in NLP, con le necessarie modifiche appositamente progettate per i modelli di diffusione applicati ai compiti di generazione di immagini, denominati Diffusion-RWKV. Similmente alla diffusione con i Transformer, il nostro modello è progettato per gestire in modo efficiente input suddivisi in patch in una sequenza con condizioni aggiuntive, scalando efficacemente e supportando sia parametri su larga scala che dataset estesi. Il suo vantaggio distintivo si manifesta nella ridotta complessità di aggregazione spaziale, rendendolo eccezionalmente abile nel processare immagini ad alta risoluzione, eliminando così la necessità di operazioni di finestratura o caching di gruppo. I risultati sperimentali su compiti di generazione di immagini condizionata e non condizionata dimostrano che Diffusion-RWKV raggiunge prestazioni pari o superiori ai modelli di diffusione basati su CNN o Transformer nelle metriche FID e IS, riducendo significativamente l'utilizzo totale di FLOP computazionali.

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.