[논문 리뷰] Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow

Xingchao Liu, Chengyue Gong외 1명|arXiv (Cornell University)|2022. 09. 07.

Model Reduction and Neural Networks인용 수 75

한 줄 요약

Rectified flow는 두 개의 실험적 분포 간에 직선 경로를 따라 운송하기 위해 ODE를 학습하며, 이는 간단한 비선형 최소제곱 목적함수를 통해 달성된다; 재귀적 Rectified Flow는 거의 직선인 흐름을 만들어 생성 및 전이 작업에서 빠른 한 단계 추론을 가능하게 한다.

ABSTRACT

We present rectified flow, a surprisingly simple approach to learning (neural) ordinary differential equation (ODE) models to transport between two empirically observed distributions π_0 and π_1, hence providing a unified solution to generative modeling and domain transfer, among various other tasks involving distribution transport. The idea of rectified flow is to learn the ODE to follow the straight paths connecting the points drawn from π_0 and π_1 as much as possible. This is achieved by solving a straightforward nonlinear least squares optimization problem, which can be easily scaled to large models without introducing extra parameters beyond standard supervised learning. The straight paths are special and preferred because they are the shortest paths between two points, and can be simulated exactly without time discretization and hence yield computationally efficient models. We show that the procedure of learning a rectified flow from data, called rectification, turns an arbitrary coupling of π_0 and π_1 to a new deterministic coupling with provably non-increasing convex transport costs. In addition, recursively applying rectification allows us to obtain a sequence of flows with increasingly straight paths, which can be simulated accurately with coarse time discretization in the inference phase. In empirical studies, we show that rectified flow performs superbly on image generation, image-to-image translation, and domain adaptation. In particular, on image generation and translation, our method yields nearly straight flows that give high quality results even with a single Euler discretization step.

연구 동기 및 목표

• 생성 모델링과 도메인 전이를 운송 맵 프레임워크 아래 통합한다.
• 분포 간의 직선적이고 최단 경로 흐름을 선호하는 확장 가능한 ODE 기반 접근법을 도입한다.
• 운송을 이끄는 속도장을 학습하기 위한 간단한 최소제곱 목적함수를 개발한다.
• 이론적 특성: 비교차 흐름, 주변 분포 보존, 축소된 볼록 운송 비용 등을 보인다.
• 이미지 생성, 이미지-대-이미지 번역, 도메인 적응에서 실험적으로 효과를 보인다.

제안 방법

• Rectified Flow를 dZ_t = v(Z_t, t) dt의 ODE로 정의하여 coupled 데이터에서 샘플된 X_1 - X_0의 거의 직선 방향을 따라 pi_0에서 pi_1으로 운송한다.
• 비선형 최소제곱 목적함수를 해결하여 속도장 v를 학습하고, X_1 - X_0와 v(tX_1 + (1 - t)X_0, t) 사이의 불일치를 t∈[0,1]에서 최소화한다.
• 교차하는 궤적을 재배선하여 비교차 흐름을 얻으면서 보간 안 X_t의 주변 분포를 보존하여 인과관계를 보장한다.
• 학습된 v로 pi_0 및 pi_1에서 순방향 및 역방향 샘플링을 사용하고, 선택적 reflow를 적용하여 점점 더 직선적인 경로를 얻는다.
• 빠른 한 단계 추론을 위해 k-rectified flow를 모방하도록 신경망 T를 학습시키는 선택적 증류.
• 일반 시간-미분 가능 X_t로 선형 보간을 대체하는 비선형 확장을 논의하고, 이를 확률 흐름 ODEs와 DDIM과 연계한다.

실험 결과

연구 질문

• RQ1 간단한 ODE 기반 운송 모델을 학습시켜 두 개의 경험적 분포 사이에서 질량을 최소한의 계산 부담으로 이동시킬 수 있는가?
• RQ2 Rectified Flow가 비교차이고 결정론적 결합을 만들어 볼록 비용에 대한 운송 비용을 줄이는가?
• RQ3 재귀적 Rectification이 점점 더 직선적인 궤적을 만들어 정확한 단일 단계 또는 소수 단계 추론을 가능하게 하는가?
• RQ4 Rectified Flow가 생성 및 전이 작업에서 확률 흐름 ODEs와 잡음 제거 확산 모델과 어떤 관련이 있고 어떻게 다른가?
• RQ5 Rectified Flow가 이미지 생성, 이미지-대-이미지 번역, 도메인 적응에 효과적인가?

주요 결과

• Rectified Flow는 적대적 학습이나 가능도 기반 추론 없이 경험적 분포 간 운송을 학습하는 간단한 최적화 기반 방법을 제공한다.
• 결과 흐름은 비교차이며 모든 시점에서 선형 보간 과정의 주변 법칙을 보존한다.
• Rectification은 모든 볼록 비용 c에 대해 입력 결합보다 낮은 볼록 운송 비용을 갖는 결합을 만들어낸다.
• 재귀적 Rectification은 점점 더 직선적인 흐름을 만들어 조잡한 시간 이산화로도 정확한 시뮬레이션과 잠재적으로 한 단계 추론을 가능하게 한다.
• 실험적으로, Rectified Flow는 고품질 이미지 생성 및 번역 결과를 달성하며, 단일 단계 CIFAR10 FID 점수 4.85 및 리콜 0.51의 최첨단 성능을 보고했다.