디퓨전

디노이징 스코어 매칭: 노이즈를 섞은 기록으로 스코어를 배운다

역방향 SDE에 필요한 새 재료는 각 노이즈 레벨의 스코어 ∇log p_t 하나였지만, 우리는 진짜 스코어를 모르고 손에 있는 것은 이미지들뿐이다. 그렇다면 모르는 스코어를 신경망 s_θ(x, t)에게 어떻게 흉내 내라고 가르칠까? 이미지에 노이즈를 직접 섞어 본 기록만 있으면, 신경망에게 스코어를 가르칠 수 있다.

스코어를 배우는 법: 과거 대회의 기록

진짜 문제는 드래곤볼의 정확한 위치를 모르고, 가진 것이라고는 과거 대회의 기록, 곧 데이터 점에 노이즈를 섞어 본 기록뿐이라는 것이다. 기록 한 줄은 「이 볼에서 출발해 안개 속으로 이만큼 떠밀려 간 사람이 이 지점에 서 있었다」는 내용이다. 그 사람에게 "돌아가는 방향"은 자기가 떠나온 볼(데이터 점) 쪽으로 분명하다. 이 기록들로 신경망을 훈련시켜, 임의의 (지점, 노이즈 레벨) 쌍에서 "스코어(레이더)가 가리켜야 할 방향"을 추정하게 만든다. 기록은 이미지 한 장을 골라 노이즈를 섞기만 하면 한 줄이 생기므로 얼마든지 만들 수 있다.

디노이징 스코어 매칭

기록 하나는 (x₀, z, t) 세 개로 이루어진다. 볼 x₀, 노이즈 z, 노이즈 레벨 t. 그 사람이 서 있던 곳은 x_t = x₀ + σ_tz 이고, 그 한 사람에게 맞는 방향은 가우시안 커널의 스코어(조건부 스코어)다.

여기서 조심할 것이 있다. 타깃 −(x_t − x₀)/σ_t² 는 "이 사람은 볼 x₀에서 왔다"를 알 때의 방향인 조건부 스코어인데, 신경망은 x₀를 보지 못하고 x_t와 t만 본다. 같은 지점에 서로 다른 데이터 점(볼)에서 온 사람들이 겹쳐 서 있으면, 신경망은 그들의 방향을 하나로 답할 수밖에 없다. 제곱 오차를 최소로 하는 하나의 답은 평균(조건부 기댓값)이다.

두 번째 등호가 트위디 공식을 스코어 쪽으로 적은 것이다. 조건부 스코어들의 평균이 주변 스코어(어느 데이터 점에서 왔는지 모르는 채 모두 합친 분포 p_t의 스코어)다. 그래서 진짜 스코어를 한 번도 보지 않고도, 노이즈를 섞어 본 기록만으로 진짜 스코어를 배운다. 이 학습법을 디노이징 스코어 매칭이라고 부른다.

파이썬

위 그림의 숫자를 샘플로 확인하는데, 데이터는 다시 봉우리가 둘인 가장 단순한 분포, −1과 +1 두 점(반반)이다.

import numpy as np
rng = np.random.default_rng(0)

# 디노이징 스코어 매칭의 최소해 = 주변 스코어
# 데이터 ±1, 노이즈 σ=0.5. 조건부 타깃 −(x_σ − x)/σ² 를 x_σ 구간별로 평균낸다.
sigma, n = 0.5, 4_000_000
x = rng.choice([-1.0, 1.0], n)
xs = x + sigma * rng.normal(size=n)
target = -(xs - x) / sigma**2
for y in [0.0, 0.3, 1.0]:
    near = np.abs(xs - y) < 0.02
    print(f"y={y}: 조건부 타깃 평균 {target[near].mean():6.3f}   "
          f"주변 스코어 {(np.tanh(y / sigma**2) - y) / sigma**2:6.3f}")
# y=0.0: 조건부 타깃 평균 -0.028   주변 스코어  0.000
# y=0.3: 조건부 타깃 평균  2.112   주변 스코어  2.135
# y=1.0: 조건부 타깃 평균 -0.003   주변 스코어 -0.003

y = 0.3의 조건부 타깃은 +2.8(+1에서 온 사람)과 −5.2(−1에서 온 사람) 둘뿐인데, 평균은 2.1이다. 구간 폭 0.02 때문에 생기는 작은 편향을 빼면 주변 스코어와 맞는다.

ML에서: 실제 모델은 노이즈를 맞힌다

이미지 디퓨전 모델은 보통 스코어 대신 섞인 노이즈 자체를 맞히도록 학습한다. 호(Jonathan Ho) 외(2020)의 DDPM은 디노이징 스코어 매칭을 노이즈 z를 맞히는 문제 ‖ε_θ − z‖²로 쓴다. ε_θ는 신경망이 추정한 노이즈이고, 이 노이즈 추정값과 스코어 사이는 ε_θ = −σ_t s_θ 로 옮겨진다. 노이즈 레벨마다 노이즈 예측 오차 ‖ε_θ − z‖²에 어떤 비중을 줄지는 정하기 나름인데, 호 외는 이론이 말하는 비중을 떼어 내고 모든 수준을 같은 비중으로 더한 단순한 손실이 더 보기 좋은 그림을 만든다는 것을 실험으로 보였다. 노이즈 추정값을 스코어나 원래 이미지의 추정값으로 바꾸는 계산은 아래 문제 13에서 한다.

역사: 트위디 공식을 학습 목표로

노이즈 낀 점에서 조건부 방향을 평균하면 주변 스코어가 된다는 관계는 통계학의 트위디 공식에서 먼저 쓰였다. 이 관계를 학습 목표로 쓴 것은 아포 히바리넨(Aapo Hyvärinen)의 스코어 매칭(2005)과 파스칼 뱅상(Pascal Vincent)의 디노이징 스코어 매칭(2011)이다.

호기심 상자: 모든 노이즈 레벨을 합치면 KL로 돌아온다

디노이징 스코어 매칭의 사영은 무엇으로 재는가

제곱 오차의 최솟값이 조건부 기댓값이라는 사실은 "x_t의 함수들이 이루는 공간 위로의 직교사영"이라는 기하학적 사실이다. 큰 공간의 대상을 작은 부분공간의 가장 가까운 점으로 옮긴다는 점에서 분포를 모델 면 위로 옮기는 사영과 모양이 같지만, 가까움을 재는 기준이 다르다. 분포를 모델 면 위로 옮기는 두 사영, 곧 KL(q‖p)을 줄이는 e-사영과 KL(p‖q)을 줄이는 m-사영은 KL 발산(두 분포가 같은 사건에 매기는 확률이 얼마나 다른지 재는 양)으로 가까움을 잰다. 스코어 매칭은 두 스코어(벡터장)의 제곱 차이를 p_t로 평균 낸 양, 피셔 발산으로 잰다.

$$\begin{aligned} \textcolor{#c2398a}{D}_{F}(\textcolor{#d6479f}{p}_t \,\|\, \textcolor{#7f8f10}{q}_t) = {} & \mathbb{E}_{\textcolor{#d6479f}{p}_t} \big\| \nabla_{\textcolor{#6f8fa6}{x}} \log \textcolor{#d6479f}{p}_t(\textcolor{#6f8fa6}{x}) \\ & \quad - \nabla_{\textcolor{#6f8fa6}{x}} \log \textcolor{#7f8f10}{q}_t(\textcolor{#6f8fa6}{x}) \big\|^2 \end{aligned}$$

$$\begin{array}{ll} \textcolor{#c2398a}{D}_{F} & \text{피셔 발산} \\ \textcolor{#d6479f}{p}_t & \text{노이즈를 더한 데이터 분포} \\ \textcolor{#7f8f10}{q}_t & \text{신경망 스코어가 정하는 모델 분포} \\ \textcolor{#6f8fa6}{x} & \text{데이터 공간의 점} \end{array}$$

그러면 피셔 발산은 KL과 아무 관계가 없을까 싶지만, 둘 사이에는 아주 깊은 관계가 있다. p와 q에 똑같이 노이즈를 더해 가면 둘은 점점 구분이 안 되고 KL은 줄어드는데, 그 KL이 줄어드는 속도가 정확히 피셔 발산의 절반이다.

$$\begin{aligned} \tfrac{d}{d\textcolor{#0068c0}{\tau}}\, \textcolor{#c2398a}{\mathrm{KL}}(\textcolor{#d6479f}{p}_\tau \| \textcolor{#7f8f10}{q}_\tau) &= -\tfrac{1}{2} \textcolor{#c2398a}{D}_{F}(\textcolor{#d6479f}{p}_\tau \| \textcolor{#7f8f10}{q}_\tau) \\ \textcolor{#c2398a}{\mathrm{KL}}(\textcolor{#d6479f}{p} \| \textcolor{#7f8f10}{q}) &= \tfrac{1}{2} \int_0^\infty \textcolor{#c2398a}{D}_{F}(\textcolor{#d6479f}{p}_\tau \| \textcolor{#7f8f10}{q}_\tau)\, d\textcolor{#0068c0}{\tau} \end{aligned}$$

$$\begin{array}{ll} \textcolor{#0068c0}{\tau} & \text{노이즈의 분산 } \sigma^2 \\ \textcolor{#c2398a}{\mathrm{KL}} & \text{KL 발산} \\ \textcolor{#c2398a}{D}_{F} & \text{피셔 발산} \\ \textcolor{#d6479f}{p}_\tau,\ \textcolor{#7f8f10}{q}_\tau & \text{노이즈를 더한 두 분포} \end{array}$$

드 브루인 항등식(엔트로피(분포의 평균 놀라움)가 노이즈에 따라 변하는 속도가 피셔 정보(스코어 크기의 제곱을 평균 낸 값)라는 식)의 KL판이다(시웨이 류(Siwei Lyu), 2009). 두 번째 식은 첫 번째 식을 τ = 0부터 ∞까지 적분한 것이고, 노이즈가 무한히 크면 둘 다 같은 넓은 정규분포에 가까워져 KL이 0이다. 모든 노이즈 레벨에서 스코어 매칭을 알맞은 비중으로 합하면 KL을 줄이는 일, 다시 말해 최대우도(모델이 데이터에 주는 확률을 가장 크게 하기)가 된다. 노이즈 레벨 하나의 스코어 매칭은 KL로 재는 사영과 가까움의 기준이 다르지만, 모든 노이즈 레벨을 합치면 같은 KL로 돌아온다.

파이썬

import numpy as np
rng = np.random.default_rng(0)

# 드 브루인 꼴: 노이즈가 커질 때 KL 이 줄어드는 속도 = −½ × 피셔 발산
# p = N(0,1), q = N(1,4) 에 똑같이 분산 τ 의 노이즈를 더한다.
def kl(m1, v1, m2, v2):
    return 0.5 * (np.log(v2 / v1) + (v1 + (m1 - m2)**2) / v2 - 1)
tau, h, n = 0.5, 1e-4, 4_000_000
dkl = (kl(0, 1 + tau + h, 1, 4 + tau + h) - kl(0, 1 + tau - h, 1, 4 + tau - h)) / (2 * h)
y = rng.normal(0, np.sqrt(1 + tau), n)
fisher_div = np.mean((-y / (1 + tau) + (y - 1) / (4 + tau))**2)
print(f"dKL/dτ = {dkl:.4f}   −½·피셔 발산 = {-0.5 * fisher_div:.4f}")
# dKL/dτ = -0.1728   −½·피셔 발산 = -0.1728

변분추론으로 보기: DDPM의 ELBO

호 외(2020)는 같은 모델을 변분추론으로 유도했다. 데이터의 로그 우도가 넘을 수 없는 아래 경계인 ELBO(증거 하한)를 올리는 방식이다. 순방향 과정 x₀ → x₁ → … → x_T를 잠재변수(직접 관측되지 않고 모델 안에만 있는 변수)로 보고, 역방향 모델 q_θ(x_t−1 | x_t)를 학습한다. (문헌에서는 순방향을 q, 모델을 p_θ로 쓰지만, 이 책은 p = 현실, q = 모델을 지키므로 이름을 바꿨다.)

$$\begin{aligned} & \log \textcolor{#7f8f10}{q}_{\textcolor{#2e9e6e}{\theta}}(\textcolor{#6f8fa6}{x}_0) \ge \textcolor{#b58a00}{\mathrm{ELBO}} \\ & \textcolor{#b58a00}{\mathrm{ELBO}} = \mathbb{E}\Big[\log \textcolor{#7f8f10}{q}_{\textcolor{#2e9e6e}{\theta}}(\textcolor{#6f8fa6}{x}_0 \mid \textcolor{#6f8fa6}{x}_1) \\ & \quad - \textstyle\sum_{t=2}^{T} \textcolor{#c2398a}{\mathrm{KL}}\big(\textcolor{#d6479f}{p}(\textcolor{#6f8fa6}{x}_{t-1} \mid \textcolor{#6f8fa6}{x}_t, \textcolor{#6f8fa6}{x}_0) \\ & \qquad\quad \,\|\, \textcolor{#7f8f10}{q}_{\textcolor{#2e9e6e}{\theta}}(\textcolor{#6f8fa6}{x}_{t-1} \mid \textcolor{#6f8fa6}{x}_t)\big) \\ & \quad - \textcolor{#c2398a}{\mathrm{KL}}\big(\textcolor{#d6479f}{p}(\textcolor{#6f8fa6}{x}_T \mid \textcolor{#6f8fa6}{x}_0) \,\|\, \textcolor{#139696}{\pi}(\textcolor{#6f8fa6}{x}_T)\big)\Big] \end{aligned}$$

$$\begin{array}{ll} \textcolor{#7f8f10}{q}_{\textcolor{#2e9e6e}{\theta}} & \text{역방향 모델} \\ \textcolor{#2e9e6e}{\theta} & \text{신경망의 가중치} \\ \textcolor{#b58a00}{\mathrm{ELBO}} & \text{증거 하한} \\ \textcolor{#d6479f}{p}(x_{t-1} \mid x_t, x_0) & \text{데이터 점 } x_0 \text{ 를 알 때} \\ \textcolor{#c2398a}{\mathrm{KL}} & \text{KL 발산} \\ \textcolor{#139696}{\pi} & \text{출발 분포 } N(0, I) \\ \textcolor{#6f8fa6}{x}_0, \dots, \textcolor{#6f8fa6}{x}_T & \text{마르코프 체인} \end{array}$$

DDPM의 ELBO는 잠재변수의 사후분포를 정확히 못 구하니 하한인 ELBO를 올린다는 점에서 변분추론과 같은 구조다. 가운데 KL 항들은 정규분포끼리의 KL이라 평균의 차이만 남고, 정리하면 앞의 디노이징 스코어 매칭(노이즈 맞히기)에 시간별 비중을 붙인 것이 된다. 위 「ML에서」 소절에서 호 외가 떼어 낸 비중이 이 시간별 비중이다. 킹마(Diederik Kingma)와 가오(Ruiqi Gao)(2023)는 흔히 쓰는 디퓨전 손실들이 모두 "노이즈 레벨별 ELBO에 비중을 붙여 적분한 것"으로 읽힌다는 것을 보였다. KL = ½∫D_F dτ 적분식이 말한 「노이즈 레벨을 모두 합치면 KL」을 변분추론 쪽에서 다시 말한 것이다.

대응표

이 절의 결론

“디퓨전 모델은 마법이 아니다. 데이터 분포에 노이즈를 더해 뭉갰다가, 스코어 함수가 가리키는 방향을 따라 노이즈를 거꾸로 걷어 낸다(매니폴드로의 귀환). 그 스코어는 노이즈를 섞어 본 기록으로 만든 조건부 스코어의 평균이고, 모든 노이즈 레벨을 합치면 최대우도로 돌아온다.”

문제 11. 우산을 몇 개 들여놓을까

편의점 점장이 내일 팔릴 우산 수를 하나의 숫자로 예측해 본사에 보내야 한다. 내일 비가 올 확률은 30%이고 비가 오면 100개, 안 오면 20개가 팔리는데, 본사는 보낸 숫자와 실제 판매량의 차이를 제곱해서 점수를 깎는다. (가) 어떤 숫자를 보내야 깎이는 점수의 기댓값이 가장 작은가? (나) 확률이 더 큰 경우(비가 안 옴)의 판매량 20개를 보내면 기댓값으로 얼마나 깎이는가?

함께 풀기

비가 안 올 확률이 70%니까 20개죠. 더 그럴듯한 경우를 고르는 게 맞잖아요.

20개를 보냈는데 비가 오면 얼마나 깎여요?

80개 차이의 제곱이니까 6,400이요. 30% 확률이니 기댓값으로 1,920이에요.

제곱 오차의 기댓값을 가장 작게 하는 숫자는 평균이야. 0.3 × 100 + 0.7 × 20 = 44개. 깎이는 기댓값은 0.3 × 56² + 0.7 × 24² = 1,344라서 20개보다 작아.

44개는 실제로는 절대 안 나오는 판매량인데요. 그래도 그게 제일 덜 깎이네요. 20개처럼 비 오는 날 크게 틀리는 예측보다, 44개처럼 양쪽으로 적당히 틀리는 예측이 제곱 오차에서는 이득이구나.

답이 둘 중 하나로 갈리는데 어느 것인지 모르면, 제곱 오차는 두 답을 확률로 섞은 값을 고르게 해요.

문제 12. 조건부 타깃과 주변 스코어

데이터가 두 점 ±1(반반)이고 표준편차 σ = 0.5인 노이즈를 더한다. (가) x_σ = −0.2에 선 사람의 조건부 타깃 −(x_σ − x₀)/σ² 는 어떤 값들을 가질 수 있는가? (나) 디노이징 스코어 매칭을 완벽하게 풀었을 때 신경망이 x = −0.2에서 내놓는 값은?

함께 풀기

(가)는 x₀가 +1이면 −(−0.2 − 1)/0.25 = 4.8, −1이면 −(−0.2 + 1)/0.25 = −3.2예요. (나)는 우산 문제처럼 두 타깃을 섞으면 되겠죠. 데이터 점이 반반이니까 0.5 × 4.8 + 0.5 × (−3.2) = 0.8이에요.

신경망은 이 사람이 어느 데이터 점에서 왔는지 보고 답해요?

아니요, x = −0.2만 봐요. 그러니까 섞는 거고요.

그럼 −0.2에 서 있는 사람 가운데 +1에서 온 사람과 −1에서 온 사람이 정말 반반일까요? 우산 문제에서 섞는 비율은 무엇이었어요?

내일 비가 올 확률이었어요. 여기서는 처음 데이터 점이 반반이었던 게 아니라, −0.2에 서 있다는 걸 본 뒤의 확률로 섞어야 하네요.

그 확률은 어느 쪽이 더 커요?

−0.2는 −1에 더 가까우니까 −1에서 온 사람이 더 많아요. +1에서 왔을 사후확률이 (1 + tanh(−0.2/0.25))/2 = 0.168이니까 0.168 × 4.8 + 0.832 × (−3.2) = −1.856이에요.

주변 스코어 (tanh(y/σ²) − y)/σ²에 넣어도 −1.856이야. 트위디 공식을 거꾸로 쓴 거네. 0.8이었으면 스코어가 +1 쪽을 가리켰을 텐데, 실제 스코어 −1.856은 가까운 점 −1 쪽을 가리켜.

팀 과제에서 누가 쓴 문단인지 모르는 채로 「이 문단 고칠 방향」을 한 줄만 적어야 하면, 팀원 각자가 원하는 방향을 그럴듯한 만큼 섞어 적는 게 제일 덜 틀리잖아요. 팀원이 데이터 점, 섞은 한 줄이 주변 스코어네요.

그래요. 한 사람에게 맞는 방향과 신경망이 배우는 방향은 달라요.

문제 13. 노이즈 추정값을 스코어로

이미지 디퓨전 모델처럼 노이즈를 맞히도록 학습한 신경망이 있다. 데이터는 두 점 ±1(반반), 노이즈 σ = 0.5이고, 신경망은 완벽하게 학습되었다. 노이즈 낀 점 x_σ = −0.6을 넣었더니 섞인 노이즈의 추정값으로 ε̂ = 0.767을 내놓았다. (가) 이 점 x_σ = −0.6의 스코어는? (나) 원래 점의 평균 위치 E[x₀ | x_σ = −0.6]은? (다) (가)의 스코어와 (나)의 평균 위치가 주변 스코어 (tanh(y/σ²) − y)/σ², 사후 평균 tanh(y/σ²)과 맞는지 확인하라.

함께 풀기

x_σ = x₀ + σz니까 원래 점은 노이즈를 더하면 돼요. −0.6 + 0.5 × 0.767 = −0.216이요.

노이즈 추정값이 양수예요. 섞인 노이즈가 이 점을 어느 쪽으로 밀어 놓았다는 뜻이에요?

오른쪽으로요. 그러면 원래 점은 왼쪽에 있었어야 하니까 빼야 하네요. x₀ = x_σ − σε̂ = −0.6 − 0.384 = −0.984예요. 식을 x₀에 대해 풀면서 부호를 놓쳤어요.

(가)는 ε̂ = −σs에서 s = −0.767/0.5 = −1.535예요. 주변 스코어로는 (tanh(−2.4) + 0.6)/0.25 = −1.535로 같아요.

그럼 트위디로 원래 점을 구해도 같아?

−0.6 + 0.25 × (−1.535) = −0.984. (나)와도 맞아. tanh(−2.4) = −0.984이기도 하고.

노이즈를 맞히든, 스코어를 맞히든, 원래 점을 맞히든 ε̂ = −σs, x₀ = x + σ²s 로 서로 바뀌는 같은 정보구나. 문제 4(운이 섞인 시험 점수)에서 「운」을 추정하면 「실력」의 추정값도 정해졌던 것처럼요.

그래요. 실제 모델이 셋 가운데 무엇을 맞히느냐는 학습이 얼마나 안정적이냐의 문제이지, 배우는 정보가 다른 게 아니에요.