[Study] Rectified Flow

Rectified Flow

디퓨전 모델의 가장 큰 약점 중 하나는 속도다. 노이즈에서 이미지로 가기 위해 보통 50 ~ 1000 번의 신경망 호출이 필요하다. 이미지 한 장을 만드는 데 시간이 오래 걸리는 근본 원인이 여기 있다. 2022년 Xingchao Liu 등이 제안한 Rectified Flow 는 이 문제를 정면으로 공격한다. 출발점(노이즈) 과 도착점(이미지) 사이의 경로를 가능한 한 직선으로 만들면, 적은 스텝만으로도 도착할 수 있다는 것이다.

아이디어 자체는 단순하지만 영향력은 컸다. 이 접근이 나중에 Stable Diffusion 3, Flux 같은 최신 이미지 생성 모델의 학습 패러다임으로 자리 잡았다. 디퓨전이라는 이름을 그대로 쓰지만, 내부적으로는 더 이상 전통적인 디퓨전이 아니라 Flow Matching 계열의 방식이 표준이 된 셈이다.

---

1. 출발점 — 디퓨전은 왜 느린가

전통적인 디퓨전 모델의 학습 방식은 다음과 같다. 깨끗한 이미지에 노이즈를 조금씩 더해가는 전방 과정(forward process) 을 정의하고, 모델이 그 과정을 거꾸로 되돌리는 역방향 과정(reverse process) 을 학습한다. 추론 시에는 순수 노이즈에서 출발해 모델을 반복 호출하며 한 단계씩 노이즈를 빼낸다.

문제는 이 역방향 경로가 휘어 있다는 것이다. 노이즈 공간의 한 점에서 이미지 공간의 한 점으로 가는 길이 직선이 아니라 구불구불한 곡선이다. 곡선을 정확히 따라가려면 작은 스텝으로 여러 번 끊어 가야 한다. 큰 스텝을 한 번에 뛰면 곡선에서 벗어나 엉뚱한 곳에 도착한다.

노이즈  ⤳  중간 1  ⤳  중간 2  ⤳ ··· ⤳  이미지

이게 디퓨전이 50 ~ 1000 스텝을 필요로 하는 이유다. 경로가 휘어 있으니, 경로 위의 점들을 자주 다시 측정해서 방향을 갱신해야 정확히 도착할 수 있다. 한 번에 도약하려고 하면 곡선의 휘어짐 때문에 누적 오차가 커진다.

그렇다면 경로 자체를 직선으로 만들 수는 없을까. 노이즈와 이미지를 잇는 길이 처음부터 곧다면, 큰 스텝 한 번에 도약해도 똑같이 도착할 수 있을 것이다. Rectified Flow 가 던진 질문이 바로 이것이다.

---

2. 핵심 아이디어 — 점 두 개를 직선으로 잇는다

Rectified Flow 의 출발은 놀랍도록 단순하다. 노이즈 분포에서 점 하나를 뽑고, 이미지 분포에서 점 하나를 뽑은 다음, 그 둘을 직선으로 잇는다. 그 직선을 모델이 학습할 정답 경로로 삼는다.

수식으로 보면 더 직관적이다. 시간 t 가 0 에서 1 로 흐른다고 할 때, 시점 t 에서의 위치 x_t 는 :

x_t  =  (1 − t) · x₀  +  t · x₁

여기서 x₀ 은 노이즈, x₁ 은 이미지다. t=0 일 때 순수 노이즈에 있고, t=1 일 때 이미지에 도착한다. 그 사이는 두 점을 잇는 직선상의 보간이다. 끝.

속도는 모든 시점에서 일정하다

이 직선 위에서 시간에 대한 미분을 취하면, 즉 "속도" 를 구하면 매우 단순한 결과가 나온다.

속도 v  =  x₁ − x₀

시간 t 와 무관하게 일정한 값이다. 출발점에서 도착점을 향하는 화살표 하나로, 처음부터 끝까지 같은 방향, 같은 속도로 움직인다는 뜻이다. 이게 직선의 정의이기도 하다.

비유로 풀어보자. 곡선 도로를 따라 운전한다고 생각해보면, 핸들을 계속 꺾어야 해서 매 순간 도로를 다시 봐야 한다. 반면 직선 도로 위에서는 핸들을 한 번 맞춰두면 그대로 끝까지 갈 수 있다. Rectified Flow 는 노이즈에서 이미지로 가는 길을 직선 도로로 깔아두는 작업이다.

---

3. 학습

여기가 처음에 가장 헷갈리는 부분이다. "속도를 학습한다" 는데, 그 속도의 정답은 어디에 있는 걸까. 정답이 들어 있는 별도의 데이터셋이 있는 걸까. 아니다. 정답 속도는 학습 매 스텝마다 그 자리에서 즉석으로 만들어진다.

우리가 가진 것과 만들어낼 것

학습이 시작되기 전, 우리 손에 있는 것은 단 하나다. 이미지 데이터셋. 그게 전부다(LAION 등).

나머지 — 노이즈 샘플, 시간값, 중간 위치, 정답 속도 — 는 전부 학습 도중에 즉석으로 만들어낸다. 노이즈는 정규분포에서 그냥 뽑으면 된다. 시간값도 [0,1] 구간에서 무작위로 뽑는다. 중간 위치는 노이즈와 이미지를 직선 보간하면 나온다. 그리고 정답 속도는? 노이즈에서 이미지를 빼면 끝이다. 별도 학습 데이터가 필요 없다.

한 학습 스텝을 구체적인 숫자로 따라가 보자

간단한 1차원 예시로 학습 한 스텝이 어떻게 이뤄지는지 정확히 보자. 실제로는 이미지 차원이 수만 ~ 수백만이지만, 차원과 무관하게 똑같이 작동한다.

1. 이미지 데이터셋에서 이미지 하나를 뽑는다. 가령 x₁ = 7.
2. 정규분포 N(0,1) 에서 노이즈 하나를 뽑는다. 가령 x₀ = 2.
3. [0,1] 에서 시간을 무작위로 뽑는다. 가령 t = 0.3.
4. 중간 위치를 보간으로 계산한다. x_t = 0.7 · 2 + 0.3 · 7 = 3.5.
5. 정답 속도를 계산한다. v_target = x₁ − x₀ = 7 − 2 = 5.
6. 모델에 x_t=3.5 와 t=0.3 을 넣어 예측 속도를 받는다. 가령 v_pred = 4.2.
7. 손실 = (4.2 − 5)² = 0.64. 이 손실을 줄이는 방향으로 모델 가중치 업데이트.

중요한 점은 5단계다. 정답 속도 5 는 이 학습 스텝에서 즉석으로 만들어진 값이다. 어디 데이터셋에 저장되어 있던 게 아니다. 노이즈와 이미지를 짝지어 빼는 단순 연산이 정답 자체다.

전체 학습 루프

위 한 스텝을 수백만 번 반복하는 게 학습 전체다. 의사코드로 보면 다음과 같다.

# 학습 루프 (한 epoch 안의 한 스텝) for step in training_steps: # --- 입력 준비: 정답 데이터셋이 따로 없다 --- x1 = sample_from_dataset() # 이미지 한 장 x0 = randn_like(x1) # 노이즈 (같은 모양) t = uniform(0, 1) # 시간 무작위 # --- 중간 위치와 정답 속도, 즉석에서 계산 --- xt = (1 - t) * x0 + t * x1 v_target = x1 - x0 # 이게 정답 속도다 # --- 모델 예측 + 손실 --- v_pred = model(xt, t) loss = mse(v_pred, v_target) loss.backward() optimizer.step()

이 코드 어디에도 "속도 데이터셋" 이 없다. v_target 은 매 스텝마다 새로 뽑은 노이즈와 이미지에서 즉석으로 만들어진다. 이것이 Rectified Flow 학습 그림의 전부다.

이미지 차원에서는 어떻게 되나

위 예시는 1차원이지만, 실제 이미지는 차원이 매우 크다. 잠재 공간에서 작업한다고 해도 차원은 수만 ~ 수십만 단위다. 그러나 모든 연산이 차원별로 독립적이라 그림은 동일하다.

변수	설명 (이미지 차원 기준)
x1	이미지(또는 잠재 표현). 예: shape (4, 64, 64) 텐서.
x0	같은 모양의 가우시안 노이즈. shape (4, 64, 64).
t	스칼라 하나. 예: 0.3.
xt	(1−t) · x0 + t · x1. shape (4, 64, 64).
vtarget	x1 − x0. shape (4, 64, 64).
vpred	모델 출력. shape (4, 64, 64).
loss	두 텐서의 원소별 차이 제곱의 평균.

즉, v_target 은 입력과 동일한 모양의 텐서이고, "이 노이즈에서 이 이미지로 향하는 화살표" 를 차원마다 적어둔 것이라고 보면 된다. 모델은 그 화살표 텐서를 예측하도록 학습된다.

한 줄로 요약하면, Rectified Flow 의 학습은 "이미지 데이터셋이 있을 때, 매 스텝마다 노이즈 하나를 뽑고, 시간 하나를 뽑고, 두 점을 보간한 위치 x_t 를 만들고, 그 위치에서 모델이 이미지 − 노이즈 를 예측하도록 가르치는 것" 이다. 정답은 그때그때 즉석에서 만들어진다.

---

4. 디퓨전 학습과의 본질적 차이

3장의 학습 루프를 본 다음 디퓨전과 비교하면 차이가 더 선명해진다. 둘 다 "노이즈와 이미지의 보간을 만들어 모델에 넣는다" 는 점은 같지만, 모델이 예측하는 대상이 다르다.

방식	모델이 예측하는 대상
디퓨전	"여기에 어떤 노이즈가 섞여 있나?" 를 예측. 정답: 섞은 노이즈 자체. v_target 자리에 노이즈 x0 가 들어간다.
Rectified Flow	"여기서 어디로 가야 하나?" 의 방향을 예측. 정답: 이미지에서 노이즈를 뺀 화살표. v_target = x1 − x0.

학습 루프 자체의 구조는 거의 동일하다. 한 줄만 다르다. 디퓨전은 손실을 mse(model(x_t, t), x_0) 로 계산하고, Rectified Flow 는 mse(model(x_t, t), x_1 − x_0) 로 계산한다. 그 한 줄의 차이가 모델이 학습하는 경로의 기하학적 형태를 곡선에서 직선으로 바꿔놓는다.

---

5. 왜 경로가 직선에 가까워지는가

여기서 한 가지 미묘한 부분이 있다. 학습 데이터를 만들 때 "노이즈 한 점과 이미지 한 점을 직선으로 이었다" 고 해서, 학습된 모델이 만드는 실제 경로가 자동으로 직선이 되지는 않는다. 왜 그런지 보자.

학습 중 같은 위치 x_t 가 여러 다른 (x_0, x_1) 쌍에서 만들어질 수 있다. 어떤 쌍에서는 정답 속도가 화살표 A, 다른 쌍에서는 화살표 B 다. 모델은 같은 입력 x_t 에 대해 서로 다른 정답을 동시에 학습해야 하는 모순적인 상황에 놓인다.

이때 모델이 할 수 있는 최선은 그 위치에서 가능한 모든 정답 속도들의 평균을 예측하는 것이다. MSE 손실의 수학적 성질이 정확히 그렇게 만든다. 그리고 이 "평균 속도장" 이 학습 후 모델이 따라가는 경로를 결정한다. 흥미로운 점은 이 평균 속도장이 만드는 경로는 학습 데이터 직선들이 서로 교차하며 만들어내는 흐름보다도 더 매끈하고 직선에 가깝다는 것이다.

왜 그럴까. 여러 직선들이 서로 교차하며 만들어내는 평균 흐름은 개별 직선들의 휘어짐을 상쇄시킨다. 모델이 학습하는 평균 속도장은 자연스럽게 매끈해지고, 그 위에서 흐르는 경로는 처음 학습 데이터의 직선들보다도 더 곧고 안정적인 길이 된다. 이게 Rectified Flow 의 핵심 통찰이다.

Reflow — 경로를 더 곧게 펴기

한 번 학습된 모델로 노이즈에서 이미지를 생성한 뒤, 그렇게 만들어진 (노이즈, 이미지) 쌍을 새 학습 데이터로 써서 다시 학습할 수 있다. 이걸 Reflow 라고 부른다.

1. 1차 학습 — 무작위 (노이즈, 이미지) 쌍으로 모델을 학습한다.
2. 경로 따라가기 — 학습된 모델로 노이즈에서 이미지를 생성한다. 이 과정에서 만들어진 (노이즈, 이미지) 쌍이 모델의 경로 끝점들이다. 1차 학습과 달리 "무작위 매칭" 이 아니라 "모델이 자연스럽게 매칭한 쌍" 이다.
3. 2차 학습 — 그 새 쌍들로 모델을 다시 학습한다. 모델이 만든 경로 끝점들끼리 잇는 직선은 1차 학습 데이터의 직선들보다 서로 덜 교차한다. 결과적으로 모델은 더 깔끔한 직선 흐름을 학습한다.
4. 반복 가능 — 이 과정을 여러 번 반복하면 경로가 점점 더 직선에 가까워진다.

Reflow 를 거친 모델은 1 ~ 4 스텝의 매우 적은 추론 스텝으로도 좋은 이미지를 생성할 수 있다. 경로가 거의 직선이라 큰 스텝으로 도약해도 누적 오차가 작기 때문이다.

---

6. 추론

학습된 Rectified Flow 모델로 이미지를 생성하는 과정은 단순하다. 노이즈 x_0 에서 출발해, 모델이 알려주는 속도 방향을 따라 t 를 0 에서 1 까지 움직이면 된다.

x₀ (노이즈)  →  + v · Δt  →  + v · Δt  → ··· →  x₁ (이미지)

매 스텝마다 현재 위치에서 모델에 속도를 묻고, 그 속도 방향으로 작은 시간 간격 Δt 만큼 이동한다. 경로가 거의 직선이기 때문에 Δt 를 크게 잡아도 도착점에서 크게 벗어나지 않는다. 이게 추론 스텝 수를 획기적으로 줄일 수 있는 직접적인 이유다.

전통적인 디퓨전은 이상적인 결과를 위해 50 ~ 1000 스텝이 필요했다. Rectified Flow 는 기본 학습만으로 25 ~ 50 스텝, Reflow 를 거치면 1 ~ 4 스텝까지 줄어든다. 모델 호출 횟수가 줄어드는 만큼 추론 시간이 같은 비율로 줄어든다.

---

7. 디퓨전과의 비교 

구분	전통 디퓨전 (DDPM 등)	Rectified Flow
경로 형태	휘어진 확률 경로	(거의) 직선
모델이 예측	섞인 노이즈	속도 벡터 (이미지 − 노이즈)
전방 과정	점진적 노이즈 추가 (수식 복잡)	두 점 사이 선형 보간
학습 손실	노이즈 예측 MSE (스케일링 필요)	속도 예측 MSE (단순)
추론 스텝	50 ~ 1000	25 ~ 50 (기본), 1 ~ 4 (Reflow)
대표 사용처	SD 1.5, SDXL	SD3, Flux