Diffusion-RWKV: Skalierung von RWKV-ähnlichen Architekturen für Diffusionsmodelle

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

April 6, 2024

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

cs.AI

Zusammenfassung

Transformer haben Fortschritte in den Bereichen Computer Vision und Natural Language Processing (NLP) vorangetrieben. Allerdings stellt die erhebliche Rechenkomplexität Einschränkungen für ihren Einsatz in Aufgaben mit langem Kontext dar, wie z.B. der Erzeugung hochauflösender Bilder. Dieses Paper stellt eine Reihe von Architekturen vor, die aus dem RWKV-Modell aus dem NLP adaptiert wurden, mit erforderlichen Anpassungen, die speziell für das Diffusion-Modell für Bildgenerierungsaufgaben entwickelt wurden, bekannt als Diffusion-RWKV. Ähnlich wie bei der Diffusion mit Transformers ist unser Modell darauf ausgelegt, Patch-Inputs effizient in einer Sequenz mit zusätzlichen Bedingungen zu verarbeiten, während es auch effektiv skaliert und sowohl große Parameter als auch umfangreiche Datensätze bewältigt. Sein herausragender Vorteil zeigt sich in seiner reduzierten räumlichen Aggregationskomplexität, was es außergewöhnlich geschickt bei der Verarbeitung hochauflösender Bilder macht und somit die Notwendigkeit für Fensterung oder gruppierte Zwischenspeicheroperationen beseitigt. Experimentelle Ergebnisse sowohl bei bedingten als auch unbedingten Bildgenerierungsaufgaben zeigen, dass Diffusion-RWKV eine Leistung erbringt, die mit oder über bestehenden CNN- oder Transformer-basierten Diffusionsmodellen in FID- und IS-Metriken liegt, während der Gesamtberechnungsaufwand signifikant reduziert wird.

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.