[Gen AI/Flow] Scaling Rectified Flow Transformers for High-Resolution Image Synthesis

기존 디퓨전 모델과 그 한계

• Diffusion 모델 : 역방향 SDE (Stochastic Differential Equation) 기반으로 노이즈에서 이미지를 복원하는 과정을 통해 새로운 이미지를 생성함. (먼저 이미지에 노이즈를 점점 추가해서 완전히 망가뜨림 → 그 다음 이 망가진 걸 거꾸로 되돌리는 방법을 학습함)

이때 reverse 과정은 노이즈에서 점진적으로 denoise해야 하므로 수많은 time step이 필요함.

이 경로는 곡선처럼 휘어 있어서 역으로 학습할 때 누적 오차가 큼.

---

논문의 핵심 목표

1. 더 빠르고 효율적인 이미지 생성 모델.
2. 고해상도 이미지도 잘 만들 수 있도록.
3. 기존 diffusion 모델보다 더 단순한 구조로 높은 성능.

---

Rectified Flow

• 일반적인 Diffusion 모델은 ‘노이즈 → 깨끗한 이미지’로 점점 복원해가며 학습함.
• Rectified Flow: 이미지의 시작(x0)과 끝(노이즈)을 직선으로 연결하는 경로를 따라 학습함. 그래서 더 간단하고 효율적.

⇒ 학습하려는 경로는 어떤 속도로 움직이는가. 모델이 시점 t에서 어디로 이동해야 하는지

초기 → 목표를 잇는 직선 경로 상에서의 '속도 벡터'를 학습.

 

⇒ 선형 스케일링을 통해 시간 t에 따른 중간 latent 표현을 만드는 공식

 

원본 이미지 x0에 일정한 비율로 노이즈를 섞어서 중간 이미지 zt를 만든다.

 

시간 t가 작을수록 원본에 가까운 이미지,

t가 클수록 노이즈 쪽에 가까운 이미지

 

 

손실함수

⇒ 모델이 예측한 흐름이 실제로 z가 시간이 지나며 어떤 방향으로 이동해야 하는지를 잘 추측하는지를 평가

⇒ 이미지 생성 경로의 중간 지점을 구성하는 방법을 표현

⇒ 단순히 균일 분포에서 t를 뽑으면 모델이 모든 시간에 대해 동등하게 학습함.

하지만 어떤 시간 구간 (t가 작을 때 or 클 때)이 더 중요할 수 있기 때문에 logit-normal 분포를 통해 중요 구간에 더 많은 학습 자원을 배정할 수 있도록 함. → “어디서 더 많은 학습을 할 것인가”

---

Rectified Flow Transformer (RFT)

이 논문은 기존의 Rectified Flow에 Transformer 구조를 붙여서 새롭게 확장.

1. U-ViT 아키텍처: U-Net 구조(이미지 생성 모델에서 자주 쓰이는 구조)와 Vision Transformer(이미지를 패치로 나눠서 Transformer로 처리)를 합친 형태. 이미지의 저해상도 정보와 고해상도 정보를 동시에 처리할 수 있음.
2. Flow Matching 학습: 일반적인 노이즈 예측 대신 이미지 사이의 ‘경로’ 자체를 예측하도록 학습함. 그래서 더 빠른 샘플링이 가능함.
3. Conditioning: CLIP, T5 같은 텍스트 모델을 통해 문장을 벡터로 바꿔서 사용.

self.transformer(
    hidden_states=latents,                         # [B, N, C] 지금까지 처리한 이미지 상태
    timestep=timestep / 1000,                      # 시간 정보를 모델에 넣음
    guidance=guidance,                             # 텍스트 조건 강도 조절
    pooled_projections=pooled_prompt_embeds,       # CLIP 기반 텍스트 임베딩
    encoder_hidden_states=prompt_embeds,           # T5 기반 텍스트 임베딩 
    txt_ids=text_ids,                              # text token 위치
    img_ids=latent_image_ids,                      # 이미지 위치 정보
    joint_attention_kwargs=...,                    # 조건부 요소
    return_dict=False
)

B: 배치 크기

N: 패치 개수 (예: 이미지가 64x64 패치로 나눠지면 N = 4096)

C: 각 패치의 채널 수

---

성능

논문에서는 1024×1024 해상도의 이미지를 생성할 수 있고, FID 점수도 경쟁 모델보다 낮게 나옴 (더 좋음).

30 steps 이하로 고품질 이미지를 생성함 → 매우 빠름.

---

아키텍처