스코어 ∇log p: '노이즈를 예측한다'가 곧 스코어 추정인 이유

확산 모델의 손실은 한 줄입니다.

loss = F.mse_loss(model(x_t, t), noise)

모델이 하는 일은 "이 흐릿한 이미지에 섞인 잡음이 무엇이었나"를 맞히는 것뿐입니다. 그런데 학습이 끝나면 이 모델로 새 이미지를 만들어 냅니다. 잡음 맞히기 연습이 어떻게 이미지의 분포를 배우는 일이 되는지가 처음에는 전혀 안 보입니다.

같은 시기에 나온 다른 계보의 모델은 아예 다른 것을 학습한다고 말했습니다. 데이터 분포의 로그밀도의 기울기 $\nabla_x \log p(x)$ 를 맞히고, 그 기울기를 따라 올라가며 표본을 만든다는 것이었습니다. 그런데 두 계보의 코드를 나란히 놓으면 손실이 상수배만큼만 다릅니다.

이 글은 그 상수배를 유도합니다. 지난 글이 손실의 뼈대를 세웠다면, 여기서는 그 뼈대에서 나온 목적식의 정체가 무엇인지를 봅니다. U-Net 구조와 DDPM 구현은 확산 모델의 기초가 다루고, 이 글은 목적식만 맡습니다.

스코어 — 밀도 대신 기울기

1차원 봉우리에서 읽기

스코어 함수(score function)는 로그밀도를 데이터에 대해 미분한 것입니다.

값의 뜻은 그래디언트 그대로입니다 — 그 자리에서 밀도가 가장 빠르게 커지는 방향과 그 가파름입니다.

봉우리에서는 화살표가 0이고, 꼬리에서는 안쪽을 가리키며 길어집니다. 이 화살표장만 알면 어디서 출발하든 밀도가 높은 쪽으로 따라 올라갈 수 있습니다. 표본을 만드는 일이 결국 그것이라, 밀도 자체가 없어도 생성이 됩니다.

1차원에서는 부호와 크기 둘뿐이라 읽기가 쉽습니다. 봉우리 왼쪽이면 양수(오른쪽으로 밀림), 오른쪽이면 음수(왼쪽으로 밀림), 꼭대기에서 0입니다. 밀도의 기울기가 아니라 로그밀도의 기울기라는 점만 조심하면 됩니다 — 로그를 씌우면 꼬리 쪽에서 값이 눌리지 않아, 밀도가 거의 0인 자리에서도 화살표가 살아 있습니다. 생성이 아무 데서나 출발해도 되는 것이 이 성질 덕입니다.

두 봉우리의 화살표장

봉우리가 둘이면 화살표장이 골짜기를 경계로 갈립니다. 왼쪽 영역의 점은 왼쪽 봉우리로, 오른쪽 영역의 점은 오른쪽 봉우리로 밀립니다. 두 봉우리가 대칭이면 정확한 가운데에서 양쪽 힘이 상쇄되어 화살표가 0이 되는데, 이것은 봉우리가 아니라 골짜기의 꼭대기입니다.

스코어가 0인 자리가 봉우리와 골짜기를 구별하지 못한다는 것은 나중에 표본을 만들 때 실제로 문제가 됩니다 — 뒤에서 랑주뱅 갱신을 세울 때 이 자리로 다시 옵니다.

통계학의 스코어와의 차이

이름이 같아 헷갈리는 자리가 하나 있습니다. 통계학에서 스코어라 하면 보통 로그가능도를 파라미터로 미분한 $\nabla_\theta \log p_\theta(x)$ 를 가리킵니다. 최대가능도 추정에서 이 값이 0이 되는 $\theta$ 를 찾고, 그 분산이 피셔 정보량입니다.

생성 모델에서 쓰는 것은 데이터 $x$ 로 미분한 쪽입니다. 미분하는 변수가 다르니 뜻도 다릅니다 — 파라미터 미분은 "모수를 어느 쪽으로 옮기면 이 표본이 더 그럴듯해지나"이고, 데이터 미분은 "이 자리에서 어느 쪽으로 옮기면 더 그럴듯한 표본이 되나"입니다. 이 글에서 스코어라고 하면 언제나 뒤쪽입니다.

사라지는 정규화 상수

적분이 막는 자리

스코어를 쓰는 결정적인 이유는 따로 있습니다. 밀도를 신경망으로 모형화하려면 넓이가 1이어야 한다는 제약을 지켜야 합니다. 보통 이렇게 씁니다.

$\tilde{p}_\theta$ 는 넓이를 신경 쓰지 않고 아무렇게나 만든 양수 함수이고, $Z_\theta$ 는 그것을 1로 맞추는 정규화 상수입니다. 그런데 $x$ 가 이미지면 이 적분은 픽셀 수만큼의 차원에서 도는 적분이라 계산할 방법이 없습니다. 밀도를 직접 다루는 모형이 어려운 이유가 이것입니다.

미분이 지우는 자리

스코어를 취하면 그 벽이 사라집니다.

$Z_\theta$ 는 $x$ 와 무관한 상수라 $x$ 로 미분하면 그냥 없어집니다.

숫자로도 한 줄이면 확인됩니다.

import numpy as np
mu, s = 1.0, 2.0
logp      = lambda x: -0.5*np.log(2*np.pi*s*s) - 0.5*((x-mu)/s)**2   # 정규화된 것
logp_tilde = lambda x: -0.5*((x-mu)/s)**2 + 12.345                   # 상수 아무거나
h, x = 1e-5, 0.7
print((logp(x+h)-logp(x-h))/(2*h), (logp_tilde(x+h)-logp_tilde(x-h))/(2*h), -(x-mu)/s**2)

0.07499999999938112 0.07499999998827889 0.07500000000000001

세 값이 같습니다. 마지막 값은 정규분포의 스코어 공식인데, 지수부만 미분하면 바로 나옵니다.

정규분포의 스코어는 평균으로 되돌아가는 벡터이고, 크기는 분산에 반비례합니다. 이 한 줄이 뒤에서 전부를 결정합니다.

그 대신 포기하는 것

공짜는 아닙니다. $Z_\theta$ 를 안 구하고 넘어간 대가는 밀도값 자체를 영영 모른다는 것입니다. 스코어는 기울기라, 적분해서 $\log \tilde p$ 까지는 복원해도 거기서 $\log Z$ 만큼의 차이는 끝내 남습니다.

그래서 못 하는 일이 셋 있습니다.

이상탐지처럼 "이 표본의 가능도가 얼마인가"를 수로 답해야 하는 일이 특히 걸립니다. 되찾으려면 스코어장을 시간축으로 적분하는 확률흐름 ODE를 따로 풀어야 하고, 그것이 다음 글의 주제 중 하나입니다. 생성만 할 것이라면 치를 필요 없는 값이라, 확산 모델은 기꺼이 이 대가를 냅니다.

디노이징 스코어 매칭

정답을 모른다는 문제

목표는 신경망 $s_\theta(x)$ 가 데이터 분포의 스코어를 맞히게 하는 것입니다. 자연스러운 손실은 이것입니다.

그런데 이 식은 쓸 수 없습니다. 정답인 $\nabla_x \log p_{\text{data}}$ 를 모르기 때문입니다. 그것을 알면 애초에 학습할 이유가 없습니다.

우리가 섞은 잡음

빠져나가는 길이 디노이징 스코어 매칭(denoising score matching)입니다. 아이디어는 이렇습니다 — 데이터의 스코어는 모르지만, 데이터에 우리가 직접 섞은 가우시안 잡음의 스코어는 정확히 압니다. 앞 절의 한 줄이 그것입니다.

재매개변수화로 세운 폐형을 그대로 씁니다.

$x_0$ 를 고정하면 이것은 평균 $\sqrt{\bar{\alpha}_t}x_0$, 분산 $1-\bar{\alpha}_t$ 인 정규분포이므로 스코어가 공식으로 나옵니다.

그리고 분자는 정의상 $\sqrt{1-\bar{\alpha}_t}\,\varepsilon$ 입니다. 대입하면

우리가 섞은 그 잡음이, 부호를 뒤집고 상수로 나눈 것이 곧 조건부 스코어입니다. 잡음은 우리가 뽑아서 알고 있으니 정답이 손에 있습니다.

조건부에서 주변으로

그런데 방금 얻은 것은 $x_0$ 를 아는 상태의 스코어입니다. 우리가 원하는 것은 $x_0$ 를 모르는 상태, 즉 잡음 낀 데이터 전체의 분포 $q(x_t)$ 의 스코어입니다. 둘은 다른 값입니다.

다리를 놓는 것은 MSE의 성질 하나입니다. 어떤 함수 $f$ 로

를 최소화하면 답은 언제나 조건부 평균 $f^\star(x_t) = \mathbb{E}[Y \mid x_t]$ 입니다. 회귀가 하는 일이 원래 그것입니다.

여기에 $Y = -\varepsilon/\sqrt{1-\bar{\alpha}_t}$ 를 넣습니다. 그러면 학습이 수렴한 신경망은

가 되고, 디노이징 스코어 매칭 정리는 이 조건부 평균이 정확히 $\nabla_{x_t} \log q(x_t)$ 라고 말합니다.

말로 하면 이렇습니다. 같은 $x_t$ 자리에는 서로 다른 $x_0$ 에서 서로 다른 잡음을 타고 온 표본들이 섞여 있습니다. 그중 하나만 보면 그 잡음은 주변 분포의 스코어와 전혀 다른 값입니다. 그런데 그 자리에 모인 것들을 평균하면 개별 사연이 지워지고 주변 분포의 기울기만 남습니다.

수치로 확인하기

직접 재 보면 보입니다. $x_0$ 가 $-2$ 또는 $+2$ 인 데이터에 $\bar{\alpha} = 0.6$ 만큼 잡음을 섞고, 같은 $x_t$ 자리에 모인 표본들의 $-\varepsilon/\sqrt{1-\bar{\alpha}}$ 를 평균했습니다.

xt = np.sqrt(ab)*x0s + np.sqrt(1-ab)*np.random.randn(n)   # x0s 는 −2 또는 +2
target = -(xt - np.sqrt(ab)*x0s) / (1 - ab)               # = −eps / sqrt(1−ab)
for xq in [-2.0, -0.5, 0.0, 0.8, 2.0]:
    m = np.abs(xt - xq) < 0.02                            # 그 자리에 모인 표본만
    print(f"x_t={xq:>5}  표본평균 {target[m].mean():8.4f}   해석적 {qscore(xq):8.4f}")

x_t= -2.0  표본평균   1.1266   해석적   1.1270
x_t= -0.5  표본평균  -2.4715   해석적  -2.4652
x_t=  0.0  표본평균   0.0141   해석적   0.0000
x_t=  0.8  표본평균   1.8595   해석적   1.8572
x_t=  2.0  표본평균  -1.1269   해석적  -1.1270

봉우리 자리($x_t = \pm 2$ 근처)에서는 안쪽으로 되돌리는 값이 나오고, 두 봉우리 사이 골짜기($x_t = -0.5$)에서는 가까운 봉우리 쪽으로 강하게 밀어냅니다. $x_t = 0$ 은 두 봉우리의 정확한 가운데라 힘이 상쇄되어 0입니다. 400만 표본으로 잰 값이 해석적 스코어와 소수 둘째 자리까지 맞습니다.

세 목표가 하나인 이유

ε 예측과 스코어의 상수배

이제 처음의 손실 한 줄로 돌아갑니다.

$\varepsilon_\theta$ 가 $\varepsilon$ 을 맞히도록 학습하는데, 방금 본 것은 $-\varepsilon/\sqrt{1-\bar{\alpha}_t}$ 가 스코어라는 사실이었습니다. $t$ 가 정해지면 $\sqrt{1-\bar{\alpha}_t}$ 는 그냥 숫자이므로

로 서로 옮겨 갈 수 있습니다. 두 손실은 $t$ 마다 상수를 곱한 것만 다르고, 최적해는 같은 함수입니다.

x̂₀ 예측이라는 다른 옷

같은 이유로 세 번째 후보도 같은 것입니다. 폐형을 $x_0$ 에 대해 풀면

이라 "깨끗한 원본을 맞히기"도 같은 함수의 다른 옷입니다. 논문마다 $\varepsilon$ 예측, $x_0$ 예측, $v$ 예측을 고르는 것은 다른 것을 배우는 것이 아니라 손실에 $t$ 마다 어떤 가중치를 붙일지를 고르는 것입니다. 그 가중치가 학습 초반의 안정성과 어느 $t$ 에 노력이 몰리는지를 바꿉니다.

v-예측의 정의와 쓰임

방금 이름만 대고 지나간 네 번째 옷을 적어 둡니다. v-예측은 다음 양을 맞히는 방식입니다.

$\varepsilon$ 과 $x_0$ 를 같은 계수로 섞되 부호를 엇갈리게 둔 것입니다. $x_t = \sqrt{\bar\alpha_t}x_0 + \sqrt{1-\bar\alpha_t}\varepsilon$ 과 나란히 놓으면 계수 행렬이 회전행렬 꼴이라, $x_t$ 와 $v$ 를 알면 $x_0$ 와 $\varepsilon$ 이 되돌려집니다.

굳이 이 옷을 입히는 이유는 $t$ 의 양 끝에서 안정적이기 때문입니다. $t$ 가 0에 가까우면 $\bar\alpha_t \approx 1$ 이라 $x_t$ 가 거의 $x_0$ 이고, $x_0$ 예측은 입력을 그대로 베끼면 되는 시시한 문제가 됩니다. 반대로 $t$ 가 $T$ 에 가까우면 $\bar\alpha_t \approx 0$ 이라 $x_t$ 가 거의 잡음이고, 이번에는 $\varepsilon$ 예측이 베끼기가 됩니다. 한쪽이 시시해지는 구간에서 다른 쪽은 $1/\sqrt{\bar\alpha_t}$ 같은 큰 수로 나뉘어 손실이 폭발합니다.

$v$ 는 두 항을 $\sqrt{\bar\alpha_t}$ 와 $\sqrt{1-\bar\alpha_t}$ 로 섞어 두어 어느 끝에서도 한쪽 항만 남지 않습니다. 계수 제곱합이 1이라 크기도 일정하게 유지됩니다. 증류로 스텝 수를 크게 줄이는 방법들이 $v$ 를 쓰는 것이 이 때문입니다.

스코어로 표본 만들기

화살표를 따라 올라가는 갱신

여기까지 "화살표장만 있으면 따라 올라가면 된다"고만 하고 방법을 안 줬습니다. 실제 갱신식은 한 줄입니다.

이것을 랑주뱅 동역학(Langevin dynamics)이라고 합니다. 스코어 방향으로 $\epsilon$ 만큼 한 걸음 가고, 거기에 $\sqrt{2\epsilon}$ 배로 키운 가우시안 잡음을 한 번 더 얹는 것이 전부입니다. $\epsilon$ 이 충분히 작고 걸음이 충분히 많으면 이 반복이 만들어 내는 점들의 분포가 $p$ 자체로 수렴합니다. 밀도값은 한 번도 안 쓰고 화살표만 씁니다.

잡음 항 없이 갇히는 자리

$\sqrt{2\epsilon}z$ 를 빼면 남는 것은 로그밀도에 대한 경사상승법입니다. 그러면 어디로 가는지는 정해져 있습니다 — 출발점에서 가장 가까운 봉우리로 올라가 멈춥니다.

이것은 표본이 아니라 최빈값입니다. 사람 얼굴을 학습했다면 매번 "가장 평균적인 얼굴" 한 장만 나오고, 봉우리가 둘인 데이터에서는 출발점이 어느 쪽 골짜기에 떨어졌는지가 결과를 다 정합니다. 앞에서 스코어가 0인 자리가 봉우리와 골짜기를 구별 못 한다고 했는데, 잡음이 없으면 골짜기 꼭대기에 정확히 놓인 점은 영영 거기 머무르기까지 합니다.

잡음 항은 그 두 가지를 한꺼번에 고칩니다. 봉우리에서도 계속 흔들리므로 한 점에 고이지 않고, 이따금 골짜기를 넘어 다른 봉우리로 건너갑니다. 얼마나 자주 건너가는지가 봉우리 높이의 비율과 맞아떨어지는 것이 이 갱신식의 요점입니다 — $\sqrt{2\epsilon}$ 라는 계수가 정확히 그 비율을 맞추는 값입니다.

확산 모델의 역과정으로

실제 확산 모델의 역과정은 이 갱신을 그대로 쓰지 않고, $t$ 를 $T$ 에서 0으로 내리면서 잡음 크기를 스케줄에 따라 줄여 가며 돕니다. 처음에는 큰 잡음으로 넓게 훑어 어느 봉우리로 갈지 정하고, 뒤로 갈수록 잡음을 줄여 그 봉우리 안에서 세부를 맞추는 식입니다. 잡음이 처음부터 작으면 가까운 봉우리에 갇히고, 끝까지 크면 세부가 안 잡힙니다.

$\epsilon$ 을 얼마로 두는지, 몇 걸음에 나눠 걷는지가 곧 품질이 되는데 — 그 걸음이 연속 시간의 방정식을 이산화한 것이라서 그렇습니다. 다음 글이 그 자리를 맡습니다.

처음의 질문에는 이제 답이 됩니다. 잡음 맞히기가 분포를 배우는 일인 이유는, 맞힌 잡음이 그 자리에서 밀도가 높아지는 방향을 가리키는 화살표와 상수배 관계이기 때문입니다.

연습 문제

답은 문항을 눌러 펼칩니다. 연습 2는 $\bar{\alpha}_t = 0.36$ 이라 $\sqrt{\bar{\alpha}_t} = 0.6$, $\sqrt{1-\bar{\alpha}_t} = 0.8$ 인 자리에서 풉니다.

연습 1 — 스코어를 정의에서 구하기

1. $\mu = -1$, $\sigma = 3$ 인 정규분포의 로그밀도를 $x$ 로 미분해 스코어를 구하고, $x = 2$ 에서의 값을 적으세요.

   $\log p(x) = -\log(\sqrt{2\pi}\,\sigma) - \dfrac{(x+1)^2}{18}$ 이고 앞항은 $x$ 와 무관하므로 $s(x) = -\dfrac{x+1}{9}$ 입니다. $x=2$ 에서 $-\dfrac{3}{9} = -0.333$ 이라, 평균 $-1$ 쪽인 왼쪽으로 밀립니다.
2. 세 함수 $\tilde p_1 = e^{-x^2/2}$, $\tilde p_2 = 5e^{-x^2/2}$, $\tilde p_3 = 0.01e^{-x^2/2}$ 의 스코어를 각각 구해 비교하세요.

   $\log \tilde p_i = \log c_i - \dfrac{x^2}{2}$ 이고 $\log c_i$ 는 $x$ 와 무관하므로 셋 다 $-x$ 로 같습니다. 상수배는 로그에서 덧셈 상수가 되고 미분에서 지워집니다 — 정규화 상수가 사라지는 것과 정확히 같은 이유입니다.

연습 2 — 세 예측을 서로 옮겨 적기

1. $x_t = 1.0$ 에서 신경망이 $\varepsilon_\theta = 0.5$ 를 내놓았습니다. 같은 자리의 $s_\theta$ 는 얼마인가요.

   $s_\theta = -\dfrac{\varepsilon_\theta}{\sqrt{1-\bar\alpha_t}} = -\dfrac{0.5}{0.8} = -0.625$ 입니다.
2. 같은 값에서 $\hat{x}_0$ 를 구하세요.

   $\hat{x}_0 = \dfrac{x_t - \sqrt{1-\bar\alpha_t}\,\varepsilon_\theta}{\sqrt{\bar\alpha_t}} = \dfrac{1 - 0.8\times0.5}{0.6} = \dfrac{0.6}{0.6} = 1.0$ 입니다.
3. 이번에는 $s_\theta = -1.25$ 가 주어졌습니다. 대응하는 $\varepsilon_\theta$ 와 $\hat{x}_0$ 를 구하고, $\hat{x}_0$ 를 두 공식으로 각각 구해 값이 같은지 확인하세요.

   $\varepsilon_\theta = -s_\theta\sqrt{1-\bar\alpha_t} = 1.25\times0.8 = 1.0$ 입니다. 스코어 쪽 공식으로는 $\hat x_0 = \dfrac{1 + 0.64\times(-1.25)}{0.6} = \dfrac{0.2}{0.6} = 0.333$ 이고, $\varepsilon$ 쪽 공식으로는 $\dfrac{1 - 0.8\times1.0}{0.6} = \dfrac{0.2}{0.6} = 0.333$ 으로 같습니다. 두 공식이 같은 함수의 다른 옷이라는 것이 이렇게 확인됩니다.

연습 3 — MSE의 최소해

1. $\mathbb{E}\big[(f(x) - Y)^2\big]$ 를 최소로 하는 $f$ 가 $f(x) = \mathbb{E}[Y \mid x]$ 임을 보이세요. (자리마다 따로 최소화해도 된다는 것부터 씁니다.)

   기대값을 $x$ 로 먼저 쪼개면 $\mathbb{E}_x\big[\mathbb{E}[(f(x)-Y)^2 \mid x]\big]$ 이고, 안쪽 값은 자리 $x$ 마다 따로 정해지므로 각 자리에서 따로 최소화하면 전체가 최소가 됩니다. 그 자리에서 $f(x)$ 는 상수 $c$ 이므로 $\mathbb{E}[(c-Y)^2\mid x]$ 를 $c$ 로 미분해 $2\big(c - \mathbb{E}[Y\mid x]\big) = 0$, 즉 $c = \mathbb{E}[Y \mid x]$ 입니다. 두 번 미분하면 2로 양수라 최소가 맞습니다. 본문에서 학습이 수렴한 신경망이 조건부 평균으로 간다고 한 근거가 이 세 줄입니다.

정리

• 스코어 함수는 $\nabla_x \log p(x)$, 파라미터가 아니라 데이터로 미분한 로그밀도의 기울기다. 밀도가 커지는 방향을 가리키고 봉우리에서 0이다.
• 정규화 상수가 미분에서 사라진다. $\nabla_x \log(\tilde{p}/Z) = \nabla_x \log \tilde{p}$ 이므로 계산 불가능한 고차원 적분을 피할 수 있다. 스코어를 쓰는 결정적 이유다.
• 정규분포의 스코어는 $-(x-\mu)/\sigma^2$ — 평균으로 되돌아가는 벡터이고 크기는 분산에 반비례한다.
• 데이터의 스코어는 모르지만 우리가 섞은 잡음의 스코어는 정확히 안다. 그것이 디노이징 스코어 매칭의 출발점이다.
• 폐형 forward에서 $\nabla_{x_t}\log q(x_t \mid x_0) = -\varepsilon / \sqrt{1-\bar{\alpha}_t}$ 다. 정답이 손에 있으므로 회귀 문제가 된다.
• MSE의 최소해는 조건부 평균이므로, 수렴한 신경망은 $\mathbb{E}[-\varepsilon/\sqrt{1-\bar{\alpha}_t} \mid x_t]$ 를 내놓는다. 그 값이 주변 분포의 스코어 $\nabla \log q(x_t)$ 다.
• 표본 하나의 $-\varepsilon$ 은 스코어와 전혀 다른 값이다. 같은 $x_t$ 에 모인 것들을 평균해야 개별 사연이 지워지고 기울기만 남는다. 두 봉우리 데이터로 재 보니 400만 표본이 해석적 값과 소수 둘째 자리까지 일치했다.
• $t$ 를 고정하면 $\sqrt{1-\bar{\alpha}_t}$ 는 상수이므로 $\varepsilon$ 예측·스코어 예측·$x_0$ 예측은 같은 함수의 다른 옷이다. 다른 것은 $t$ 마다의 손실 가중치뿐이고, 그것이 안정성을 가른다.
• v-예측 $v = \sqrt{\bar\alpha_t}\varepsilon - \sqrt{1-\bar\alpha_t}x_0$ 는 두 항을 회전행렬 꼴로 섞어, $t$ 의 어느 끝에서도 한쪽이 베끼기가 되거나 손실이 폭발하지 않게 한다.
• 표본은 랑주뱅 갱신 $x \leftarrow x + \epsilon s(x) + \sqrt{2\epsilon}z$ 로 만든다. 잡음 항을 빼면 가장 가까운 봉우리로 올라가 멈추므로 표본이 아니라 최빈값이 나온다.
• 정규화 상수를 포기한 대가는 밀도값을 못 준다는 것이다. 생성은 되지만 가능도 비교와 bits-per-dim 보고는 확률흐름 ODE를 따로 풀어야 한다.

읽어주셔서 감사합니다. 😊