스코어: 흐름을 거슬러 퍼지기 전으로

밀도 전체를 무늬로 나누어 되돌리는 길이 막혔다면, 점 하나하나를 되돌릴 수는 없을까? 그러려면 퍼지는 동안 각 자리의 확률(확률 질량)이 어느 쪽으로 얼마나 빨리 움직이는지 알아야 한다.

연속 방정식과 픽의 법칙

이를 위해 열 방정식을 전혀 다르게 읽어 보자. 퍼지는 양 u를 합이 1인 확률밀도 ρ로 두면, 밀도는 어디서 생기거나 사라지지 않고 옆으로 흘러갈 뿐이다. 각 자리의 질량이 속도 v로 움직인다고 하면 한 점의 밀도는 들어온 흐름에서 나간 흐름을 뺀 만큼 바뀌는데, 이것이 연속 방정식이다. 사람이 오가는 광장에서 한 구역의 인원이 들어온 사람에서 나간 사람을 뺀 만큼만 바뀌는 것과 같다.

이제 흐름이 어느 쪽으로 얼마나 생기는지만 정하면 된다. 푸리에와 픽의 법칙은 흐름의 밀도 ρv가 밀도의 기울기에 비례해 높은 곳에서 낮은 곳으로 향한다고 말한다. 두 식을 합치면 열 방정식이 되고, 흐름의 속도만 따로 떼어 내면 다음과 같다.

\frac{\partial \textcolor{#e377c2}{\rho}}{\partial t} = -\nabla\cdot(\textcolor{#e377c2}{\rho}\,\textcolor{#1b9e77}{\mathbf{v}}), \qquad \textcolor{#e377c2}{\rho}\,\textcolor{#1b9e77}{\mathbf{v}} = -\textcolor{#800080}{D}\,\nabla\textcolor{#e377c2}{\rho} \quad\Longrightarrow\quad \textcolor{#1b9e77}{\mathbf{v}} = -\textcolor{#800080}{D}\,\nabla \ln \textcolor{#e377c2}{\rho}

\begin{array}{ll} \textcolor{#e377c2}{\rho}(\textcolor{#1b9e77}{\mathbf{x}}, \textcolor{#a67a63}{t}) & \text{확률밀도 (합이 1인 퍼지는 양)} \\ \textcolor{#1b9e77}{\mathbf{v}}(\textcolor{#1b9e77}{\mathbf{x}}, \textcolor{#a67a63}{t}) & \text{그 자리의 질량이 움직이는 속도} \\ \textcolor{#e377c2}{\rho}\,\textcolor{#1b9e77}{\mathbf{v}} & \text{흐름의 밀도: 단위 시간에 단위 넓이를 지나는 양 (픽의 법칙)} \\ \nabla\cdot & \text{발산(벡터장의 divergence, KL 발산과 이름만 같다): 한 점에서 나가는 흐름 − 들어오는 흐름} \\ \nabla \ln \textcolor{#e377c2}{\rho} & \text{스코어: 로그 밀도의 위치 기울기} \\ \textcolor{#800080}{D},\ \textcolor{#a67a63}{t} & \text{확산 계수, 시간} \end{array}

∇ln ρ를 확산 모델의 용어로 스코어 (로그 밀도가 가장 가파르게 오르는 방향과 그 가파름, score)라 한다. 통계학에서 스코어는 흔히 매개변수 θ에 대한 로그우도의 기울기를 가리키지만, 확산 모델에서는 위치 x에 대한 기울기를 가리킨다. 식이 말하는 것은 간단하다. 각 자리의 질량은 로그 밀도가 내려가는 쪽으로 움직이고, 밀도가 자기 크기에 비해 가파르게 줄어드는 곳일수록 빨리 움직인다. 분산이 σ²인 가우시안이라면 스코어가 −x/σ²이므로 v = Dx/σ²이고, 원점에서 먼 자리일수록 빨리 바깥으로 밀려난다. 이때 x/σ는 시간이 지나도 변하지 않는다. 가우시안 구름은 모든 점이 표준화한 위치를 지키며 고르게 부풀 뿐이다.

흐름으로 읽은 열 방정식(D = 1/2, 처음 분산 1이라 σ² = 1 + t). 점마다 속도 v = Dx/σ²로 바깥으로 밀려나 먼 점일수록 빠르지만, x/σ는 끝까지 같아 가우시안 구름은 모양을 지킨 채 부푼다.

두 봉우리를 되찾기

각 자리의 질량이 움직이는 속도를 알았으니, 이제 되돌리는 물음으로 돌아가 보자. 이 흐름을 거꾸로 따라가면 퍼지기 전의 분포를 되찾을 수 있을까? 되찾을 수 있다. 어느 시간에서 점들을 v와 반대로, 곧 +D∇ln ρ 쪽으로(ρ는 그 시간의 밀도) 조금씩 옮기면 한 걸음 앞의 분포가 된다. 이 길은 불안정하지 않다. 점 하나하나가 매끄러운 속도장(자리마다 속도가 하나씩 정해진 것)을 따라 움직일 뿐, 반올림 오차 아래로 가라앉은 촘촘한 무늬를 끝 모양에서 되살릴 필요가 없기 때문이다. 대신 모든 시간의 스코어를 알아야 한다. 촘촘한 무늬의 정보는 끝 모양에서는 사라졌지만 이른 시간의 스코어 속에는 남아 있다. ±2에 봉우리를 둔 두 봉우리 분포(각 표준편차 0.5)를 D = 1/2로 t = 8까지 퍼뜨리면 가운데가 가장 높은 한 봉우리가 되는데, 여기서 뽑은 점 2만 개를 정확한 스코어로 t = 0까지 거꾸로 흘려 보내면 |x|의 평균이 1.996, 표준편차가 0.498로 두 봉우리가 그대로 돌아온다. 확산 모델이 거꾸로 가는 길을 찾을 때 열 방정식을 거꾸로 푸는 대신 스코어를 신경망으로 배우는 이유가 여기에 있다.

±2에 봉우리를 둔 두 봉우리 분포(각 표준편차 0.5)를 D = 1/2로 퍼뜨리면 t = 8에 한 봉우리가 된다(왼쪽). t = 8에서 뽑은 점 2만 개를 정확한 스코어를 따라 +D∂ln ρ/∂x 쪽으로 t = 0까지 옮기면 |x|의 평균 1.996, 표준편차 0.498로 두 봉우리가 돌아온다(오른쪽).

코드: 열 방정식을 거꾸로 풀기와 흐름을 거꾸로 따라가기

두 봉우리 분포를 D = 1/2로 t = 8까지 퍼뜨린 뒤 두 방법으로 되돌린다. 첫째는 푸리에 성분마다 e^(Dk²t)를 곱하는 것이고, 둘째는 점들을 스코어를 따라 거꾸로 옮기는 것이다.

import numpy as np

D, T, s0 = 0.5, 8.0, 0.5                     # D = 1/2이면 시간 t가 곧 더해진 분산 2Dt
mus = np.array([-2.0, 2.0])                  # 처음 분포: ±2에 봉우리, 각 표준편차 0.5

def rho(x, t):                               # 열 방정식의 정확한 해: 봉우리마다 분산이 s0² + 2Dt
    v = s0**2 + 2 * D * t
    return sum(np.exp(-(x - m)**2 / (2 * v)) for m in mus) / (2 * np.sqrt(2 * np.pi * v))

def score(x, t):                             # ∂ ln ρ / ∂x (넘침을 막으려고 로그 가중치로 계산)
    v = s0**2 + 2 * D * t
    lw = np.array([-(x - m)**2 / (2 * v) for m in mus]); w = np.exp(lw - lw.max(0)); w /= w.sum(0)
    return ((mus[:, None] - x) * w).sum(0) / v

# (1) 거꾸로 푸는 열 방정식: 푸리에 성분마다 e^(+Dk²T)를 곱해 되돌린다
x = np.linspace(-20, 20, 800, endpoint=False); k = 2 * np.pi * np.fft.fftfreq(800, d=x[1] - x[0])
U_T = np.fft.fft(rho(x, T))
for kc in (2.0, 2.5, 3.0):                   # 파수 kc 위의 성분은 버리고 되돌린다
    keep = np.abs(k) < kc
    back = np.real(np.fft.ifft(np.where(keep, U_T * np.exp(D * np.where(keep, k, 0)**2 * T), 0)))
    print(f"kc = {kc}: x = 2에서 {back[np.argmin(abs(x - 2))]:10.3f} (참값 {rho(2.0, 0):.3f}), 가장 작은 값 {back.min():10.3f}")

# (2) 흐름으로 되돌리기: 입자를 dx/d(−t) = +D ∂ln ρ_t/∂x 로 t = 8에서 0까지 옮긴다 (룽게–쿠타 4차)
rng = np.random.default_rng(0)
xs = mus[rng.integers(0, 2, 20000)] + np.sqrt(s0**2 + 2 * D * T) * rng.standard_normal(20000)
start = xs.copy(); h = T / 4000; t = T
f = lambda x, t: D * score(x, t)
for _ in range(4000):
    k1 = f(xs, t); k2 = f(xs + h / 2 * k1, t - h / 2); k3 = f(xs + h / 2 * k2, t - h / 2); k4 = f(xs + h * k3, t - h)
    xs = xs + h / 6 * (k1 + 2 * k2 + 2 * k3 + k4); t -= h
print(f"t = 8의 표준편차 {start.std():.3f} → t = 0: |x|의 평균 {np.abs(xs).mean():.3f}, |x|의 표준편차 {np.abs(xs).std():.3f}, "
      f"|x| < 1인 비율 {np.mean(np.abs(xs) < 1):.4f}")
print(f"부호가 그대로인 입자 {np.mean(np.sign(xs) == np.sign(start)):.3f}, 순서가 그대로 {np.all(np.argsort(xs) == np.argsort(start))}")
# kc = 2.0: x = 2에서      0.294 (참값 0.399), 가장 작은 값     -0.049
# kc = 2.5: x = 2에서      0.307 (참값 0.399), 가장 작은 값     -0.057
# kc = 3.0: x = 2에서   9244.304 (참값 0.399), 가장 작은 값 -10235.627
# t = 8의 표준편차 3.477 → t = 0: |x|의 평균 1.996, |x|의 표준편차 0.498, |x| < 1인 비율 0.0232
# 부호가 그대로인 입자 1.000, 순서가 그대로 True

첫째 방법은 파수 2.5까지만 살려도 봉우리 높이가 참값의 77%에 그치고 음수 밀도가 생기며, 3까지 살리면 10⁴ 크기로 폭발한다. 둘째 방법은 두 봉우리를 되찾고, 모든 점의 부호와 순서를 지킨다.

같은 t = 8의 밀도에서 되돌린 두 결과. 왼쪽은 푸리에 성분마다 e^(+Dk²t)를 곱하고 파수 2와 2.5 위를 잘라 낸 것으로, 봉우리가 참값 0.399의 74%·77%에 그치고 가운데에 음수 밀도가 생긴다(3까지 살리면 ±10⁴로 폭발해 그리지 않았다). 오른쪽은 점 2만 개를 정확한 스코어를 따라 거꾸로 옮긴 것으로, 히스토그램이 처음 분포(분홍 선)와 겹친다.

ML에서: 분산 폭발형 확산

분산 폭발형(variance exploding, VE) 확산 모델은 데이터 샘플에 표준편차 σ(t)인 정규분포 잡음을 더하고 σ(t)를 시간에 따라 키운다. 잡음 섞인 샘플의 분포는 데이터 분포와 분산 σ²(t)인 가우시안의 합성곱이므로, 열핵의 성질에 따라 확산 계수가 ½dσ²/dt인 열 방정식을 따른다. CIFAR-10(가로세로 32픽셀짜리 컬러 이미지 데이터셋)이라면 픽셀마다 빨강·초록·파랑 세 값이 있어 이미지 하나가 32 × 32 × 3 = 3072개의 숫자이므로, 3072차원 데이터 공간 위의 열 방정식이다.

\textcolor{#1b9e77}{\mathbf{x}}_t = \textcolor{#1b9e77}{\mathbf{x}}_0 + \textcolor{#4c5d80}{\sigma}(\textcolor{#a67a63}{t})\,\textcolor{#1f6066}{\mathbf{z}} \;\Longrightarrow\; \frac{\partial \textcolor{#e377c2}{p_t}}{\partial t} = \frac{1}{2}\frac{d\textcolor{#4c5d80}{\sigma}^2}{dt}\,\nabla^2 \textcolor{#e377c2}{p_t}, \qquad \frac{d\textcolor{#1b9e77}{\mathbf{x}}}{dt} = -\frac{1}{2}\frac{d\textcolor{#4c5d80}{\sigma}^2}{dt}\,\nabla \ln \textcolor{#e377c2}{p_t}(\textcolor{#1b9e77}{\mathbf{x}})

\begin{array}{ll} \textcolor{#1b9e77}{\mathbf{x}}_t & \text{시간 t의 잡음 섞인 샘플 (t는 잡음을 키워 가는 시각)} \\ \textcolor{#1b9e77}{\mathbf{x}}_0,\ \textcolor{#1f6066}{\mathbf{z}} & \text{데이터 샘플, 표준정규 잡음} \\ \textcolor{#4c5d80}{\sigma}(\textcolor{#a67a63}{t}) & \text{잡음의 표준편차 (CIFAR-10에서 0.01부터 50까지)} \\ \textcolor{#e377c2}{p_t} & \text{잡음 섞인 샘플의 분포 = 데이터 분포와 가우시안의 합성곱} \\ \tfrac{1}{2}\,d\textcolor{#4c5d80}{\sigma}^2/dt & \text{확산 계수 D의 자리} \\ \nabla \ln \textcolor{#e377c2}{p_t}(\textcolor{#1b9e77}{\mathbf{x}}) & \text{그 시간 분포의 스코어} \\ d\textcolor{#1b9e77}{\mathbf{x}}/dt & \text{확률 흐름 ODE의 속도 (흐름으로 읽은 열 방정식의 v)} \end{array}

「분산 폭발」이라는 이름은 이 과정이 열 방정식처럼 퍼지기만 할 뿐 어디에서도 멈추지 않아 분산이 σ²(t)만큼 끝없이 커진다는 뜻이다. 송과 동료들(2021)은 오른쪽 식을 확률 흐름 ODE(ODE는 ordinary differential equation의 머리글자로, 변수 하나, 여기서는 시간으로만 미분하는 상미분 방정식)라 부르고 이 길을 따라 결정론적으로 샘플을 만들었는데, 이 식은 흐름으로 읽은 열 방정식의 속도 v = −D∇ln ρ에 D = ½dσ²/dt를 넣은 것이다.

가장 큰 잡음을 데이터 두 점 사이 거리의 최댓값만큼 잡으라는 송과 에르몬의 권고는, 퍼진 분포가 데이터의 봉우리 구조를 잊고 하나의 가우시안과 거의 구별되지 않을 때까지 퍼뜨리라는 뜻이다. 그래야 생성할 때 그 가우시안에서 출발해도 된다. 확률 흐름 ODE를 소개한 송과 동료들(2021)의 논문 요약 첫 문장 「데이터에서 잡음을 만들기는 쉽다. 잡음에서 데이터를 만드는 것이 생성 모델링이다」를, 이 장을 마친 독자는 「앞으로 가는 열 방정식은 무늬를 e^(−Dk²t)로 줄이기만 하니 쉽고, 거꾸로 가는 열 방정식은 e^(Dk²t)로 불안정하니 스코어를 배워 흐름을 거꾸로 따라간다」로 읽게 된다.

문제 10. 스코어를 따라 거꾸로 흐르기

±2에 봉우리를 둔 두 봉우리 분포(각 표준편차 0.5, 문제 9와 같은 분포)를 D = 1/2로 퍼뜨리면 시간 t의 밀도는 분산이 0.25 + t인 두 가우시안의 평균이라 스코어 ∂ln ρ/∂x를 정확히 계산할 수 있다. 이번에는 t = 32까지 퍼뜨린 분포에서 점 2만 개를 뽑는다. (가) 점들을 흐름을 거꾸로 따라 t = 0까지 옮기면 무엇이 되는가? (나) 정확한 스코어 대신, 그 시간의 분포와 평균·분산이 같은 가우시안의 스코어를 쓰면? (다) 점들은 왜 서로 앞지르지 않는가? x = 0에 있던 점은 어디로 가는가? (라) 퍼진 분포를 모른다고 치고 분산이 같은 가우시안에서 점을 뽑아 출발하면 어떻게 되는가? 퍼뜨린 시간 T = 2, 8, 32에서 비교하라. (풀어 본 뒤 위젯 4의 「문제 10 불러오기」로 (가)와 (나)를 확인해 보자.)

t = 32면 표준편차가 6이나 되도록 퍼졌는데, 봉우리가 둘이었다는 흔적이 남아 있기나 할까?

흐름 속도가 v = −D ∂ln ρ/∂x니까 이걸 그대로 쓰고 시간만 32에서 0으로 줄여 가며 옮겼어요. 그런데… 표준편차가 6.0에서 55가 됐어요. 되돌리려던 게 오히려 아홉 배로 퍼졌어요.

시간을 거꾸로 가니까 속도의 부호도 바꿔야지. dx/dt = v를 t가 줄어드는 쪽으로 풀면 한 걸음마다 x에서 v × (시간 간격)을 빼야 하는데, 너는 더했어.

아, 그러네. +D ∂ln ρ/∂x 쪽으로 옮기면… |x|의 평균이 1.997, 표준편차가 0.497이고, |x| < 1인 점이 2.2%야. 처음 분포에서 |x| < 1일 확률이 2.3%니까 두 봉우리가 그대로 돌아왔어!

문제 9에서는 t = 2만 퍼뜨려도 6보다 촘촘한 무늬를 잃었는데, 여기서는 t = 32에서 출발해도 왜 될까요?

문제 9는 t = 2의 밀도 하나에서 반올림 아래로 가라앉은 무늬를 되살려야 했어요. 여기서는 시간마다 스코어를 따로 받아요. 가는 무늬의 정보가 끝 모양에는 없지만, t가 작을 때의 스코어에는 들어 있으니까요. 그래서 얼마나 오래 퍼뜨렸는지는 상관이 없었네요. 흔적은 끝 모양이 아니라 흐름 쪽에 남아 있었어요.

그럼 (나)로 확인해 봐요. 스코어의 어느 부분이 그 정보를 담고 있을까요?

평균·분산이 같은 가우시안의 스코어는 −x/(4.25 + t)예요. 이걸로 옮기면 표준편차 2.05인 한 봉우리가 나와요. |x| < 1인 점이 37%고요. 평균과 분산은 맞는데 두 봉우리는 사라졌어요.

가우시안 스코어로는 모든 점이 같은 비율로 쪼그라들기만 하니까 한 봉우리가 둘로 갈라질 수가 없네. 갈라지게 하는 건 t가 작을 때 스코어가 x = 0 근처의 점들을 양쪽 봉우리로 밀어내는 부분이야. 두 봉우리라는 정보는 스코어의 모양에 있어.

이번에도 점들의 순서와 부호를 확인해 봐요. 점들이 서로 앞지르지 않는다면, 왜 그럴까요?

순서는 한 번도 안 바뀌었어요. 부호도 전부 그대로고요. t = 32에서 오른쪽 절반에 있던 점은 모두 오른쪽 봉우리로 갔어요.

1차원에서 매끄러운 속도장을 따라가는 점들은 서로 추월할 수 없으니까. 두 점이 어느 순간 같은 자리에서 만나면 그 뒤로는 같은 속도로 영원히 붙어 다녀야 하는데, 거꾸로 따라가 보면 처음부터 같은 점이었어야 해.

x = 0에 있던 점은 스코어가 0이라 계속 0에 있겠네요. 근데 t = 32에서는 0이 가장 높은 곳이었는데, t = 0에서는 두 봉우리 사이 골짜기예요.

가장 흔한 자리에서 출발한 점이 가장 드문 자리에 도착하는 거네. 0 근처 점들은 조금만 오른쪽이어도 오른쪽 봉우리로 가 버리니까.

그래요. 흐름은 점을 옮기면서 그 둘레의 넓이를 늘이거나 줄이고, 밀도는 그만큼 바뀌어요. 0 둘레의 좁은 구간이 넓게 늘어나서 골짜기의 낮은 밀도가 되는 거예요.

점 하나의 도착점은 출발점으로 정해지고, 밀도는 넓이가 늘고 주는 만큼 바뀌고… 연속 정규화 흐름에서 로그밀도가 발산만큼 바뀌던 거랑 같은 모양이에요.

마지막으로 (라)요. 실제 생성 모델은 퍼진 분포를 정확히 모르니까 가우시안에서 출발해요.

위젯에서 출발점을 분산이 같은 가우시안으로 바꾸고 정확한 스코어로 되돌리면, T = 2에서는 |x|의 표준편차가 0.573이고 |x| > 3.5인 점이 0.93%예요. 처음 분포에서는 0.13%인데요. T = 8이면 표준편차가 0.508, T = 32면 0.500으로 맞아요. 퍼진 분포와 그 가우시안 사이의 KL이 T = 2에서 0.038, 8에서 0.0013, 32에서 0.000013이에요.

출발점이 틀린 만큼 도착점도 틀리는 거네요. 그러니까 가장 큰 잡음을 데이터 거리의 최댓값만큼 잡으라는 게, 봉우리가 몇 개였는지 잊어버릴 만큼 퍼뜨려야 가우시안에서 출발해도 된다는 뜻이군요.

맞아요. 확산 모델은 앞으로는 열 방정식을 따라 퍼뜨리고, 거꾸로는 배운 스코어로 흐름을 거슬러 올라가요. 앞으로 가는 쪽은 식 하나로 끝나지만, 거꾸로 가는 쪽은 모든 시간의 스코어가 필요하고, 그게 신경망이 배우는 거예요.

같은 방정식인데, 밀도를 거꾸로 풀려고 하면 막히고 흐름을 거꾸로 따라가면 열리네.