기술 요약: 경계 조건을 통한 Rectified Flow 개선

문제 정의

Rectified Flow (RF), 또는 Flow Matching으로도 알려진 이 방식은 노이즈 분포를 타겟 데이터 분포로 변환하기 위해 상미분 방정식(ODE)을 통해 속도장(velocity field)을 학습하는 강력한 생성 모델링 프레임워크로 부상했습니다. 텍스트-투-이미지 및 비디오 생성 작업에서의 성공에도 불구하고, 저자들은 표준 RF 구현체에서 중요한 한계를 식별했습니다: 학습된 속도장이 이론적으로 도출된 경계 조건을 충족하지 못하는 경우가 많다는 점입니다.

구체적으로, 노이즈($t=0$)와 데이터($t=1$) 사이의 선형 보간 하에서, 최적의 속도장은 다음을 만족해야 합니다:

1. 좌측 경계 ($t=0$): $\mathbf{v}(\mathbf{x}, 0) = \mathbb{E}[\mathbf{X}_1] - \mathbf{x}$
2. 우측 경계 ($t=1$): $\mathbf{v}(\mathbf{x}, 1) = \mathbf{x}$

표준 신경망 파라미터화 방식은 이러한 제약 조건을 내재적으로 강제하지 않습니다. 논문에서 시각화된 바와 같이, $t=1$에서의 예측 속도는 기대되는 항등 사상(identity map)에서 벗어나는 경우가 빈번합니다. 이러한 위반은 확률적 샘플링(예: Langevin dynamics) 중에 특히 치명적입니다. 왜냐하면 스코어 함수 추정(Tweedie의 공식을 통한)은 $(1-t)$로 나누는 과정을 포함하는데, $t=1$ 근처에서 속도가 부정확하면 스코어 추정치가 발산하기 때문입니다. 이는 샘플링을 불안정하게 만들어, 특히 더 큰 스텝 사이즈나 노이즈 스케일을 사용할 때 이미지가 과도하게 매끄러워지거나(over-smoothed) "만화 같은" 아티팩트를 생성하는 결과를 초축래합니다.

방법론

이를 해결하기 위해, 저자들은 **경계 강제 Rectified Flow 모델(Boundary-enforced Rectified Flow Model)**을 제안합니다. 이 접근 방식은 두 가지 구별되는 변형을 통해 속도장의 파라미터화에 경계 제약 조건을 명시적으로 통합합니다:

1. 마스크 기반 경계 RF 모델 (Mask-based Boundary RF Model)

이 변형은 속도장을 경계 항과 학습 가능한 신경망 출력의 선형 결합으로 구조화하여 두 경계 조건($t=0$ 및 $t=1$)을 모두 강제합니다:
$$\mathbf{v}(\mathbf{x}, t) = g(t) \cdot (C - \mathbf{x}) + f(t) \cdot \mathbf{x} + h(t) \cdot m_\theta(\mathbf{x}, t)$$
여기서:

• $m_\theta(\mathbf{x}, t)$는 표준 신경망(예: U-Net 또는 DiT)입니다.
• $C = \mathbb{E}[\mathbf{X}_1]$입니다.
• $f(t), g(t), h(t)$는 특정 경계에서 소멸하거나 활성화되도록 설계된 스칼라 함수입니다 (예: $g(1)=f(0)=h(0)=h(1)=0$ 및 $g(0)=f(1)=1$).
• 저자들은 "표준 코사인(Standard Cosine)" 선택지를 제안합니다: $g(t) = \cos(\frac{\pi}{2}t)$, $f(t) = \sin(\frac{\pi}{2}t)$, $h(t) = \sin(\pi t)$.

이 구조는 신경망의 출력과 관계없이 설계에 의해 경계 조건이 충족됨을 보장합니다.

2. 차분 기반 경계 RF 모델 (Subtraction-based Boundary RF Model)

$t=1$ 조건을 강제하는 것만으로도 상당한 경험적 이득을 얻을 수 있다는 관찰에 착안하여, 저자들은 더 단순한 단일 경계 변형을 제안합니다:
$$\mathbf{v}(\mathbf{x}, t) = \mathbf{x} + m_\theta(\mathbf{x}, t) - m_\theta(\mathbf{x}, 1)$$
이 공식은 $t=1$에서 $m_\theta(\mathbf{x}, 1)$ 항이 상쇄되기 때문에 $\mathbf{v}(\mathbf{x}, 1) = \mathbf{x}$를 본질적으로 만족합니다. 이 방식은 설계의 단순함과 유연성을 제공하지만, 매 평가 시마다 $m_\theta(\mathbf{x}, 1)$을 계산하기 위한 추가적인 순전파(forward pass)가 필요하므로 계산 비용이 발생합니다.

확률적 샘플링에 미치는 영향

$\mathbf{v}(\mathbf{x}, 1) = \mathbf{x}$를 강제함으로써, 경계 RF 모델은 Tweedie 공식의 분자($t\mathbf{v}(\mathbf{x}, t) - \mathbf{x}$)가 $t \to 1$로 갈 때 0에 접근하도록 보장합니다. 이는 스코어 함수 $\nabla \log \rho_t$가 발산하는 것을 방지하여, SDE 샘플러를 안정화하고 더 큰 노이즈 스케일이나 스텝 사이즈에서도 견고한 생성을 가능하게 합니다.

주요 기여

1. 한계 식별: 본 논문은 제약되지 않은 속도장이 Rectified Flow의 이론적 경계 조건을 위반하여, 부정확한 ODE 궤적과 불안정한 확률적 샘플링을 초래한다는 점을 강조합니다.
2. 새로운 파라미터화: 이러한 조건들을 최소한의 코드 수정으로 강제할 수 있는 두 가지 실용적인 변형(마스크 기반 및 차분 기반)을 제공하는 Boundary RF 모델을 도입했습니다.
3. 스코어 함수의 안정화: 경계 강제가 $t=1$ 근처에서 스코어 함수를 직접적으로 안정화하여, 바닐라 RF에서 발생하는 발산 문제를 완화함을 입증했습니다.
4. 경험적 검증: 다양한 데이터셋과 모델 규모에 걸쳐 생성 품질과 샘플링 안정성이 일관되게 향상됨을 보여주는 포괄적인 실험을 수행했습니다.

실험 결과

저자들은 Diffusion Transformer (DiT) 아키텍처를 사용하여 CIFAR-10 및 ImageNet (256$\times$256 및 512$\times$512)에서 모델을 평가했습니다.

• 정량적 개선:

  • ImageNet 256$\times$256에서, 차분 기반 경계 RF 모델은 바닐라 RF 베이스라인(각각 6.87 및 6.79)에 비해 ODE 샘플링 시 6.32, SDE 샘플링 시 6.38의 FID를 달성했습니다.
  • 마스크 기반 모델은 ODE 6.63, SDE 6.18의 FID를 달축했습니다.
  • 두 데이터셋 모두에서 Inception Score (IS), Precision, Recall에서도 개선이 관찰되었습니다.
  • 차분 기반 변형이 일관되게 가장 높은 IS 점수를 기록했습니다.

• 샘플링 견고성:

  • 다양한 노이즈 스케일($c$)을 사용하는 Overshooting Sampler 실험에서, 경계 RF 모델은 바닐라 RF에 비해 확률성이 증가함에 따라 훨씬 더 안정적인 성능을 보였으며 샘플 품질의 저하가 적었습니다.
  • 정성적 결과는 경계 RF 모델이 바닐라 RF가 과도하게 매끄러워지는 경향이 있는 확률적 샘플링 설정에서 더 선명한 이미지와 적은 아티팩트, 그리고 더 나은 객체 충실도를 생성함을 보여주었습니다.

• 확장성:
  경계 RF 모델의 이점은 더 큰 모델(DiT-L, DiT-XL) 및 더 높은 해상도(512$\times$512)로 확장할 때도 일관되게 나타났습니다.

의의 및 주장

본 논문은 Boundary RF 모델이 Rectified Flow의 근본적인 한계에 대한 간단하고 즉시 구현 가능한 솔루션을 제공한다고 주장합니다. 경계 조건을 명시적으로 강제함으로써 이 방법은 다음을 달성합니다:

1. 생성 품질 향상: FID 및 기타 지표에서 일관되게 바닐라 RF보다 우수한 성능을 보입니다.
2. 안정적인 확률적 샘플링 가능: 종단 시간 근처에서 스코어 함수의 발산을 해결하여, 더 견고한 샘플링 전략을 가능하게 합니다.
3. 쉬운 적용성: 최소한의 아키텍처 변경만으로 기존 RF 모델(U-Net 또는 DiT 사용 모델)에 통합될 수 있습니다.

저자들은 이 연구를 Rectified Flow의 실무 배포, 특히 고품질의 확률적 샘플링이 필요한 시나리오를 위한 필수적인 개선 작업으로 규정합니다. 또한, 연구의 초점은 RF에 있지만, 경계 조건을 강제하는 핵심 개념은 약간의 수정을 거쳐 확산 모델(diffusion models)에도 확장될 수 있음을 시사합니다. 향후 연구로는 적응형 경계 강제 및 텍스트-투-비디오 생성에 대한 적용을 탐구하는 것이 제안되었습니다.