[논문 리뷰] Diffusion-RWKV: Scaling RWKV-Like Architectures for Diffusion Models

Zhengcong Fei, Mingyuan Fan외 3명|arXiv (Cornell University)|2024. 04. 06.

Simulation Techniques and Applications인용 수 5

한 줄 요약

Diffusion-RWKV는 RWKV 백본을 확산 기반 이미지 합성을 위해 적응시키며, Transformer 기반 확산 모델에 비해 선형 시간 복잡도와 더 낮은 FLOPs로 경쟁력 있는 이미지 품질을 달성합니다.

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.

연구 동기 및 목표

• RWKV 아키텍처를 확산 기반 이미지 생성 작업에 맞게 적응하는 것을 탐구한다.
• 안정성과 확장성을 보장하기 위한 conditioning, skip 연결, 그리고 모델 확장성에 대해 조사한다.
• 여러 데이터셋에서 픽셀 및 잠재 공간 표현에 대한 경험적 기준선을 제공한다.
• 품질과 효율성 측면에서 CNN/Transformer 확산 베이스라인과의 성능을 비교한다.

제안 방법

• Bi-RWKV를 백본으로 사용하여 패치화된 이미지 토큰을 양방향으로 선형 복잡도 방식으로 처리한다.
• 표준 self-attention을 양방향 RWKV 기반 메커니즘으로 교체하고, 사방향 공간 시프트와 전역 선형 주의(global linear attention)를 포함한다.
• 타임스텝 정보 및 클래스 정보를 위해 in-context 토큰, adaLN, 또는 adaLN-Zero를 통한 conditioning을 도입한다.
• 패치화하고 임베드한 이미지 토큰에 위치 임베딩을 더해 토큰 시퀀스를 구성한다.
• 얕은 브랜치와 깊은 브랜치 상태를 선형 프로젝션 전에 연결(concatenate)하는 skip-connection 프레임워크를 적용한다.
• 최종 Bi-RWKV 출력을 선형 디코더로 디코딩하여 DDPM 기반 샘플링을 위한 노이즈와 대각 분산을 예측한다.

Figure 2 : Overall framework of diffusion models with RWKV-like architectures. (a) The Diffusion-RWKV architecture comprises $L$ identical Bi-RWKV layers, a patch embedding, and a projection layer. A skip connection is established between shallow and deep stacked Bi-RWKV layers for information flow.

실험 결과

연구 질문

• RQ1 RWKV-like 백본으로 구축된 확산 모델이 Transformer 기반 확산 모델과 비교해 고해상도 이미지 생성 작업에서 어떤 성능을 보이나요?
• RQ2 어떤 아키텍처 선택(패치 크기, skip 연결, conditioning)이 품질, 속도, 확장성의 균형을 가장 효과적으로 이루나요?
• RQ3 다양한 해상도에서 낮은 FLOPs와 메모리 사용으로 경쟁력 있는 FID/IS를 Diffusion-RWKV가 달성할 수 있나요?
• RQ4 CIFAR-10, CelebA, ImageNet 같은 데이터셋에서 깊이/너비의 모델 확장이 성능과 효율성에 어떤 영향을 미치나요?

주요 결과

• Diffusion-RWKV는 유사한 학습 설정에서 CNN/Transformer 확산 모델과 동등하거나 더 나은 FID 결과를 달성합니다.
• 더 작은 패치 크기와 긴 skip 연결 연결이 생성 품질과 학습 효율을 개선합니다.
• AdaLN-Zero conditioning은 in-context conditioning에 비해 더 우수한 FID 성능과 효율성을 제공합니다.
• 더 큰 Bi-RWKV 모델은 증가하는 FLOPs와 함께 더 나은 FID를 제공하여 DiT 베이스라인과 유사한 확장 가능한 개선을 보여줍니다.
• ImageNet 256x256에서 DRWKV-H/2는 일부 최첨단 모델에 비해 총 FLOPs가 낮으면서도 경쟁력 있는 FID를 달성합니다.
• 512x512에서 DRWKV-H/2는 최상위 방법에 근접하면서도 계산 부담을 줄여 여전히 경쟁력을 유지합니다.

Figure 3 : Ablation experiments and model analysis for different designs with DRWKV-S/2 model on the CIFAR10 dataset. We report FID metrics on 10K generated samples every 50K steps. We can find that: (a) Patch size. A smaller patch size can improve the image generation performance. (b) Skip operatio