Scaling Rectified Flow Transformers for High-Resolution Image Synthesis

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.

한 줄 요약

고해상도 텍스트-이미지 생성을 위해 Rectified Flow Transformer를 확장한 연구.

핵심 기여도

Rectified Flow 모델의 노이즈 샘플링 기법을 퍼셉션 관련 스케일로 재가중치하여 성능 향상 (예: `rf/lognorm(0.00, 1.00)`).
텍스트와 이미지 토큰 간 양방향 정보 흐름을 지원하는 새로운 트랜스포머 아키텍처 `MM-DiT` 제안.
8B 파라미터, 5×10²² FLOPs 규모의 확장 실험 수행. `d=38` 모델은 Parti-prompts 벤치마크에서 인간 선호도 평가에서 SOTA 모델을 상회.

핵심 아이디어

Rectified Flow는 데이터와 노이즈를 직선으로 연결하는 생성 모델로 이론적으로 우수하지만, 실제 적용에서는 확산 모델에 비해 성능이 뒤처져 있었다. 본 연구는 이 문제를 해결하기 위해 노이즈 스케일 샘플링 기법을 퍼셉션에 집중적으로 재가중치하여 훈련 과정을 개선했다. 특히, `rf/lognorm(0.00, 1.00)`과 같은 로그정규 분포 기반의 샘플링 방식이 기존의 균일 샘플링(`rf`)보다 훨씬 높은 성능을 보였다.

또한, 기존 텍스트-이미지 생성 모델이 고정된 텍스트 표현을 사용하는 반면, 본 연구는 텍스트와 이미지 토큰 간 양방향 정보 흐름을 가능하게 하는 `MM-DiT`라는 새로운 트랜스포머 아키텍처를 제안했다. 이는 텍스트 이해력, 타이포그래피 정확도, 인간 선호도 평가에서 개선된 성능을 보였다.

기술적 접근법

**Rectified Flow 노이즈 샘플링 개선**: `rf/lognorm(0.00, 1.00)`과 같은 로그정규 분포 기반의 샘플링 기법을 도입하여 훈련 과정에서 퍼셉션 관련 스케일에 집중.
**MM-DiT 아키텍처**: 텍스트와 이미지 토큰에 대해 별도의 가중치를 사용하고, 양방향 정보 흐름을 지원하는 트랜스포머 기반 구조.
**대규모 훈련**: 256×256 해상도 이미지와 16프레임 동영상에 대해 500k 스텝, 140k 스텝 동안 훈련. 최대 8B 파라미터, 5×10²² FLOPs 사용.
**샘플링 스텝 수**: 5단계, 50단계 등 다양한 샘플링 설정에서 성능 평가.

주요 결과

`rf/lognorm(0.00, 1.00)`은 5단계 및 50단계 샘플링 모두에서 평균 순위 2~3위를 기록하며, 기존 LDM-Linear(`eps/linear`)보다 훨씬 높은 성능.
`d=38` 모델은 Parti-prompts 벤치마크에서 `visual aesthetics`, `prompt following`, `typography generation` 카테고리에서 인간 선호도 평가에서 SOTA 모델을 상회.
GenEval에서 최대 모델은 DALL-E 3를 포함한 기존 SOTA 모델을 상회하며, 전체 평가 점수에서 1위.
훈련 손실(`validation loss`)이 낮아질수록 CompBench, GenEval, 인간 평가 점수와 강한 상관관계를 보임.

의의 및 한계

**의의**: Rectified Flow 모델의 노이즈 샘플링 기법을 개선함으로써 기존 확산 모델과의 성능 격차를 줄였으며, `MM-DiT` 아키텍처는 텍스트-이미지 생성의 다중 모달성을 효과적으로 활용.
**한계**: 연구는 대규모 실험을 기반으로 하지만, 실제 산업적 적용 시 필요한 추론 속도, 메모리 효율성, 멀티플랫폼 호환성 등은 명시되지 않음. 또한, 폐쇄형 모델과의 비교는 제한적임.

실용적 활용

고해상도 텍스트-이미지 생성, 디지털 콘텐츠 자동 생성, 광고 및 게임 산업에서의 시각 자료 생성에 활용 가능.
`MM-DiT` 아키텍처는 멀티모달 생성 모델 개발에 기초가 될 수 있으며, 텍스트-비디오 생성 등으로 확장 가능.