[논문 리뷰] Scaling Rectified Flow Transformers for High-Resolution Image Synthesis

Patrick Esser, Sumith Kulal외 15명|arXiv (Cornell University)|2024. 03. 05.

Computer Graphics and Visualization Techniques인용 수 84

한 줄 요약

논문은 노이즈-스케일 편향 샘플링을 도입하고, 별도의 모달리티 가중치를 가진 다모달 텍스트-이미지 트랜스포머 백본 MM-DiT를 제시하며, 고해상도 이미지 합성을 위한 보정 흐름 모델의 성능을 SOTA 확산 모델과 경쟁하거나 이를 능가하는 결과로 보여주는 스케일링 연구를 수행한다.

ABSTRACT

Diffusion models create data from noise by inverting the forward paths of data towards noise and have emerged as a powerful generative modeling technique for high-dimensional, perceptual data such as images and videos. Rectified flow is a recent generative model formulation that connects data and noise in a straight line. Despite its better theoretical properties and conceptual simplicity, it is not yet decisively established as standard practice. In this work, we improve existing noise sampling techniques for training rectified flow models by biasing them towards perceptually relevant scales. Through a large-scale study, we demonstrate the superior performance of this approach compared to established diffusion formulations for high-resolution text-to-image synthesis. Additionally, we present a novel transformer-based architecture for text-to-image generation that uses separate weights for the two modalities and enables a bidirectional flow of information between image and text tokens, improving text comprehension, typography, and human preference ratings. We demonstrate that this architecture follows predictable scaling trends and correlates lower validation loss to improved text-to-image synthesis as measured by various metrics and human evaluations. Our largest models outperform state-of-the-art models, and we will make our experimental data, code, and model weights publicly available.

연구 동기 및 목표

• 고해상도 이미지 합성을 위한 보정 흐름(formulations)들을 편향된 노이즈 스케일 toward perceptually relevant content로 향하도록 하여 동기를 부여하고 개선한다.
• 텍스트와 이미지 토큰 간 양방향 정보 흐름을 가능하게 하는 확장 가능한 다모달 트랜스포머 백본을 개발하여 텍스트-이미지 생성 성능을 향상한다.
• 데이터셋과 샘플링 설정 전반에서 확산 모델과 보정 흐름 변형을 체계적으로 비교하여 우수한 학습 및 샘플링 전략을 식별한다.
• 제안된 모델의 8B 매개변수까지의 확장 동작을 보여주고 검증 손실과 이미지-텍스트 평가 지표 간의 상관관계를 평가한다.

제안 방법

• 노이즈 스케일 재가중치를 조정하여 training을 지각적으로 관련된 타임스텝 쪽으로 편향시키고, 중간 타임스텝을 강조하는 가중치가 가능한 노이즈 예측 목적(L_w)으로 제시된다.
• RF, EDM, LDM 스타일 스케줄을 포함한 변형들을 비교하고, 로짓-정규(logit-normal), 모드 기반(mode-based), CosMap 타임스텝 분포와 같은 맞춤 SNR 샘플러를 사용한다.
• MM-DiT를 도입하여 이미지와 텍스트 모달리티에 대해 두 개의 별도 가중치 세트를 가진 다모달 확산 백본을 구성하고 교차 주의 및 MLP 처리 중 양방향 상호작용을 가능하게 한다.
• 고해상도 데이터에서 QK 정규화를 사용한 사전 학습 및 미세조정을 통해 주의 집중을 안정화하고 bf16 정밀도에서의 학습을 가능하게 하며, 더 나은 재구성을 위한 잠재 채널 확장(d=16)을 활용한다.
• 향상된 오토인코더(잠재 공간 d=16)와 CogVLM으로 생성된 합성 자막과 원본 자막의 혼합(50/50)을 사용하여 합성 자막을 포함한 대규모 텍스트-이미지 학습 데이터를 보강하고, 확장 가능한 모달리티별 확산 백본을 활용한다.

실험 결과

연구 질문

• RQ1 중간의 지각적으로 관련 스케일로의 타임스텝 샘플링 편향이 고해상도 이미지 합성을 위한 보정 흐름의 성능을 표준 확산 형식과 비교해 향상시키는가?
• RQ2 MM-DiT와 같은 separate image/text token streams를 갖춘 다모달 확산 백본이 전통적 백본(DiT, CrossDiT, UViT)보다 텍스트-이미지 생성에서 우수한가?
• RQ3 보정 흐름 기반 모델에서 스케일링 추세가 어떻게 나타나며, 낮은 검증 손실이 자동 및 인간 평가에서 텍스트-이미지 성능 향상으로 어떻게 이어지는가?
• RQ4 대규모 T2I 모델의 GenEval 스타일 메트릭에 데이터 전처리 및 자막 증강(합성 + 원본 자막)의 효과는 어떤가?
• RQ5 훈련 안정화 기술(QK 정규화, 혼합 정밀도 미세조정, 다양한 종횡비를 위한 위치 인코딩)이 고해상도 미세조정에 필수적인가?

주요 결과

• 중간 타임스텝을 강조하는 노이즈 샘플링 전략(rf/lognorm(0.00, 1.00))은 CLIP 및 FID 지표에서 강력한 성능을 보이고 종종 SOTA 확산 모델을 능가하거나 일치한다.
• 타깃 타임스텝 샘플링이 포함된 보정 흐름 변형은 여러 설정에서 LDM-Linear 및 EDM 베이스라인을 능가하며 특히 낮은 샘플링 단계에서 우수하다.
• 텍스트와 이미지 모달리티에 대해 별도 가중치를 가진 MM-DiT 다모달 백본은 CC12M에서 검증 손실, CLIP, FID에서 vanilla DiT, CrossDiT, UViT를 크게 능가한다.
• 오토인코더의 잠재 채널을 d=16으로 증가시키면 재구성 지표가 향상되고 더 나은 확장 가능성을 지원한다; 더 높은 용량이 더 나은 이미지 품질과 상관관계가 있다.
• 합성(CogVLM 생성) 자막과 원본 자막을 50/50로 혼합하면 GenEval 점수가 향상되어 합성 자막이 학습 데이터를 효과적으로 보강할 수 있음을 시사한다.
• QK 정규화를 통한 훈련 안정화로 bf16 혼합 정밀도에서 안정적인 미세조정이 가능해 고해상도 스케일링 및 더 나은 주의 동작을 촉진한다.
• 8B 매개변수까지의 확장 실험에서 검증 손실이 낮을수록 자동 및 인간 평가에서 텍스트-이미지 성능이 향상된다.