1-4. 플로우 매칭

노이즈와 그림을 직선으로 잇는다. 그리고 그 직선 위를 움직이는 속도를 배운다. Flux, SD3, Qwen-Image, Wan 등 지금 쓰는 모델은 거의 전부 이 방식이다. 0-1에서 본 식이 여기서 주인공이 된다.

DDPM에서 무엇이 달라졌나

1-3에서 샘플링은 경로를 따라 걷는 일이라고 했다. 경로가 휘어 있으면 걸음을 잘게 쪼개야 한다. 그러면 생각이 여기에 닿는다. 처음부터 경로를 곧게 만들면 되지 않나.

플로우 매칭은 섞는 식을 가장 단순한 것으로 잡는다.

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

0-1의 식이다. t = 0이면 그림, t = 1이면 노이즈다. 그 사이는 둘을 잇는 직선이다.

속도

직선 위에서 t가 조금 커지면 x_t는 어느 쪽으로 얼마나 움직일까. 식을 t로 미분하면 나온다.

v = ε − x₀

이것이 속도다. 그림에서 노이즈를 향하는 화살표다. 한 직선 위에서는 어디서나 같다. 직선이니까.

학습

신경망이 이 속도를 맞히게 한다.

손실 = ‖ v_θ(x_t, t) − (ε − x₀) ‖²

절차는 0-6과 똑같다. 맞히는 대상만 다르다.

반복:
    그림 x₀ 과 노이즈 ε 을 뽑는다
    t 를 0~1 에서 뽑는다
    x_t = (1 − t)·x₀ + t·ε
    신경망이 x_t, t 를 보고 속도를 예측한다
    예측과 (ε − x₀) 의 제곱 오차를 줄인다

0-6의 코드에서 한 줄만 바꾸면 된다.

loss = F.mse_loss(model(xt, t), eps)        # 0-6: 노이즈를 맞힌다
loss = F.mse_loss(model(xt, t), eps - x0)   # 플로우 매칭: 속도를 맞힌다

생성

순수 노이즈(t = 1)에서 출발해서, 속도의 반대 방향으로 걸어 t = 0까지 온다. 1-3의 Euler 방법이다.

x = torch.randn(n, 1, 28, 28)          # 순수 노이즈
dt = 1.0 / steps
for i in range(steps):
    t = torch.full((n,), 1.0 - i * dt)
    x = x - dt * model(x, t)           # 속도의 반대쪽으로 한 걸음

이 여섯 줄이 이미지 생성기다. 실제 도구도 뼈대는 이것이다.

돌려 본 결과

전체 코드는 1-4-flow.py다. 0-6의 작은 U-Net을 10에폭 학습한 뒤, 같은 노이즈 16개에서 스텝 수만 바꿔 생성했다.

위에서부터 1스텝, 4스텝, 20스텝, 50스텝이다.

• 1스텝: 흐릿한 덩어리. 모든 숫자의 평균이다. 1-2에서 말한 "한 번에 되돌리면 평균이 나온다"가 눈에 보인다.
• 4스텝: 숫자의 형태가 잡힌다. 획이 아직 무르다.
• 20스텝: 획이 또렷하다.
• 50스텝: 20스텝과 거의 같다. 더 올려도 얻는 것이 없다.

모델이 작고 조건(어떤 숫자를 그릴지)도 주지 않았기 때문에 숫자가 아닌 모양도 섞여 있다. 그래도 노이즈에서 손글씨 획이 만들어진다.

직선인데 왜 한 걸음에 안 되나

한 쌍(그림 하나, 노이즈 하나)의 경로는 직선이다. 그런데 1스텝 결과는 뭉개졌다. 이유가 있다.

공간의 한 자리 x_t를 지나는 직선은 하나가 아니다. 서로 다른 그림과 노이즈의 쌍이 같은 자리를 지나간다. 신경망은 그 자리에서 하나의 답만 낼 수 있다. 그래서 그 자리를 지나는 모든 직선의 속도를 평균 낸 값을 배운다.

• t = 1 근처: 어떤 그림으로든 갈 수 있다. 평균 속도는 "모든 그림의 평균" 쪽을 가리킨다.
• t = 0 근처: 갈 수 있는 그림이 거의 하나로 좁혀졌다. 속도가 정확하다.

그래서 실제로 걸어가는 경로는 완전한 직선이 아니다. 처음에는 대략적인 방향으로 가다가, 갈수록 특정 그림 쪽으로 휜다. 그래도 DDPM의 경로보다 훨씬 곧다. 20~30스텝이면 충분한 이유다.

경로를 더 곧게: 정류 흐름

경로를 더 펴는 방법이 있다.

1. 학습된 모델로 노이즈에서 그림을 많이 생성한다. (노이즈, 그 노이즈에서 나온 그림) 쌍이 생긴다.
2. 그 쌍을 직선으로 이어 다시 학습한다.

처음에는 그림과 노이즈를 아무렇게나 짝지었다. 이번에는 "실제로 이어지는" 짝이다. 직선끼리 엇갈리는 일이 줄어들어 경로가 곧아진다. 이것을 정류 흐름(rectified flow)이라고 한다. 반복하면 몇 스텝만으로 생성할 수 있게 된다. 1-7의 증류로 이어지는 생각이다.

세 가지 예측은 서로 바뀐다

속도를 알면 원본과 노이즈를 계산할 수 있다.

x₀ = x_t − t·v
ε  = x_t + (1 − t)·v

그래서 "노이즈를 맞히는 모델"과 "속도를 맞히는 모델"은 본질적으로 같은 정보를 담고 있다. 표현이 다를 뿐이다. 1-2의 ε 예측, x₀ 예측, v 예측과 같은 이야기다.

실제 모델에서 보이는 설정

shift(시프트). 큰 그림은 노이즈에 강하다. 같은 양의 노이즈를 섞어도 1024 그림은 256 그림보다 형태가 더 많이 남는다. 픽셀이 많아서 노이즈가 평균되어 지워지기 때문이다. 그래서 해상도가 높을수록 걸음을 노이즈가 많은 쪽으로 몰아 준다. 이 정도를 정하는 값이 shift다. 플로우 매칭 모델의 워크플로에 이 값을 넣는 노드가 들어 있다.

t를 고르게 뽑지 않는다. 학습할 때 중간 구간의 t를 더 자주 뽑는다. 양 끝은 쉽고 가운데가 어렵기 때문이다.

샘플러는 단순한 것으로 충분하다. 경로가 곧아서 Euler만으로 잘 된다. 플로우 매칭 모델의 권장 설정이 대개 Euler인 이유다.

정리: DDPM과 플로우 매칭

뼈대는 같다. 노이즈를 섞고, 되돌리는 방향을 배우고, 그 방향을 따라 걷는다.

확인 문제

1. x₀ = 2, ε = −0.7일 때 속도 v는?
2. 위에서 t = 0.5일 때 x_t는 0.65다. 속도를 써서 x₀을 되찾아 보자.
3. 경로가 직선인데도 1스텝 생성이 뭉개지는 이유는?

풀이

1. −2.7이다. ε − x₀ = −0.7 − 2.
2. x₀ = x_t − t·v = 0.65 − 0.5 × (−2.7) = 2다.
3. 같은 자리를 지나는 직선이 여러 개라서, 모델은 그 속도들의 평균을 배운다. 노이즈 근처에서 그 평균은 모든 그림의 평균 쪽을 가리킨다.