이 기사는 Rectified Flow에 대해 배우기 위해 작성한 개인적인 노트의 일부를 정리한 것입니다.

내용은 논문이나 공개 자료를 바탕으로 Codex를 활용하며 정리하고 있지만, 내용 및 구성의 확인과 정리는 저자 본인이 직접 수행하고 있습니다.

현재 인터넷상에는 Rectified Flow에 대한 해설이 그리 많지 않습니다. 본 기사는 입문을 위해 이론에 무게를 두지 않고, 직관적인 이해에 무게를 두고 있습니다.

수학이 전혀 필요 없다는 뜻은 아니지만, 편미분 방정식이나 학습 이론에 대해 깊게 파고드는 것이 아니라, 대략적으로 이해하여 코드를 읽고 쓸 수 있도록 하는 것을 목적으로 만들었습니다.

또한, 본 자료의 전문과 코드는 GitHub에 올려두었습니다.

본 자료는 Rectified Flow를 배우고 싶은 분들을 대상으로 합니다.

특히,

• 딥러닝 (Deep Learning)의 기초를 한 차례 학습한 분
• 논문을 읽거나 PyTorch로 구현을 하기 위한 발판이 되는 해설을 찾고 있는 분
• 수식뿐만 아니라 직관적인 이해도 얻고 싶은 분

을 주요 대상으로 합니다.

엄밀한 이론보다는 "왜 그 식이 되는가", "어떤 의미를 갖는가"를 중시하여 설명하고 있습니다.

단, 수학에 대해서는 다음을 전제로 합니다.

• 편미분에 대해 이해하고 있음
• 미분 방정식(Differential Equation)을 읽을 수 있음
• 통계 기호를 읽을 수 있음

또한, 본 기사는 현대적인 관점에서 Flow Matching과 Rectified Flow의 관련성에 대해서도 설명하고 있습니다. 이 때문에 과거에 작성한 Flow Matching 해설 기사를 읽어보시는 것을 추천합니다.

특히 중요한 것은 (1), (2), (3)입니다. path나 CFM을 이해하고 있는 것이 전제입니다.

이 장에서 다루는 개념: 현재의 통일적인 관점에서는, Flow Matching은 "path와 교사 속도(teacher velocity)를 결정하고, 속도장(velocity field)을 학습한다"는 일반적인 프레임워크로 읽을 수 있다. Rectified Flow는 직선 path를 사용한 속도 회귀(velocity regression)에서 시작하여, 생성 시의 실제 궤적(actual trajectory)을 더욱 직선적으로 만들어 적은 NFE(Number of Function Evaluations)로도 추적하기 쉽게 만드는 것을 중시하는 관련 기법이다. 1-Rectified Flow의 loss만 보면, 본 교재의 Linear Path CFM과 동일한 형태이다. 차이점은 학습 후의 궤적 직선성을 평가하고, 그 flow가 만드는 coupling으로 쌍(pair)을 다시 만들어 재학습하는 Reflow까지 포함하여 설계한다는 점에 있다.

Flow Matching과 Rectified Flow의 차이는 처음에 다음의 세 층으로 나누면 읽기 쉽다. 단, 이는 연구사의 순서를 나타내는 설명이 아니라, 본 교재에서 개념을 정리하기 위한 관점이다. 연구사에 대해서는 다시 기술하겠다.

넓은 프레임워크로서의 Flow Matching

Flow Matching은 확률 경로(probability path) $p_t$와 교사 속도를 설계하고, 그 속도를 뉴럴 네트워크 $v_ heta(t,x)$로 회귀하는 일반적인 학습 방법이다. path는 직선이어도 좋고, Gaussian path여도 좋다. coupling도 독립 coupling, OT coupling, 기타 설계를 선택할 수 있다. -
첫 번째 학습으로서의 1-Rectified Flow

1-Rectified Flow는 양 끝점 쌍 $(x_0,x_1)$을 직선으로 연결하고, 교사 속도 $x_1-x_0$를 학습한다. 이 부분만 보면 본 교재의 Linear Path CFM과 거의 동일하다. 따라서 여기서 새로운 loss가 갑자기 등장하는 것은 아니다. -
연구상의 목표로서의 Rectified Flow/Reflow

Rectified Flow에서 중요한 것은 학습된 ODE의 실제 궤적을 직선에 가깝게 만들어, 적은 NFE로도 생성하기 쉽게 하는 것이다. 이를 위해 학습된 flow로 source 점을 실제로 이동시키고, 도착점을 사용하여 새로운 쌍을 다시 만든다. 이 재쌍(re-pairing) 절차가 Reflow이다.

즉, Flow Matching과의 차이점은 "속도 회귀를 그만두는 것"이 아니다. 동일한 속도 회귀 형태를 사용하면서, 어떤 coupling으로 학습할지, 학습 후의 궤적을 어떻게 진단할지, 다음 학습 쌍을 어떻게 만들지를 설계 대상으로 삼는 것이다.

Reflow가 없는 Rectified Flow와 Linear Path CFM을 수식만으로 명확하게 구별하는 것은 어렵다. 둘 다 양 끝점을 직선으로 연결하고, 그 직선 속도를 뉴럴 네트워크로 회귀하는 형태로 쓸 수 있기 때문이다. 차이는 주로 연구상의 문맥과, 그 이후에 무엇을 문제로 다루는지에서 나타난다.

따라서, 본문에서 "Rectified Flow는 Linear Path CFM과 같은 형태로 읽힐 수 있다"라고 말할 때, 그것은 수식상의 첫 번째 학습(initial training)에 관한 이야기이다. 반면, "Rectified Flow다움"이라고 말할 때는 수송을 직선으로 만들고, 적은 NFE (Number of Function Evaluations)로 생성하기 쉽게 하며, 학습된 flow로 쌍(pair)을 다시 만드는 연구상의 목표까지 포함하여 읽는 것이다.

본 교재의 최소 단위인 Linear Path CFM과 1-Rectified Flow를 비교하면, 첫 번째 학습에서는 많은 구성 요소가 동일하다.

따라서, "Rectified Flow는 CFM loss를 새로운 식으로 교체한 기법"이라고 이해하면 어긋나게 된다. 첫 번째 학습은 Linear Path CFM과 같은 형태이다. 그 위에서, 생성 시의 궤도가 휘는 원인을 coupling을 만드는 방식에서 다시 생각하고, Reflow를 통해 쌍(pair)을 업데이트하는 부분이 Rectified Flow의 핵심적인 읽기 포인트이다.

"Flow Matching이라는 일반적인 프레임워크가 먼저 존재하고, 거기서 Rectified Flow가 특수한 사례로 파생되었다"라고 읽는 것은 연구사(research history) 측면에서 지나치게 단순화된 것이다.

실제로, Rectified Flow를 제안한 "Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow"와, Flow Matching을 일반적인 학습 원리로 제시한 "Flow Matching for Generative Modeling"은 둘 다 2022년 arXiv에 공개된 매우 가까운 시기의 관련 연구이다. Rectified Flow 논문은 수송과 궤도의 직선화, Reflow를 통한 rectification을 중심으로 한다. 반면, Flow Matching 논문은 확률 경로(probabilistic path)와 속도장(velocity field) 회귀를 보다 일반적인 생성 모델링 원리로 정리하고 있다.

그러므로, 본 교재에서 "Rectified Flow를 Flow Matching의 직선 path 버전으로 읽을 수 있다"라고 설명하는 것은 사후적인 통합적 정리이다. 역사적으로 "Flow Matching이 먼저 있었고, 그 특수한 사례로서 Rectified Flow가 태어났다"라고 주장하는 것이 아니다.

연구용으로 읽을 때는 다음 두 가지를 구분한다.

통합적인 이해:
Rectified Flow의 첫 번째 학습은 Linear Path CFM과 같은 형태로 쓸 수 있다

연구사·방법론:
...

이 구분을 두면, 동일한 수식이 나오는데도 논문명이나 구현명이 다른 이유를 오해하기 쉬워지지 않는다.

source 점과 target 점을 직선으로 연결하면, 쌍(pair)마다의 교사 화살표(teacher arrow)는 직선이 된다. 하지만 수많은 직선이 중간에 교차하면, 같은 장소에 서로 다른 방향의 교사 화살표가 모이게 된다. velocity model은 그것들을 평균하여 학습하기 때문에, 생성 시 solver가 따라가는 실제 궤도는 휘어질 수 있다.

도로에 비유하자면, 최초의 Flow Matching은 무작위로 선택한 출발지와 목적지를 자로 연결하고, 그 선 위의 화살표를 대량으로 보여주는 학습이다. 선이 교차하는 교차로에서는 모델이 여러 진행 방향을 하나로 합쳐야만 한다. 그 결과, 완성된 도로는 개별적인 자의 선보다 휘어질 수 있다.

Rectified Flow는 일단 완성된 도로를 실제로 달려보고, "이 출발점은 이 flow를 통해 어디에 도착했는가"를 새로운 대응 관계로 사용한다. 그 대응을 통해 다시 직선 교사(teacher)를 만들면, 서로 모순되는 화살표가 줄어들고 실제 도로도 직선에 가까워지기 쉬워진다. 이 쌍(pair)을 다시 만드는 과정이 다음 장의 Reflow이다.

1회차:
독립적으로 만든 쌍
-> 직선 교사 화살표로 학습
...

여기서 말하는 "직선"은 훈련용 path가 직선이라는 것만을 의미하지 않는다. 학습된 속도장에서 입자를 생성했을 때의 궤도가, 멀리 돌아가지 않고 끝점으로 향한다는 것을 의미한다.

제1부의 Linear Path CFM와 1回目の Rectified Flow (1-Rectified Flow)는 둘 다

x_t=(1-t)x_0+tx_1,\qquad u_t=x_1-x_0

를 사용한 속도 회귀 (velocity regression)이다. 코드상에서도 두 loss는 동일한 처리로 읽을 수 있다.

# Linear Path CFM로 읽는 경우
loss_fm = conditional_flow_matching_loss(
model, LinearPath(), x0, x1
...

rectified_flow_loss는 내부에서 Linear Path의 CFM loss를 호출한다. RectifiedFlowMLP도 구조적으로는 MLPVelocity와 동일한 작은 MLP이다. 이름을 나누어 놓은 것은 새로운 모델 구조를 도입했기 때문이 아니라, 궤도 직선성 (trajectory straightness)과 Reflow를 중시하는 실험임을 명시하기 위해서이다.

src/train_rectified_flow_2d.py는 1-Rectified Flow를 2D toy 데이터로 학습하는 스크립트이다.

python .\src\train_rectified_flow_2d.py --steps 2000 --batch 512 --device cpu

모델은 RectifiedFlowMLP이다.

model = RectifiedFlowMLP()
loss = rectified_flow_loss(model, x0, x1)

학습 후에는 생성된 샘플뿐만 아니라 궤도 (trajectory)를 추출하여 straightness_ratio를 계산한다.

trajectory = euler_solve(
lambda t, x: model(t, x),
x0,
...

즉, 첫 번째 코드 변경의 핵심은 loss가 아니라, 궤도를 저장하고, 직선성을 평가하며, 다음 Reflow로 넘길 수 있도록 만드는 것이다.

일반적인 Flow Matching에서는 처음에 준비한 소스 (source) 점과 타겟 (target) 점의 쌍을 사용하여 학습한다. 예를 들어 독립 결합 (independent coupling)이라면, 소스에서 가져온 점과 타겟에서 가져온 점을 상당히 기계적으로 연결한다. 그 쌍에 대해 직선 경로 (linear path)를 만들고, 각 시점마다 "이쪽으로 나아가라"는 교사 화살표 (teacher arrow)를 부여한다.

이 방법은 단순하고 다루기 쉽다. 반면, 쌍을 맺는 방식이 좋지 않으면 도중에 많은 직선이 교차하게 된다. 같은 장소 근처에 오른쪽으로 가는 화살표, 왼쪽으로 가는 화살표, 위로 가는 화살표가 모이게 된다. 뉴럴 네트워크 (neural network)는 이것들을 하나의 속도장 (velocity field)으로 학습하기 때문에, 생성 시의 입자는 평균화된 화살표를 따라 휘어진 길을 진행할 수 있다.

Reflow의 발상은 여기서 한 번 멈춰 서는 것이다. 첫 번째 학습으로 만들어진 flow를 실제로 사용하여, 소스 점을 생성 측으로 흘려보내 보는 것이다.

여기서 "흘려보낸다"는 것은, 소스 점 $x_0$를 학습된 속도장 $v_ heta(t,x)$ 위에 놓고, 생성 시와 마찬가지로 ODE solver를 사용하여 조금씩 움직인다는 의미이다. 예를 들어 Euler sampler라면, 현재 위치에서 속도를 예측하고 그 방향으로 조금 나아간다. 이 업데이트를 $t=0$부터 $t=1$까지 반복하면, 소스 점은 어떤 종점에 도착한다. Reflow에서는 이 종점을 "이 소스 점이 현재 학습된 flow를 통해 도착한 장소"로 취급한다.

그러면 "이 소스 점은 이 학습된 flow를 통해 이 장소에 도착했다"라는 대응 관계를 얻을 수 있다. Reflow에서는 그 대응 관계를 다음 훈련 쌍으로 사용한다.

일반적인 Flow Matching과의 차이점은 다음과 같이 정리할 수 있다.

그림: 일반적인 Flow Matching과 Reflow를 사용하는 Rectified Flow의 학습 방법 차이

그림 읽는 법: 왼쪽은 일반적인 Flow Matching이다. 학습 전에 결정된 source와 target의 쌍을 사용하여, 해당 쌍으로부터 직선 path와 지도 속도(teacher velocity)를 만들어 속도장(velocity field)을 학습한다. 생성 궤도가 휘어지더라도, 통상적으로는 그 결과를 사용하여 다음 학습 쌍을 다시 만들지 않는다. 오른쪽은 Reflow를 사용하는 Rectified Flow이다. 먼저 첫 번째 속도장을 학습하고, 학습된 flow를 생성 시와 마찬가지로 ODE solver로 따라간다. 얻어진 도착점을 사용하여 쌍을 다시 만들고, 그 새로운 쌍으로 다음 속도장을 학습한다. 따라서 차이점은 loss 식에 있는 것이 아니라, 다음 학습에 사용할 coupling을 업데이트하느냐의 여부에 있다.

중요한 점은, Reflow 후에도 "직선 path로 속도를 학습한다"라는 학습 형태 그 자체는 크게 변하지 않는다는 것이다. 변하는 것은 어떤 source 점과 어떤 도착점을 쌍으로 연결할 것인가이다. 즉, Reflow는 새로운 loss를 도입한다기보다, 학습된 flow를 사용하여 다음 교재 데이터를 다시 만드는 조작이다.

이 절에서는 직관만 파악하면 된다. 수식으로서의 marginal preserving property는 다음 절에서 다룬다.

Rectified Flow의 이론 논문에서 중요한 것은 marginal preserving property, 즉 "각 시점의 주변 분포(marginal distribution)를 보존한다"는 관점이다.

먼저, 임의의 coupling으로부터 쌍 $(X_0, X_1)$을 추출하고, 직선 보간(linear interpolation)

X_t = (1-t)X_0 + tX_1

을 고려한다. 이때 $X_t$의 분포를 $p_t$라고 쓴다.

Rectified Flow는 조건부 속도 $X_1-X_0$를 그대로 모든 점에 붙이는 것이 아니라, 동일한 중간점 $x$에 도달할 가능성이 있는 쌍들의 속도를 평균하여 주변 속도장(marginal velocity field)을 만든다.

v(t,x) = \mathbb{E}[X_1 - X_0 \mid X_t=x].

이 식은 본 교재에서 여러 번 등장했던 "제곱 오차는 조건부 평균을 학습한다"라는 내용과 동일하다. 뉴럴 네트워크 $v_ heta(t,x)$는 이 $v(t,x)$를 회귀(regression)로 근사하고 있다.

이 속도장으로 ODE

\frac{dZ_t}{dt}=v(t,Z_t),\quad Z_0 \sim p_0

를 풀면, 이상적으로는 $Z_t$의 주변 분포가 $p_t$와 일치하게 된다. 특히 $t=1$에서는 $Z_1$의 분포가 target 분포 $p_1$이 된다.

이 부분이 중요하다. Rectified Flow는 개별 샘플 궤도를 원래의 직선 보간과 완전히 동일하게 만드는 방법이 아니다. 많은 직선 보간이 만드는 시점별 분포를 유지하도록 평균 속도장으로 ODE를 만드는 방법이다. 따라서 "훈련 쌍별 직선"과 "학습된 ODE의 입자 궤도(particle trajectory)"는 구분해서 읽어야 한다.

Euler sampler는 짧은 직선 화살표를 이어 붙여 ODE를 근사한다.

x_{k+1} = x_k + \Delta t\, v_\theta(t_k, x_k).

실제 궤도가 휘어 있다면, 적은 스텝으로는 굴곡을 놓치기 쉽다. 반대로 궤도가 거의 직선이라면, 거친(coarse) 스텝으로도 끝점에 가까워지기 쉽다.

Rectified Flow의 "straight is fast"라는 생각은 바로 여기에 있다. 속도장의 궤도가 곧으면 NFE를 줄여도 생성 품질이 떨어지기 어렵다. 연구 측면에서 이는 sampling 가속화를 위한 설계 원리가 된다.

단, 주의할 점이 있다. 훈련 시에 직선 path를 사용하더라도, 학습된 주변 속도장의 궤도가 완전한 직선이 된다고 단정할 수는 없다. 많은 쌍으로부터 오는 지도 속도를 평균하기 때문에, 동일한 지점 근처에서 여러 방향의 속도가 섞일 수 있다.

이 주의사항은 marginal preserving property와 모순되지 않는다. 주변 분포를 올바르게 운반하는 것과 각 입자의 궤도가 완전히 직선인 것은 별개의 성질이다. Rectified Flow 연구에서는 이 두 가지를 혼동하지 않는 것이 중요하다.

이 장에서 다루는 개념: Reflow는 학습된 Rectified Flow를 통해 source 표본을 실제로 수송하고, 그 시작점과 끝점을 새로운 학습 쌍으로 삼아 다음 Rectified Flow를 학습하는 반복 조작이다. 모델 구조의 이름이 아니라, 학습 데이터가 되는 coupling을 업데이트하는 절차이다.

Reflow는 학습된 flow를 사용하여 새로운 훈련 쌍을 다시 만드는 조작이다.

첫 번째 학습에서는 source에서 점을 하나 뽑고, target에서 점을 하나 뽑아 두 점을 연결했다. 이 단계의 쌍(pair)은 데이터의 실제 대응 관계를 알고 있는 것이 아니다. 많은 경우, "source의 이 점과 target의 이 점을 연결하기로 한다"라는 학습상의 임시 대응이다.

Reflow에서는 먼저 학습된 모델로 source 점을 실제로 흘려보낸다(flow). 그러면 모델 스스로가 "이 source 점이라면, 생성 과정에서 이 근처에 도착한다"라고 대응 관계를 만든다. 그 도착점을 다음 학습에서 사용할 새로운 상대(target)로 취급한다.

일반적인 Flow Matching과의 차이점은 쌍(pair)을 결정하는 타이밍에 있다.

일반적인 Flow Matching:
학습 전에 쌍을 결정함
-> 그 쌍을 직선으로 연결함
...

직관적으로는, 첫 번째 모델이 "이 source 점은 대략 이 근처로 간다"라고 알려주는 것이다. 그 대응 관계를 사용하여 쌍을 다시 맺음으로써, 다음 학습에서는 더 정직한(straight) 직선을 교사(teacher)로 삼을 수 있다. 여기서 변하는 것은 loss 식 자체가 아니라, 훈련 데이터로 보여주는 "출발점과 도착점의 쌍"이다.

Reflow는 학습된 ODE가 만드는 사상(mapping)을 사용하여 새로운 coupling을 유도한다.

첫 번째 학습된 모델로 source 측의 점을 $t=0$에서 $t=1$까지 흘려보낸다.

\frac{dX_t}{dt} = v_\theta(t, X_t),\quad X_0 = x_0.

그 종점을 새로운 target 측의 점으로 취급한다.

Z_0 = X_0,\quad Z_1 = \mathrm{ODESolve}(v_\theta, X_0).

이 쌍 $(Z_0, Z_1)$은 source 측의 주변 분포(marginal distribution)를 유지한다. 이상적인 Rectified Flow라면 target 측의 주변 분포도 유지한다. 즉, Reflow는 "target 분포를 버리고 모델의 출력만을 믿는" 조작이 아니라, 분포를 운반하는 flow가 유도한 coupling을 다음 훈련에 사용하는 조작이다.

Qiang Liu의 "Rectified Flow: A Marginal Preserving Approach to Optimal Transport"에서 강조되는 점은 이 조작이 OT(Optimal Transport, 최적 운송)와 관련이 있다는 점이다. 대략적으로 Rectified Flow는 임의의 초기 coupling에서 출발하여, marginal을 유지한 채로 운송 비용(transport cost)이 더 작고 더 직선적인 coupling에 가까워지는 것을 목표로 한다.

단, 여기서 "Rectified Flow는 항상 진정한 OT map을 배운다"라고 읽어서는 안 된다. 유한한 데이터, 유한한 모델, 유한한 스텝의 solver, 그리고 최적화 오차가 존재하는 구현에서는 OT와 관련된 근사적인 coupling을 얻게 된다. 본 교재에서는 Reflow를 다음과 같이 정의한다.

이론적 역할: marginal을 유지하는 flow로부터 새로운 coupling을 유도함
구현적 역할: 학습된 모델로 source를 흘려보내 다음 훈련 쌍을 만듦
실험적 확인: NFE, 생성 품질, 궤도 직선성, 계산 비용을 확인함

이렇게 정리하면 OT-CFM과의 차이도 명확해진다. OT-CFM은 샘플 간의 거리 비용으로부터 coupling을 선택한다. Reflow는 학습된 flow가 만든 도착점으로부터 coupling을 만든다. 둘 다 coupling 설계 방식이지만, 사용하는 정보가 다르다.

Reflow는 "궤도가 휘어지면 벌칙을 준다"라는 항을 loss에 추가하는 것이 아니다. 직선성을 직접적인 벌칙(penalty)으로 최적화하는 것이 아니라, 다음 학습에서 보여주는 쌍을 바꿈으로써 직선적인 속도장(velocity field)을 배우기 쉽게 만드는 것이다.

일반적인 Flow Matching에서 독립적인 coupling을 사용하면, source 점과 target 점의 대응은 상당히 임의적(arbitrary)이 된다. 예를 들어, 왼쪽 아래의 source 점이 오른쪽 위의 target 점과 연결되고, 근처의 다른 source 점이 왼쪽 위의 target 점과 연결될 수 있다. 이러한 쌍들을 모두 직선으로 연결하면 중간에 선들이 교차하기 쉽다.

선이 교차하면 그 교차 지점 근처에서는 다음과 같은 일이 발생한다.

같은 장소 근처에 온 점 A: 오른쪽 위로 가고 싶어 함
같은 장소 근처에 온 점 B: 왼쪽 위로 가고 싶어 함
같은 장소 근처에 온 점 C: 오른쪽 아래로 가고 싶어 함
...

Reflow에서는 먼저 학습된 flow로 source 점을 실제로 운반한다. 그러면 각 source 점에 대해 "이 flow가 실제로 도착시킨 점"을 얻을 수 있다. 이 도착점은 완전한 정답 대응은 아니지만, 적어도 현재의 flow가 만든 이동의 결과이다.

다음 학습에서는 이 source 점과 도착점을 직선으로 다시 연결한다. 그러면 이전처럼 임의로 연결된 쌍(pair)보다, 학습된 flow의 움직임을 따르는 대응이 되기 쉽다. 결과적으로 직선 path들 사이의 무리한 교차나, 같은 장소에서 모순되는 교사 화살표(teacher arrow)가 줄어들기 쉽다.

이 차이를 그림으로 나타낸다.

[IMG:1] Reflow에 의한 쌍(pair) 업데이트와 궤도 직선화의 직관

그림 읽는 법: 왼쪽은 일반적인 Flow Matching이다. 학습 전에 결정된 source와 target의 쌍을 직선으로 연결하기 때문에, 도중에 여러 개의 교사 화살표가 같은 장소로 모이게 되며, 뉴럴 네트워크(neural network)는 이를 평균하여 학습한다. 그 결과, 생성 시의 실제 궤도는 휘어질 수 있다. 오른쪽은 Reflow 이후의 학습이다. 학습된 flow로 source 점을 실제로 흘려보내고, 그 도착점을 새로운 상대로서 쌍을 다시 연결한다. 이를 통해 다음 학습에서 같은 장소에 모이는 모순된 교사 화살표가 줄어들고, 실제 궤도가 직선적이 되기 쉽다.

이 메커니즘을 한 문장으로 말하자면, Reflow는 다음과 같은 작업이다.

휘어진 생성 궤도를 관찰한다
-> 그 궤도의 시점과 종점을 새로운 쌍(pair)으로 만든다
-> 그 쌍을 직선으로 연결하는 교사(teacher)로 교체한다
...

따라서, Reflow에서 "직선이 되도록 학습시킨다"는 것은 loss에 straightness ratio를 넣는 것이 아니다. 쌍(pair)의 대응 관계를 업데이트하고, 속도 회귀(velocity regression)에서 평균화되는 화살표의 충돌을 줄임으로써, 생성 시의 실제 궤도가 직선에 가까워지도록 유도한다는 의미이다.

단, 이것이 보장되는 것은 아니다. base model의 정밀도가 낮거나, pair-steps가 너무 거칠거나, 모델 용량(capacity)이 부족하거나, 혹은 Reflow 이후의 쌍(pair)에서도 path 교차가 많은 경우, 궤도는 충분히 직선화되지 않는다. 따라서 straightness ratio로 Reflow 전후의 궤도를 평가할 필요가 있다.

src/fm_minimal/reflow.py는 Reflow에서 사용할 새로운 시점·종점 쌍을 생성한다.

@torch.no_grad()
def make_reflow_pairs(model, x0, steps=64, solver=euler_solve):
    x1_hat = solver(lambda t, x: model(t, x), x0, steps=steps)
    ...

이 함수는 학습된 모델을 사용하여 x0를 흘려보내고, 도착점 x1_hat을 반환한다. 여기서 detach()를 하는 이유는 Reflow용 쌍(pair) 생성을 다음 모델 학습의 계산 그래프(computation graph)로부터 분리하기 위해서이다.

사용법은 다음과 같다.

from fm_minimal import LinearPath, make_reflow_pairs
path = LinearPath()
# model_1은 이미 1-Rectified Flow로서 학습되었다고 가정한다.
...

이 코드에서 중요한 것은 Reflow 이후의 학습에서도 loss의 형태가 변하지 않는다는 점이다. 변하는 것은 훈련 쌍(training pair)이다.

1회차: source sample과 target sample을 연결
Reflow용 쌍 생성: source sample을 model_1로 흘려보냄
2회차: source sample과 model_1의 도착점을 연결

이런 의미에서, Reflow는 "모델 구조를 바꾸는 방법"이 아니라, "다음 학습에 사용할 coupling을 학습된 flow로부터 만드는 방법"이라고 이해하는 것이 통찰력을 얻기에 좋다.

일련의 실험은 먼저 1-Rectified Flow를 학습하고, 그 체크포인트(checkpoint)를 train_reflow_2d.py로 전달한다.

python .\src\train_rectified_flow_2d.py `--out .\_outputs\rectified_flow`
python .\src\train_reflow_2d.py ...

pair-steps는 base model을 적분하여 새로운 종점을 만드는 정밀도를 결정한다. 이는 새 모델의 학습 step 수나 평가 시의 NFE(Number of Function Evaluations)와는 별개이다. 새 모델은 base checkpoint의 model_config를 계승하므로, Reflow와 모델 구조 변경을 동시에 수행하지 않고 우선 coupling 업데이트만의 효과를 비교할 수 있다.

Rectified Flow와 Reflow는 같은 것이 아니다. Rectified Flow는 분포를 운반하는 속도장(velocity field)을 직선적으로 만들려는 학습 프레임워크(framework)이다. Reflow는 그 프레임워크 안에서, 학습된 flow를 사용하여 새로운 훈련 쌍을 만드는 조작이다.

짧게 말하자면, Rectified Flow는 "무엇을 배울 것인가"이고, Reflow는 "다음 학습 데이터를 어떻게 만들 것인가"이다. 첫 번째 Rectified Flow를 학습한 것만으로는 아직 Reflow를 수행한 것이 아니다. 학습된 모델로 source를 흘려보내고, 그 도달점을 사용하여 다음 쌍 (pair)을 만든 단계에 이르러야 Reflow가 된다.

Rectified Flow를 다룰 때, loss만으로는 "궤도 (trajectory)가 직선이 되었는지"를 알 수 없다. 그래서 생성 궤도의 길이와 시작점에서 종점까지의 직선 거리를 비교한다.

여기서의 straightness ratio는 손실 함수 (loss function)의 일부가 아니라, Reflow의 효과를 확인하기 위한 평가·진단 지표이다. Reflow 이후의 학습에서도 기본적인 loss는 직선 경로 (path) 상의 속도 회귀 (velocity regression)이다. straightness ratio를 loss에 더해서 모델을 직접 최적화하는 것이 아니다.

Reflow에서의 사용법은 다음과 같이 읽을 수 있다.

1번째 Rectified Flow를 학습한다
-> 생성 궤도 (trajectory)를 추출한다
-> straightness ratio로 궤도가 어느 정도 휘어져 있는지 측정한다
...

즉, straightness ratio는 "Reflow에 의해 궤도가 직선이 되었는가"를 보기 위한 척도이다. 값이 낮아지면, 동일한 source 지점에서 target 측으로 향하는 경로가 짧아지고 우회하는 경로가 줄어들었다고 해석할 수 있다. 값이 낮아지지 않는다면, Reflow 쌍 (pair)을 만드는 방식, pair-steps, 학습 step 수, 모델 용량(capacity), 혹은 원래 coupling의 교차(intersection)가 문제가 되고 있을 가능성이 있다.

\text{straightness ratio}
=
\frac{\text{trajectory length}}{\|X_1 - X_0\|}.

완전한 직선이라면 이 값은 1에 가깝다. 휘어지거나 우회하면 궤도 길이 (trajectory length)가 늘어나므로 1보다 커진다.

코드에서는 다음과 같이 계산한다.