플로우 매칭

최적 수송: 얽히지 않게 짝짓기

조건부 플로우 매칭에서 보았듯, 짝마다 길은 곧아도 그 길들을 평균한 흐름은 휜다. 휜 흐름을 큰 스텝으로 따라가면 빗나가서 스텝을 잘게 나눠야 하는데, 흐름이 휜 것은 한 점을 여러 (z, x) 짝의 직선이 엇갈려 지나가기 때문이었다. 그렇다면 노이즈 점과 데이터 점의 짝을 길이 덜 엇갈리게 지을 수는 없을까? 길은 여전히 직선이고, 바뀌는 것은 어느 출발점을 어느 도착점과 짝지을지다.

철새의 항로 배정: 길이 얽히지 않는 출발지–도착지 짝

수천 마리의 새가 동시에 출발할 때 출발점과 도착지를 제비뽑기로 짝지으면 길이 얽힌다. 시베리아 동쪽의 새가 한반도 서쪽으로, 시베리아 서쪽의 새가 한반도 동쪽으로 가면 중간에서 길이 교차한다.

항로를 정해 주는 사람이라면 다르게 짝을 지을 것이다. 시베리아 동쪽의 새는 한반도 동쪽으로, 시베리아 서쪽의 새는 한반도 서쪽으로. 길이 얽히지 않고 모든 새가 나는 거리의 합도 줄어든다. 실제 새가 이렇게 출발지와 도착지를 짝짓는다는 뜻은 아니고, 짝을 정하는 것은 모델을 학습시키는 우리다.

출발점을 도착점으로 옮기는 비용의 합이 가장 작아지도록 짝짓는 문제를 최적 수송(optimal transport, OT)이라 부른다.

얽힌 길이 왜 문제인가

한 점을 여러 (z, x) 짝의 직선이 지나면 그 점의 속도는 그 직선들의 속도를 평균한 값이 된다. 그런데 주변 속도장을 따르는 실제 궤적은 ODE의 해이고 ODE의 해는 한 점에서 속도가 하나뿐이어서, 실제 궤적끼리는 서로 교차할 수 없다. 그래서 조건부 직선들이 얽혀 있으면, 실제 궤적은 교차점 근처에서 평균된 속도를 따라 휜다. 조건부 직선(점선)과 실제 궤적(실선)을 함께 그린 속도장 위젯에서 실선이 휘는 이유가 이것이다. 휜 궤적은 ODE를 적은 스텝으로 풀면 크게 틀린다. 봉우리 일곱 개짜리 데이터에서 직선 조건부 경로로 배운 속도장을 2스텝으로 풀면, 봉우리별 도착 비율이 참 비율에서 총변동(두 비율 표의 차이 절댓값을 더해 반으로 나눈 값) 0.45만큼 어긋난다. 200스텝이면 0.01 남짓이다(「바람의 크기」 절의 파이썬).

노이즈와 데이터의 짝을 잘 지어 조건부 직선들이 덜 얽히면, 평균된 속도도 덜 흔들리고 궤적도 곧아진다. 두 가지 방법이 널리 쓰인다. 미니배치 최적 수송(통(Alexander Tong) 외, 2023; 풀라디안(Aram-Alexandre Pooladian) 외, 2023): 데이터 전체를 한꺼번에 짝지을 수는 없으니, 훈련 배치 안에서만 노이즈와 데이터를 최적 수송으로 짝짓는다. 렉티파이드 플로우(rectified flow, 류싱차오(Xingchao Liu)·궁청웨(Chengyue Gong)·류창(Qiang Liu), 2022): 어느 노이즈가 어느 데이터로 가야 좋은지 먼저 알 길이 없다. 그래서 한 번 학습한 흐름으로 노이즈를 데이터까지 보내고 그 (출발, 도착) 짝으로 다시 학습한다(reflow). 반복할수록 궤적이 곧아져 한 스텝 생성에 가까워진다.

짝짓기의 비용: 바서슈타인 거리 한 줄

노이즈와 데이터의 짝을 잘 지었는지 재려면 비용이 있어야 한다. 최적 수송으로 짝을 짓는 플로우 매칭(OT-CFM)이 보통 쓰는 비용은 출발점과 도착점 사이 거리의 제곱이다. 가능한 모든 짝짓기 가운데 이 비용의 평균이 가장 작은 것을 골라, 그 최솟값의 제곱근을 2-바서슈타인 거리(2-Wasserstein distance)라 부른다.

KL 발산이 「두 분포가 같은 자리에 얼마나 다른 확률을 매기는가」를 묻는다면, W₂ 는 「한 분포를 다른 분포로 옮기려면 확률을 얼마나 멀리 날라야 하는가」를 묻는다. 거리로서의 성질, 비용을 거리로 잴 때와 제곱거리로 잴 때의 차이, 최솟값을 다른 쪽에서 구하는 쌍대 문제는 최적화 책에서 다룬다.

노이즈와 데이터의 짝을 W₂ 최적으로 지으면 샘플 보간 (1 − t)z + tx의 분포는 맥캔(Robert McCann)이 1997년 논문에서 변위 보간(displacement interpolation)이라 부른 것이 된다. W₂가 재는 기하학에서 두 분포를 잇는 측지선(가장 곧은 길)이다. 분포 공간의 두 길(확률을 섞는 m-측지선, 확률을 곱해 섞는 e-측지선)과 나란히 놓이는, 세 번째 종류의 "사이"다. 플로우 매칭이 쓰는 직선은 m-측지선이 아니라, 모양만 닮은 변위 보간의 직선이다.

플로우 매칭에서는 이렇게 정리된다.

• 랜덤 짝 + 직선: 각 새가 제비뽑기한 목적지로 직진한다. 쉽지만 길이 얽힌다.
• OT 짝 + 직선: 새들의 출발-도착 짝을 최적화한다. 길이 덜 얽히고, 궤적이 곧다.

파이썬

import numpy as np
from scipy.optimize import linear_sum_assignment
rng = np.random.default_rng(1)

n = 256
z = rng.normal(size=(n, 2))                                         # 노이즈 출발점
x = rng.normal(size=(n, 2)) * 0.3 + rng.choice([-2.0, 2.0], (n, 1)) * [1, 0]  # 두 덩어리 데이터

rand = rng.permutation(n)                                           # 제비뽑기 짝
_, ot2 = linear_sum_assignment(((z[:, None] - x[None])**2).sum(-1))    # 제곱거리 합이 가장 작은 짝 (W₂)
for name, p in [("랜덤 짝", rand), ("W₂ 최적 짝", ot2)]:
    print(f"{name:8s} 제곱거리 합 {((x[p] - z)**2).sum():7.1f}")
# 랜덤 짝     제곱거리 합  1551.6
# W₂ 최적 짝  제곱거리 합   537.2

출발점과 도착점은 그대로 두고 짝만 바꿨는데 제곱거리 합이 1551.6 에서 537.2 로 3분의 1 남짓이 된다. 미니배치 최적 수송은 이 계산을 훈련 배치마다 한다.

최적 수송 문제의 첫 꼴은 가스파르 몽주(Gaspard Monge)가 1781년 공병의 흙 나르기 문제로 적었고, 레오니트 칸토로비치(Leonid Kantorovich)가 1942년 흙을 여러 곳으로 나눠 보낼 수 있게 고쳐 선형계획으로 풀리는 꼴로 만들었다. 짝짓기 Π를 결합분포로 적는 바서슈타인 거리의 정의식이 칸토로비치의 꼴이다.

이 절의 결론

“OT는 노이즈와 데이터를 잇는 직선 경로는 그대로 두고, 어느 노이즈 점을 어느 데이터 점과 짝지을지를 바꾼다. 짝의 직선들이 덜 얽히면 평균된 속도장이 덜 휘고, ODE를 적은 스텝으로 풀 수 있다. KL 발산이 '확률의 비율’을 보는 잣대라면, 바서슈타인 거리는 '확률의 이동’을 보는 잣대다.”

문제 6. 사물함까지 걷는 네 사람

복도 한 줄 위에 학생 넷이 0, 1, 2, 3 m 자리에 서 있고, 사물함 넷이 4, 5, 6, 7 m 자리에 있다. 네 사람은 동시에 출발해 같은 시각에 자기 사물함에 닿도록 각자 일정한 속력으로 걷는다. (가) 순서대로 짝(0→4, 1→5, 2→6, 3→7)과 거꾸로 짝(0→7, 1→6, 2→5, 3→4)에서 걸은 거리의 합과 거리 제곱의 합을 각각 구하라. (나) 걸린 시간의 절반이 지났을 때 두 짝에서 네 사람은 어디에 있는가? (다) 24가지 짝 가운데 거리의 합이 가장 작은 짝과 거리 제곱의 합이 가장 작은 짝은 각각 몇 개인가?

(위 짝짓기 위젯의 「문제 6 불러오기」로 네 사람과 사물함을 불러와, 시간 슬라이더를 움직이며 자기 계산과 견줄 수 있다.)

함께 풀기

순서대로 짝은 다들 4 m 씩 걸으니까 합이 16 이고, 거꾸로 짝은 7 + 5 + 3 + 1 = 16 이에요. 걷는 거리가 같으니 어느 짝으로 해도 똑같이 좋은 배정이네요.

(나)를 먼저 해 봐요. 절반 시점에 다들 어디 있어요?

순서대로 짝은 2, 3, 4, 5 m 에 한 사람씩요. 거꾸로 짝은… 0 과 7 의 가운데 3.5, 1 과 6 의 가운데 3.5, 나머지도 다 3.5 예요. 네 사람이 한 자리에서 부딪혀요.

거리 제곱의 합으로 재면 순서대로 짝은 16 × 4 = 64, 거꾸로 짝은 49 + 25 + 9 + 1 = 84 야. 제곱으로 재야 순서대로 짝과 거꾸로 짝이 갈려.

(다)는요?

모든 사물함이 모든 학생보다 오른쪽이라 거리는 늘 사물함 위치 − 학생 위치예요.

그걸 다 더하면 짝에 따라 달라져요?

사물함 위치의 합에서 학생 위치의 합을 빼는 거니까… 22 − 6 = 16 이라 짝과 상관이 없어요. 24가지 모두 최소예요. 제곱의 합으로는 순서대로 짝 하나뿐이고요.

제곱은 한 사람이 멀리 가는 걸 무겁게 벌주니까, 다들 비슷한 거리를 가는 배정만 남겨. 그게 길이 엇갈리지 않는 짝이고. 바서슈타인 거리를 왜 제곱거리로 정의하는지 보이네.

3.5 m 에서 만난 네 사람의 속력을 신경망이 하나의 속도로 배워야 한다면요?

속력이 7, 5, 3, 1 이니까 평균 4 요.

한 자리에서 모두 속도 4를 받으면, 그 점을 떠난 네 사람은 다시 제 사물함으로 나뉠 수 있을까요?

못 나뉘어요. 그래서 실제 흐름은 거꾸로 짝을 흉내 내지 못하고 휘어야 하는 거네요. 시험 끝나고 다들 한꺼번에 사물함으로 몰리면 복도 가운데서 엉키는 거랑 같아요.

문제 7. 한 스텝에도 맞는 노이즈–데이터 짝

데이터가 x ~ N(4, 0.5²), 노이즈가 z ~ N(0, 1)인 설정을 다시 본다. 노이즈와 데이터의 짝을 제비뽑기로 지으면 주변 속도장으로 오일러 방법을 한 스텝에 풀었을 때 모든 샘플이 4에 모였다. 이번에는 짝을 순서대로 짓는다. 노이즈 z에는 데이터 x = 4 + 0.5z를 짝짓는다(1차원에서 제곱거리 비용이 가장 작은 짝이다). (가) 순서대로 짓는 짝짓기에서도 z만 모으면 노이즈 분포, x만 모으면 데이터 분포인가? (나) 주변 속도장 v(y, 0)을 구하고, 오일러 한 스텝으로 만든 샘플의 분포를 구하라.

함께 풀기

(가)는 돼요. z가 N(0, 1)이면 4 + 0.5z는 N(4, 0.25)니까요. (나)는 문제 5에서 한 스텝에 다 4로 모였던 계산과 같아요. v(y, 0) = E[x] − y = 4 − y라서 또 전부 4에 모여요.

v(y, 0) = E[x] − y는 어떤 조건에서 나온 식이었죠?

t = 0에서 z가 정해져도 x는 아무것이나 될 수 있다는 조건이었어요. z와 x가 독립이라는 거요.

순서대로 지은 짝에서도 그래요?

아니요. z를 알면 x가 4 + 0.5z 하나로 정해져요. z와 x가 독립이라는 조건이 깨졌어요.

그러면 t = 0에서 위치 y = z에 있는 새의 속도는 x − z = 4 + 0.5y − y = 4 − 0.5y 하나뿐이야. 평균 낼 것도 없어.

한 스텝 가면 y + 4 − 0.5y = 4 + 0.5y, 곧 N(4, 0.25)예요. 데이터 분포 N(4, 0.5²)와 같아요. 코드로 스텝 수를 1, 2, 10으로 바꿔도 표준편차가 다 0.50이에요.

왜 스텝 수와 상관이 없을까요?

순서대로 지은 짝에서 x_t = 4t + (1 − 0.5t)z라서 z가 큰 새는 언제나 위에 있어. 어느 두 새의 길도 엇갈리지 않으니까 한 점을 지나는 (z, x) 짝이 하나뿐이고, 평균된 속도가 그 짝의 직선 속도 4 − 0.5z와 같아. 흐름선 자체가 직선이라 큰 스텝도 빗나가지 않아.

미니배치 최적 수송이 노리는 게 이거네요. 데이터 전체에서 순서대로 짓는 짝은 모르니까 배치 안에서만이라도 맞춰 주는 거고요. 조교가 답안지를 학번순으로 쌓아 두면 성적 입력할 때 엇갈릴 일이 없는 거랑 같아요.