바람의 크기: 같은 흐름을 지나는 샘플러 가족

디퓨전에서는 매 걸음에 무작위 흔들림을 더하며 걷는 역방향 SDE도, 흔들림 없이 걷는 확률흐름 ODE도 같은 분포에 닿았다. 이 무작위 흔들림을 바람이라고 불렀다. 생성 모델을 실제로 돌릴 때는 같은 신경망으로 어느 샘플러를 쓸지 골라야 한다. 바람 없이 매끈하게 흐를 것인가, 바람을 맞으며 걸을 것인가. 속도장만 배운 플로우 매칭에서도 바람을 넣을 수 있을까? 같은 주변분포 p_t의 흐름을, 바람 없이 흐를 수도 있고 바람을 맞으며 걸을 수도 있다.

바람의 크기를 고르는 식

디퓨전의 역방향 SDE와 확률흐름 ODE는 같은 p_t를 지난다. 플로우 매칭의 시간 방향(0 = 노이즈)으로 쓰면, 그 사이에 한 줄로 늘어선 가족이 있다.

\begin{aligned} d\textcolor{#6f8fa6}{x} = &\Big[\textcolor{#d07a12}{v}(\textcolor{#6f8fa6}{x}, \textcolor{#007800}{t}) + \frac{\textcolor{#0068c0}{g}(\textcolor{#007800}{t})^2}{2}\, \textcolor{#b0892a}{\nabla_{x} \log p_t}(\textcolor{#6f8fa6}{x})\Big]\, d\textcolor{#007800}{t} \\ &+ \textcolor{#0068c0}{g}(\textcolor{#007800}{t})\, d\textcolor{#8a7fb0}{W} \end{aligned}

\begin{array}{ll} \textcolor{#6f8fa6}{x} & \text{새의 위치} \\ \textcolor{#d07a12}{v} & \text{주변 속도장} \\ & \text{(플로우 매칭이 배운 것)} \\ \textcolor{#b0892a}{\nabla_{x} \log p_t} & \text{같은 경로의 주변 스코어} \\ \textcolor{#0068c0}{g}(\textcolor{#007800}{t}) & \text{바람의 크기} \\ & \text{(아무 값이나 고를 수 있다)} \\ d\textcolor{#8a7fb0}{W} & \text{브라운 운동의 증분, } dt > 0 \end{array}

g = 0이면 플로우 매칭의 ODE다. g > 0이면 바람이 퍼뜨리는 만큼 스코어 항이 도로 모아서, 매 순간의 분포는 여전히 p_t다. 포커-플랑크 방정식(분포가 시간에 따라 어떻게 흐르는지 적은 식)에 넣어 보면 g² 항이 서로 지워지는 것을 확인할 수 있다(아래 문제 10). 그리고 속도와 스코어를 잇는 다리 공식 v = x/t + ((1 − t)/t)∇log p_t 덕분에, v만 배워도 스코어를 얻으니 이 가족 전체를 쓸 수 있다.

데이터가 −1.5 와 +1.5 근처의 두 봉우리일 때, 같은 출발점 여덟 개에서 흐른 길. 위: 바람 없음(g = 0, ODE). 길이 매끈하다. 아래: 바람 g = 1(SDE). 길이 들쭉날쭉하다. 오른쪽 끝의 막대는 2000 개를 흘려 보낸 도착 위치의 분포로, 두 경우가 같은 두 봉우리를 그린다

파이썬

import numpy as np

mu = np.array([[-2.5, -1.5], [1.5, -2.0], [-1.0, 2.0], [2.5, 1.5],
               [0.0, 0.0], [-2.0, 0.5], [1.0, -0.5]])            # 봉우리 일곱 개의 중심
w = np.array([.18, .14, .16, .14, .12, .13, .13]); s0 = 0.4      # 봉우리마다의 혼합 비중, 봉우리 폭(표준편차)

def fields(y, t):
    """선형 경로 x_t=(1−t)z+tx 의 주변 속도 v 와 스코어 ∇log p_t (닫힌꼴)."""
    V = (1 - t)**2 + t**2 * s0**2                      # 성분별 분산
    d = t * mu[None] - y[:, None]
    lg = np.log(w) - (d**2).sum(-1) / (2 * V)
    g = np.exp(lg - lg.max(1, keepdims=True)); g /= g.sum(1, keepdims=True)
    score = (g[..., None] * d).sum(1) / V
    Ex = (g[..., None] * (mu[None] + t * s0**2 / V * (y[:, None] - t * mu[None]))).sum(1)
    return (Ex - y) / (1 - t), score

def sample(gs, N=200, n=4000, seed=1):
    r = np.random.default_rng(seed)
    y = r.normal(size=(n, 2))                          # 출발: z ~ N(0, I)
    ts = np.linspace(0, 0.999, N + 1)                  # t = 1 에서는 (Ex − y)/(1 − t) 가 0/0 이 되어 계산이 멈추므로 바로 앞에서 멈춘다
    for i in range(N):
        t, dt = ts[i], ts[i + 1] - ts[i]
        v, sc = fields(y, t); g = gs * (1 - t)
        y = y + (v + 0.5 * g**2 * sc) * dt + g * np.sqrt(dt) * r.normal(size=y.shape)
    k = np.argmin(((y[:, None] - mu[None])**2).sum(-1), 1)
    return np.bincount(k, minlength=7) / n

print("참 비중", w)
for gs in [0.0, 1.0, 2.0]:
    f = sample(gs)
    print(f"g={gs}    ", f.round(2), f"총변동 {0.5 * abs(f - w).sum():.3f}")
for N in [2, 4, 8]:
    f = sample(0.0, N=N)
    print(f"ODE {N}스텝", f.round(2), f"총변동 {0.5 * abs(f - w).sum():.3f}")
# 참 비중 [0.18 0.14 0.16 0.14 0.12 0.13 0.13]
# g=0.0     [0.18 0.15 0.16 0.13 0.13 0.13 0.12] 총변동 0.014
# g=1.0     [0.18 0.15 0.16 0.13 0.12 0.13 0.13] 총변동 0.010
# g=2.0     [0.18 0.14 0.16 0.14 0.13 0.13 0.14] 총변동 0.014
# ODE 2스텝 [0.05 0.02 0.08 0.02 0.47 0.14 0.23] 총변동 0.454
# ODE 4스텝 [0.12 0.08 0.13 0.07 0.25 0.15 0.19] 총변동 0.217
# ODE 8스텝 [0.15 0.12 0.15 0.1  0.18 0.14 0.16] 총변동 0.102

바람의 크기를 바꿔도 도착 비율은 그대로다. 총변동은 두 비율 표의 차이 절댓값을 더해 반으로 나눈 값이고, 0이면 두 표가 똑같다. 반면 ODE를 몇 걸음으로 푸느냐는 결과를 크게 바꾼다. 조건부 길은 직선이어도 주변 궤적이 휘어 있기 때문이다. 짝을 최적 수송으로 짓는 방법과, 한 번 배운 흐름의 (출발, 도착) 짝으로 다시 배우는 reflow가 겨냥하는 것이 바로 이 걸음 수다.

수확

“바람의 크기는 우리가 고르는 샘플러의 선택이다. 바람이 흩뜨리는 만큼 스코어가 도로 모아 주면, 한 마리의 길은 달라져도 매 순간의 분포는 같다.”

문제 10. ½은 어디서 왔나

바람을 넣은 샘플러 dx = [v + (g²/2)∇log p_t]dt + g dW 의 분포 p는 포커-플랑크 방정식 ∂p/∂t = −∇·(p × 표류) + (g²/2)∇²p 를 따른다. 표류는 dt 앞의 대괄호이고, ∇²는 두 번 미분이다(1차원이면 ∂²/∂x²). (가) p∇log p = ∇p 를 써서 g² 항이 지워지고 연속방정식 ∂p/∂t + ∇·(pv) = 0 만 남음을 보여라. (나) 스코어 항의 계수 ½을 1로 바꾸면 무엇이 남는가? 위 파이썬에서 0.5를 1로 바꿔 g = 1, 2로 돌리고, 도착 비율과 봉우리 안의 퍼짐을 예측과 견주라.

함께 풀기

디퓨전의 역방향 SDE는 표류에 스코어를 g²배 통째로 넣었잖아요. 그러니까 ½이 오히려 이상해 보여요. (나)부터 할게요. 1로 바꿔 돌렸더니 총변동이 g = 1이면 0.066, 2면 0.142예요. ½일 때 0.010, 0.014보다 커졌네요.

봉우리 하나하나는 어떻게 생겼어요?

도착점마다 가장 가까운 봉우리 중심에서 떨어진 정도를 재는 줄을 더했어요. 한 방향 표준편차로 ½이면 0.39인데, 1이면 0.35, 0.32예요. 참 봉우리 폭은 0.4니까 봉우리가 너무 좁게 뭉쳐요.

그런데 역방향 SDE엔 정말 g²가 통째로 들어 있었는데요?

디퓨전의 확률흐름 ODE 표류에는 스코어가 몇 배 들어 있었죠?

½g²예요. 아, 역방향 SDE의 g²는 확률흐름 ODE에 이미 든 ½g²에, 바람을 넣으면서 더하는 ½g²를 합친 거예요. 우리 v는 그 ODE의 표류 전체에 해당하니까, 바람 몫으로는 ½만 더해야 해요. 민준이는 같은 몫을 두 번 넣은 거야.

(가)를 식으로 할게요. 표류의 스코어 항이 만드는 부분은 −∇·(p · ½g²∇log p) = −½g²∇·∇p = −½g²∇²p예요. 바람의 확산 항 +½g²∇²p와 정확히 지워져서 ∂p/∂t = −∇·(pv)만 남아요. 연속방정식이에요.

계수가 1이면 −g²∇²p + ½g²∇²p = −½g²∇²p가 남아요. 퍼지는 방정식을 거꾸로 돌린 꼴이라 분포를 뾰족하게 모아요. 민준이가 잰 대로 봉우리가 좁아져요.

그래요. ½은 바람이 퍼뜨리는 양과 스코어가 모으는 양을 정확히 맞춘 숫자예요.

조교가 채점 기준을 고칠 때 가산점과 감점을 같은 크기로 맞춰야 반 평균이 안 바뀌는 거랑 같네요. 감점을 두 번 넣으면 점수가 한쪽으로 몰리고요.