Diffusion-RWKV: Schaalbaarheid van RWKV-achtige Architecturen voor Diffusiemodellen

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

Samenvatting

Transformers hebben vooruitgang in de computer vision en natuurlijke taalverwerking (NLP) velden versneld. Echter, de aanzienlijke rekencomplexiteit stelt beperkingen aan hun toepassing in taken met lange context, zoals het genereren van hoge-resolutie afbeeldingen. Dit artikel introduceert een reeks architecturen die zijn aangepast van het RWKV-model dat in NLP wordt gebruikt, met de nodige aanpassingen specifiek voor diffusiemodellen die worden toegepast op afbeeldingsgeneratietaken, aangeduid als Diffusion-RWKV. Net als bij diffusie met Transformers, is ons model ontworpen om efficiënt om te gaan met gepatchte invoer in een reeks met extra condities, terwijl het ook effectief kan opschalen, zowel voor grote-schaal parameters als uitgebreide datasets. Het onderscheidende voordeel komt tot uiting in de verminderde complexiteit van ruimtelijke aggregatie, waardoor het bijzonder geschikt is voor het verwerken van hoge-resolutie afbeeldingen, waardoor de noodzaak voor windowing of groepscache-operaties wordt geëlimineerd. Experimentele resultaten voor zowel conditionele als niet-conditionele afbeeldingsgeneratietaken tonen aan dat Diffusion-RWKV prestaties levert die gelijk zijn aan of beter zijn dan bestaande CNN- of Transformer-gebaseerde diffusiemodellen in FID- en IS-metrics, terwijl het totale reken-FLOP-gebruik aanzienlijk wordt verminderd.

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.