Diffusion-RWKV：拡散モデルのためのRWKV様アーキテクチャのスケーリング

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

April 6, 2024

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

cs.AI

要旨

Transformerは、コンピュータビジョンや自然言語処理（NLP）分野の進歩を促進してきました。しかし、その高い計算複雑性は、高解像度画像生成などの長文脈タスクへの応用に制約を課しています。本論文では、NLPで使用されるRWKVモデルを基に、画像生成タスクに適用する拡散モデル向けに必要な修正を加えた一連のアーキテクチャを紹介します。これをDiffusion-RWKVと呼びます。Transformerを用いた拡散モデルと同様に、本モデルは、追加条件付きのシーケンスにおけるパッチ化された入力を効率的に処理し、大規模なパラメータとデータセットに対応できるように設計されています。その特筆すべき利点は、空間集約の複雑性が低減されており、高解像度画像の処理に特に優れている点です。これにより、ウィンドウ処理やグループキャッシュ操作の必要性がなくなります。条件付きおよび無条件の画像生成タスクにおける実験結果は、Diffusion-RWKVがFIDおよびISメトリクスにおいて、既存のCNNやTransformerベースの拡散モデルと同等またはそれ以上の性能を達成しつつ、総計算FLOP使用量を大幅に削減できることを示しています。

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.