Диффузия-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

Аннотация

Трансформеры стимулировали прогресс в областях компьютерного зрения и обработки естественного языка (NLP). Однако значительная вычислительная сложность ставит ограничения на их применение в задачах с длинным контекстом, таких как генерация изображений высокого разрешения. В данной статье представлен ряд архитектур, адаптированных из модели RWKV, используемой в NLP, с необходимыми модификациями, нацеленными на модель диффузии, применяемую к задачам генерации изображений, известную как Diffusion-RWKV. Подобно диффузии с Трансформерами, наша модель разработана для эффективной обработки входных данных, разбитых на патчи в последовательности с дополнительными условиями, а также для эффективного масштабирования, способного вмещать как большое количество параметров, так и обширные наборы данных. Ее отличительным преимуществом является снижение сложности пространственной агрегации, что делает ее исключительно способной к обработке изображений высокого разрешения, тем самым устраняя необходимость в оконных или групповых кэшируемых операциях. Экспериментальные результаты как по задачам генерации изображений с условиями, так и без них показывают, что Diffusion-RWKV достигает производительности на уровне или превосходящей существующие модели диффузии на основе CNN или Трансформеров по метрикам FID и IS, при этом значительно сокращая общее использование вычислительных операций 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.