Diffusion-RWKV: Scaling RWKV-Like Architectures for Diffusion Models
Authors: name: Zhengcong Fei; name: Mingyuan Fan; name: Changqian Yu; name: Debang Li; name: Junshi Huang
Published: 2024-04-06 02:54:35

Abstract: 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.