Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow

• Latent Diffusion Model에서 Sampling 가속화를 논할 때 거의 빠지지 않는 논문.
• 최대한 직관으로 이해할 수 있도록 요약해봄

문제상황

• 굉장히 저명한 문제인 Transport mapping problem을 다룬다.
• 두 분포 $X_0∼π_0, X_1∼π_1$ 가 주어질 때, 두 분포 사이의 transport map $T$를 찾는 문제.
  • 단, $Z_1=T(Z_0)∼π_1$인 ($where Z_0∼π_0$)
• Transport mapping 문제를 잘 풀되 최대한 빠른 속도 안에 푸는 게 목표

가장 빠른 건 직선이동.

• 일단, 우리는 $X_0$로부터 시작해 시간 $t$에 따라 $X_1$ 으로 이동하는 $X_t$ 에 대해 논할 것임.
• 두 분포 사이를 잇는 직선이 있다고 그냥 상상해보자.
• 그 직선의 방향은 $X_1 - X_0$ 방향이고 우리는 이 사이에서 부유 중인 $X_t$를 모델링 하고 싶은 것임.
• 직선을 가정하였으니 $X_t$는 $X_1$과 $X_0$사이를 내분하는 점이라고 할 수 있다.
  • $X_t$가 등속 운동을 한다면, $X_t=tX_1 + (1-t)X_0$ 로 쓸 수 있다.
  • 즉, $dX_t=(X_1-X_0)dt$이다.

Recap Basic.

• 딥러닝의 본질은 결국 $y=f(x)$라는 자연현상을 근사하는 $g(x)$를 유도하는 과정이라고 감히 말하겠다.
• 그럼 우리는 직선 위에서 등속 운동하는 어떤 샘플의 속도를 근사하는 함수도 유도해낼 수 있다.
  • $dZ_t=v_\theta(Z_t,t)$ 라는 미분 방정식이 위의 등속운동에 근사하게끔 학습하면 된다!
  • 이 속도 벡터의 방향이 $X_1-X_0$와 같도록 근사하는 것.
• 즉, 우리는 아래와 같은 식으로 우리의 속도 미분 방정식을 근사시킬 수 있다.
  • $\min_{v_{\theta}} \int_0^1 \mathbb{E}\left[\left\|(X_1 - X_0) - v(X_t, t)\right\|^2\right] dt$ (단, $X_t = tX_1 + (1-t)X_0$)

Full Algorithm

• 요약하자면, Recap Basic에서 얻은 최적화 수식은 Procedure로 명명하고,
• Reflow에서 n번의 Procedure를 걸쳐 (이전 결과가 다음 Procedure의 input이 됨)
• 최종적으론 Distill을 통해 n번의 procedure를 한 번에 학습하는 모델을 만드는 것.

Main Results and Properties

• Marginal preserving property
  • 임의의 시간 $t$에서 $Z_t$의 marginal distribution은 $X_t$의 marginal distribution과 항상 같다.
  • 직선을 가정했으므로, 우리는 샘플의 분포에서 분포로 이동하는 과정을 아래와 같이 정의할 수 있다.
    • $v^X(x,t) = \mathbb{E}[X_1 - X_0 \mid X_t = x]$
    • t시점에 위치 x에 있는 $X_t$가 직선 거리에 있을 기댓값.
  • 위의 수식이 우리가 찾고자 하는 정답이다.
  • 근데, 이것이 우리가 학습하는 모델 $Z_t$는 직선 거리로 이동하도록 학습하였음.
    • $dZ_t = v^X(Z_t, t) \, dt, \quad Z_0 \sim \pi_0$
    • 즉, 위의 정답 분포와 우리가 학습시키는 분포는 같다!
      • (솔직히 말하면 직관으로 이해해서 이렇게 이해하는 것이 맞는지 확신이 없음... ㅠ)
• Reducing transport costs
  • 위의 Marginal preserving property는 어찌보면 당연해 보임.
  • 그러나 이는 얀센부등식을 사용하기 위한 초석이었음
    • 둘의 분포가 같으므로 아래의 등식은 항상 참임
    • $\mathbb{E}\left[\int_0^1 c(v_X(X_t, t)) dt\right] = \mathbb{E}\left[\int_0^1 c(v_X(Z_t, t)) dt\right]$
    • 이를 얀센 부등식을 응용하여 풀면...
  • $\mathbb{E}[c(Z_1 - Z_0)] \leq \mathbb{E}[c(X_1 - X_0)]$ 가 성립한다.
  • 즉, $Z_0$ ~$Z_1$의 convex transport cost가 항상 실제 분포의 이동보다 작다!
    • 말인즉, 항상 거의 직선 거리로의 이동을 보장한다는 것. (아래 이미지 참고)
      
• Reflow, straightening, fast simulation
  • 그러나 실제론 reflow 없이 한 방에 직선으로 가기 어려움
  • 이를 Z 분포의 이동경로가 실제로 직선인지 탐지하는 score 수식을 만들어 밝힘
    • $S(Z) = \int_0^1 \mathbb{E}\left[\| (Z_1 - Z_0) - \dot{Z}_t \|^2 \right] dt$
    • $\dot{Z}_t$ : 현재의 이동방향
    • 즉, 현재의 이동방향이 직선이면 0이 나오는 구조.
  • 그리고 저자들은 Rectified Flow가 K회 반복될수록 더 직선화됨을 수학적으로 증명함
    • $\min_{k \in \{0, \ldots, K\}} S(Z_k) \leq \frac{\mathbb{E}[\|X_1 - X_0\|^2]}{K}$
      
  • 경로가 Rectification으로 충분히 직선화 되기 전까지는 Distill 단계로 가지 말 것.

Non-Linear Extension

• Non-Linear 상황은 사실, 기존의 work들이 한 것임.
• 단순히 두 분포 사이의 $X_t$를 linear interpolation 하지 않고, 아래와 같이 표현
  • $X_t = \alpha_t X_1 + \beta_t X_0$
  • VP ODE: $\beta_t = \sqrt{1 - \alpha_t^2}$
  • sub-VP ODE: $\beta_t = 1 - \alpha_t^2$
  • 네 즉 Rectified Flow는 Diffusion 모델의 특수해 입니다.