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

디퓨전에서는 매 스텝에 무작위 노이즈를 더하며 걷는 역방향 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

바람의 크기 g 를 0, 1, 2로 바꿔도 봉우리별 도착 비율의 총변동은 0.010~0.014로 거의 같다. 총변동은 두 비율 표의 차이 절댓값을 더해 반으로 나눈 값이고, 0이면 두 표가 똑같다. 반면 ODE를 몇 스텝으로 푸느냐는, 조건부 길은 직선이어도 주변 궤적이 휘어 있기 때문에 결과를 크게 바꾼다. 노이즈와 데이터의 짝을 최적 수송으로 짓는 방법과, 한 번 배운 흐름의 (출발, 도착) 짝으로 다시 배우는 reflow가 줄이려는 것이 이 ODE 스텝 수다.

이 절의 결론

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

문제 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로 돌리고, 도착 비율과 봉우리 안의 퍼짐을 예측과 비교하라.

봉우리 하나에 작용하는 두 항: 바람(g dW)은 퍼뜨리고 스코어 항은 봉우리 쪽으로 모은다. 스코어 항의 계수 c를 ½과 1로 놓아 본다

함께 풀기

디퓨전의 역방향 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에는 없고 역방향 SDE에는 있는 게 뭐였죠?

바람이요! 역방향 SDE의 g²는 확률 플로우 ODE에 이미 든 ½g²에, 바람을 넣으면서 더하는 ½g²를 합친 거예요.

우리 v는 그 ODE의 표류 전체에 해당하니까…

응, 바람을 넣으면서 더하는 스코어 항으로는 ½g²만 더해야 해. 너는 확률 플로우 ODE에 이미 든 ½g² 스코어 항을 두 번 넣은 거야.

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

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

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

½g² 스코어 항과 바람의 확산 항이 맞물리는 건, 조교가 채점 기준을 고칠 때 가산점과 감점을 같은 크기로 맞춰야 반 평균이 안 바뀌는 것과 같네요. 감점을 두 번 넣으면 점수가 한쪽으로 몰리고요.