Diffusion-RWKV: Scaling RWKV-Like Architectures for Diffusion Models
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.
6614b5c67018e413bf5718d7
Zhengcong Fei
6614b5c67018e413bf5718d8
Mingyuan Fan
6614b5c67018e413bf5718d9
Changqian Yu
6614b5c67018e413bf5718da
Debang Li
6614b5c67018e413bf5718db
Junshi Huang
Diffusion-RWKV: Scaling RWKV-Like Architectures for Diffusion Models
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.
Diffusion-RWKV, an architecture adapted from RWKV for image generation, efficiently handles high-resolution images with reduced computational complexity compared to CNN or Transformer-based models.
RWKV diffusion model patchnified inputs spatial aggregation complexity FID IS total computation FLOP usage