개념정리 : Diffusion 모델 학습 Loss : 노이즈 예측부터 Score Matching까지

이제는 자연스럽게 학습 loss가 왜 저 모양이 되는지, 그리고 그게 왜 score matching과 연결되는지를 이어서 설명드리겠습니다.

지금까지의 흐름을 한 줄로 요약하면 이렇습니다.

• forward는 우리가 정확히 안다
• $x_0$를 안다고 가정하면, 이상적인 one-step reverse도 정확히 안다
• 실제 생성에서는 $x_0$를 모르므로, 신경망이 그 빠진 정보를 예측해야 한다
• 그 예측 타깃으로 $x_0$나 $\epsilon$을 쓸 수 있다

이제 남은 질문은 딱 두 개입니다.

1. 그 신경망을 어떤 loss로 학습시키는가?
2. 왜 그 loss가 score matching과 연결되는가?

이 두 개를 순서대로 보겠습니다.

---

1. 학습의 최종 목표는 무엇인가

생성모델의 최종 목표는 결국 데이터 같은 샘플을 잘 만들도록 하는 것입니다.

확률적으로 말하면, 모델이 정의하는 분포 $p_\theta(x_0)$가 실제 데이터 분포에 가까워지도록 하고 싶습니다.

그런데 diffusion에서는 $x_0$만 바로 모델링하지 않고,

중간 latent chain 전체를 같이 둡니다.

즉 모델의 joint distribution은 보통

$$p_\theta(x_{0:T}) = p(x_T)\prod_{t=1}^T p_\theta(x_{t-1}\mid x_t)$$

로 둡니다.

여기서

• $p(x_T)$: 거의 표준가우시안
• $p_\theta(x_{t-1}\mid x_t)$: 학습할 reverse step

입니다.

즉 모델은 전체적으로

“가우시안 노이즈에서 시작해서, step-by-step으로 데이터를 복원하는 확률과정”

을 정의합니다.

---

2. 그런데 $p_\theta(x_0)$를 직접 최대화하기는 어렵다

왜냐하면 $x_1,\dots,x_T$라는 latent 변수를 적분해야 하기 때문입니다.

$$p_\theta(x_0)=\int p_\theta(x_{0:T})dx_{1:T}$$

이 적분은 직접 다루기 어렵습니다.

그래서 diffusion에서는 forward process $q(x_{1:T}\mid x_0)$를 이용해 variational bound를 씁니다.

핵심 아이디어는:

• forward는 우리가 안다
• reverse는 학습해야 한다
• 그러니 forward를 기준분포처럼 써서 reverse를 학습하자

입니다.

전체 ELBO를 전개하면 여러 항으로 나뉘는데,

지금 단계에서 제일 중요한 항은 각 time step마다 등장하는

$$\mathrm{KL}\big(q(x_{t-1}\mid x_t,x_0) \parallel p_\theta(x_{t-1}\mid x_t)\big)$$

입니다.

즉 각 step에서 모델이 해야 할 일은

이상적인 reverse posterior

$$q(x_{t-1}\mid x_t,x_0)$$

를

학습된 reverse model

$$p_\theta(x_{t-1}\mid x_t)$$

가 잘 흉내 내도록 만드는 것

입니다.

이게 loss의 출발점입니다.

---

3. 왜 KL이 등장하는가

KL divergence는 “한 분포가 다른 분포와 얼마나 다른가”를 재는 양입니다.

$$\mathrm{KL}(q \parallel p) = \mathbb{E}_q\left[\log \frac{q}{p}\right]$$

입니다.

여기서는

$$q(x_{t-1}\mid x_t,x_0)$$

가 이상적인 정답 분포이고,

$$p_\theta(x_{t-1}\mid x_t)$$

가 우리가 만든 모델 분포입니다.

그러므로 KL을 줄인다는 것은

“모델의 reverse step이 정답 reverse step과 비슷해지도록 한다”

는 뜻입니다.

즉 loss의 본질은 아주 직관적입니다.

정답 가우시안 posterior를 모델 가우시안이 잘 따라가게 하는 것입니다.

---

4. 둘 다 가우시안이면 KL은 어떻게 생기나

이제 중요한 단순화가 나옵니다.

우리는 이상적인 posterior를

$$q(x_{t-1}\mid x_t,x_0) = \mathcal{N}\big(\tilde\mu_t(x_t,x_0),\tilde\beta_t I\big)$$

로 알고 있습니다.

그리고 모델도 보통

$$p_\theta(x_{t-1}\mid x_t) = \mathcal{N}\big(\mu_\theta(x_t,t),\sigma_t^2 I\big)$$

처럼 둡니다.

특히 분산을 고정해서 $\sigma_t^2=\tilde\beta_t$처럼 두면,

두 분포는 같은 분산을 가진 가우시안이 됩니다.

이 경우 KL은 아주 단순해집니다.

일반적으로

$$q=\mathcal{N}(m_1,\Sigma),\qquad p=\mathcal{N}(m_2,\Sigma)$$

이면

$$\mathrm{KL}(q \parallel p) = \frac{1}{2} (m_1-m_2)^\top \Sigma^{-1}(m_1-m_2)$$

입니다.

만약 $\Sigma=\sigma^2 I$이면

$$\mathrm{KL}(q \parallel p) = \frac{1}{2\sigma^2}|m_1-m_2|^2$$

가 됩니다.

즉 같은 분산을 가진 두 가우시안의 KL은 평균끼리의 제곱오차와 같다는 뜻입니다.

이게 diffusion loss가 MSE 꼴로 나오는 첫 번째 이유입니다.

---

5. 그러면 결국 무엇을 맞추면 되나

이제 바로 결론이 나옵니다.

reverse KL을 줄인다는 것은 결국

$$\mu_\theta(x_t,t) \approx \tilde\mu_t(x_t,x_0)$$

즉 모델이 예측하는 reverse mean이

이상적인 posterior mean에 가까워지도록 만들면 된다는 뜻입니다.

그래서 학습의 핵심은 사실상

reverse mean matching

입니다.

문제는, 모델이 mean을 직접 예측하게 할 수도 있지만,

앞에서 봤듯이 $x_0$나 $\epsilon$을 예측해서 mean을 구성하는 것도 가능합니다.

이제 바로 그 연결을 보겠습니다.

---

6. $\epsilon$-prediction으로 loss가 왜 MSE가 되는가

앞에서 이상적인 posterior mean을 $\epsilon$로 쓰면

$$\tilde\mu_t(x_t,\epsilon) = \frac{1}{\sqrt{\alpha_t}} \left( x_t-\frac{\beta_t}{\sqrt{1-\bar\alpha_t}}\epsilon \right)$$

라고 했습니다.

그리고 모델이 $\epsilon$를 예측한다고 하면,

$$\hat\epsilon = \epsilon_\theta(x_t,t)$$

이고, 모델 mean은

$$\mu_\theta(x_t,t) = \frac{1}{\sqrt{\alpha_t}} \left( x_t-\frac{\beta_t}{\sqrt{1-\bar\alpha_t}}\epsilon_\theta(x_t,t) \right)$$

로 정의할 수 있습니다.

이제 두 mean의 차이를 빼보겠습니다.

$$\tilde\mu_t-\mu_\theta = \frac{1}{\sqrt{\alpha_t}} \left( -\frac{\beta_t}{\sqrt{1-\bar\alpha_t}}\epsilon + \frac{\beta_t}{\sqrt{1-\bar\alpha_t}}\epsilon_\theta \right)$$

정리하면

$$\tilde\mu_t-\mu_\theta = \frac{\beta_t}{\sqrt{\alpha_t}\sqrt{1-\bar\alpha_t}} (\epsilon_\theta-\epsilon)$$

입니다.

즉 mean error는 결국 true noise와 predicted noise의 차이에 비례합니다.

그러므로 KL은

$$\mathrm{KL} \propto |\tilde\mu_t-\mu_\theta|^2 \propto |\epsilon-\epsilon_\theta(x_t,t)|^2$$

가 됩니다.

더 정확히 쓰면 time-dependent weight가 붙어서

$$L_t = \mathbb{E}\left[ \frac{\beta_t^2}{2\sigma_t^2\alpha_t(1-\bar\alpha_t)} |\epsilon-\epsilon_\theta(x_t,t)|^2 \right]$$

꼴이 됩니다.

이 식이 중요한 이유는 분명합니다.

reverse KL을 최소화하는 문제가

결국 노이즈 예측 MSE 문제로 바뀝니다.

즉 복잡한 확률모델 학습이, 놀랍게도

“네트워크가 섞인 노이즈를 맞혀라”라는 supervised regression 문제로 바뀌는 것입니다.

이게 diffusion이 학습하기 편한 큰 이유 중 하나입니다.

---

7. 그런데 실제로는 왜 더 단순한 loss를 많이 쓰나

실전에서는 방금의 가중치까지 그대로 쓰지 않고,

보통 더 단순한

$$L_{\text{simple}} = \mathbb{E}_{x_0,\epsilon,t} \big[ |\epsilon-\epsilon_\theta(x_t,t)|^2 \big]$$

를 많이 씁니다. 여기서

$$x_t=\sqrt{\bar\alpha_t}x_0+\sqrt{1-\bar\alpha_t}\epsilon$$

입니다.

즉 학습은 매우 단순합니다.

1. 깨끗한 데이터 $x_0$를 하나 뽑고
2. 시간 $t$를 뽑고
3. 가우시안 노이즈 $\epsilon$를 하나 샘플링해서 $x_t$를 만들고
4. 네트워크에 $x_t,t$를 넣어서 $\epsilon$를 맞히게 한다

이 과정만 반복하면 됩니다.

---

이 loss의 직관적 의미

이 loss는 사실상 이렇게 해석할 수 있습니다.

“현재 noisy sample을 보고, 그 안에 섞여 있는 noise component를 알아맞혀라.”

그러면 sampling할 때는 그 noise 방향을 조금씩 제거하면서

더 깨끗한 샘플 쪽으로 이동할 수 있습니다.

즉 training과 sampling의 직관이 매우 잘 맞습니다.

• 학습: 들어간 노이즈를 맞힌다
• 샘플링: 예측한 노이즈를 제거한다

---

8. $x_0$-prediction이면 loss는 어떻게 되나

비슷한 논리로 $x_0$-prediction도 가능합니다.

모델이

$$\hat x_0 = x_\theta(x_t,t)$$

를 예측한다고 하면, 이상적인 posterior mean의 $x_0$ 자리에 $\hat x_0$를 넣으면 됩니다.

그러면 mean error는 결국

$$|\hat x_0 - x_0|^2$$

에 비례하는 꼴로 바뀝니다.

즉 시간가중치가 붙은 $x_0$ MSE를 쓰는 셈입니다.

그래서 $x_0$-prediction과 $\epsilon$-prediction은 본질적으로 아주 가깝습니다.

차이는 어느 표현공간에서 regression을 하느냐입니다.

다만 $\epsilon$-prediction은 타깃이 표준가우시안 스케일이라서

보통 학습 안정성이 좋고 구현도 편합니다.

---

9. 이제 score matching과 연결해보자

여기서 아주 중요한 다리가 나옵니다.

앞에서 원본 $x_0$를 알고 있을 때의 조건 score를 구했었습니다.

$$\nabla_{x_t}\log q(x_t\mid x_0) = -\frac{x_t-\sqrt{\bar\alpha_t}x_0}{1-\bar\alpha_t}$$

그리고

$$x_t=\sqrt{\bar\alpha_t}x_0+\sqrt{1-\bar\alpha_t}\epsilon$$

이므로

$$x_t-\sqrt{\bar\alpha_t}x_0=\sqrt{1-\bar\alpha_t}\epsilon$$

입니다. 따라서

$$\nabla_{x_t}\log q(x_t\mid x_0) = -\frac{\epsilon}{\sqrt{1-\bar\alpha_t}}$$

가 됩니다.

이 식은 굉장히 중요합니다.

왜냐하면 여기서 바로 보이듯이,

조건 score는 노이즈의 반대 방향이기 때문입니다.

---

10. $\epsilon$-prediction을 score estimator로 바꿔보자

모델이 $\epsilon_\theta(x_t,t)$를 예측한다면,

그걸 이용해 score-like quantity를 정의할 수 있습니다.

$$s_\theta(x_t,t) := -\frac{\epsilon_\theta(x_t,t)}{\sqrt{1-\bar\alpha_t}}$$

그러면 진짜 조건 score는

$$\nabla_{x_t}\log q(x_t\mid x_0) = -\frac{\epsilon}{\sqrt{1-\bar\alpha_t}}$$

이므로, simple loss는

$$|\epsilon-\epsilon_\theta|^2 = (1-\bar\alpha_t) \left| s_\theta(x_t,t)-\nabla_{x_t}\log q(x_t\mid x_0) \right|^2$$

와 같습니다.

즉 $\epsilon$-prediction MSE는 사실상

조건 score를 맞히는 제곱오차

와 같은 말입니다.

단지 time-dependent scaling만 다를 뿐입니다.

이것이 “$\epsilon$-prediction과 score prediction이 본질적으로 연결된다”는 뜻입니다.

---

11. 그런데 우리가 정말 원하는 score는 어떤 score인가

여기서 미묘하지만 중요한 점이 있습니다.

방금 맞춘 것은

$$\nabla_{x_t}\log q(x_t\mid x_0)$$

즉 조건 score입니다.

이건 “원본 $x_0$를 알고 있다면”의 score입니다.

하지만 reverse dynamics에서 진짜 중요한 것은 noisy marginal distribution

$$q_t(x_t)$$

의 score, 즉

$$\nabla_{x_t}\log q_t(x_t)$$

입니다.

그러면 질문이 생깁니다.

조건 score를 학습하면, 왜 noisy marginal score와 연결되는가?

여기서 바로 denoising score matching의 핵심 정리가 등장합니다.

---

12. 조건 score의 평균이 marginal score가 된다

noisy marginal은

$$q_t(x_t)=\int q(x_t\mid x_0)q_{\text{data}}(x_0)dx_0$$

입니다.

즉 clean data에 노이즈를 섞어서 얻는 전체 noisy distribution입니다.

이제 이 로그를 $x_t$로 미분해보면

$$\nabla_{x_t}\log q_t(x_t) = \frac{\nabla_{x_t} q_t(x_t)}{q_t(x_t)}$$

이고,

$$q_t(x_t)=\int q(x_t\mid x_0)q_{\text{data}}(x_0)dx_0$$

이므로

$$\nabla_{x_t} q_t(x_t) = \int \nabla_{x_t} q(x_t\mid x_0)q_{\text{data}}(x_0)dx_0$$

입니다.

여기서

$$\nabla_{x_t} q(x_t\mid x_0) = q(x_t\mid x_0)\nabla_{x_t}\log q(x_t\mid x_0)$$

를 넣으면

$$\nabla_{x_t}\log q_t(x_t) = \frac{\int q(x_t\mid x_0)q_{\text{data}}(x_0)\nabla_{x_t}\log q(x_t\mid x_0)dx_0}{q_t(x_t)}$$

입니다.

이제 분자와 분모를 보면 Bayes를 써서

$$q(x_0\mid x_t) = \frac{q(x_t\mid x_0)q_{\text{data}}(x_0)}{q_t(x_t)}$$

이므로

$$\nabla_{x_t}\log q_t(x_t) = \int q(x_0\mid x_t)\nabla_{x_t}\log q(x_t\mid x_0)dx_0$$

즉

$$\boxed{\nabla_{x_t}\log q_t(x_t) = \mathbb{E}\big[\nabla_{x_t}\log q(x_t\mid x_0)\mid x_t\big]}$$

가 됩니다.

이 식의 의미는 매우 중요합니다.

noisy sample $x_t$ 하나를 놓고 보면,

그걸 만들 수 있었던 여러 clean source $x_0$들이 있고,

조건 score를 그 posterior에 대해 평균내면

noisy marginal 전체의 score가 된다.

즉 조건 score는 per-sample target이고,

그걸 제곱오차로 학습한 네트워크의 최적해는

실제로는 noisy marginal score를 배우게 됩니다.

이게 denoising score matching의 핵심입니다.

---

13. 그래서 score matching이라고 부르는 이유

이제 전체가 연결됩니다.

우리가 실제로 하는 학습은 보통

$$\mathbb{E}|\epsilon-\epsilon_\theta(x_t,t)|^2$$

같은 노이즈 예측 MSE처럼 보입니다.

하지만 이를 score 변수로 다시 쓰면

$$\mathbb{E} \left[ (1-\bar\alpha_t) \left| s_\theta(x_t,t)-\nabla_{x_t}\log q(x_t\mid x_0) \right|^2 \right]$$

꼴입니다.

즉 형태상으로는

노이즈가 섞인 데이터의 score를 맞추는 loss

입니다.

그리고 그 최적해는 noisy marginal score

$$\nabla_{x_t}\log q_t(x_t)$$

가 되므로, 결국 모델은 각 noise level $t$에서

“현재 noisy point를 어디로 움직이면 더 데이터다운 쪽인가”

를 배우게 됩니다.

그래서 diffusion을 score-based model로도 이해할 수 있습니다.

---

14. reverse sampling과 score의 연결을 다시 보면

sampling할 때 $\epsilon_\theta$를 예측해서 reverse mean을 만들었죠.

$$\mu_\theta(x_t,t) = \frac{1}{\sqrt{\alpha_t}} \left( x_t-\frac{\beta_t}{\sqrt{1-\bar\alpha_t}}\epsilon_\theta(x_t,t) \right)$$

여기서

$$s_\theta(x_t,t)=-\frac{\epsilon_\theta(x_t,t)}{\sqrt{1-\bar\alpha_t}}$$

를 넣으면

$$\mu_\theta(x_t,t) = \frac{1}{\sqrt{\alpha_t}} \left( x_t+\beta_t s_\theta(x_t,t) \right)$$

처럼 볼 수 있습니다.

즉 reverse mean은 대략

현재 상태 ($x_t$)에서 score가 가리키는 방향으로 조금 이동한 위치

입니다.

이게 바로 “디퓨전 샘플링은 noisy point를 score가 가리키는 고밀도 방향으로 조금씩 되돌리는 과정”이라는 해석과 연결됩니다.

---

15. 지금 단계에서 꼭 남겨야 할 핵심 문장

여기까지를 압축하면 다음 여섯 문장이 핵심입니다.

첫째, 학습의 본질은 이상적인 reverse posterior

$$q(x_{t-1}\mid x_t,x_0)$$

를 모델 reverse

$$p_\theta(x_{t-1}\mid x_t)$$

가 잘 모사하도록 만드는 것입니다.

둘째, 두 분포를 가우시안으로 두고 분산을 맞추면, KL은 평균들 사이의 제곱오차로 단순화됩니다.

셋째, posterior mean을 $\epsilon$로 표현할 수 있으므로, mean matching은 결국 $\epsilon$-prediction MSE가 됩니다.

넷째, 그래서 diffusion 학습은 놀랍게도

“현재 noisy sample에 섞인 가우시안 노이즈를 맞혀라”라는 supervised regression 문제로 쓸 수 있습니다.

다섯째, 조건 score는

$$\nabla_{x_t}\log q(x_t\mid x_0) = -\frac{\epsilon}{\sqrt{1-\bar\alpha_t}}$$

이므로 $\epsilon$-prediction은 곧 조건 score prediction과 같습니다.

여섯째, denoising score matching 정리에 의해, 조건 score를 이런 방식으로 학습하면 noisy marginal score

$$\nabla_{x_t}\log q_t(x_t)$$

를 배우게 되고, 이것이 reverse denoising 방향이 됩니다.

---

16. 한 줄 요약

가장 압축해서 말하면,

diffusion loss는 원래 “이상적인 reverse step과 모델 reverse step의 차이”를 줄이는 KL에서 출발하지만, 가우시안 구조와 $\epsilon$-parameterization 덕분에 결국 “섞인 노이즈를 맞히는 MSE”로 바뀌며, 이 MSE는 본질적으로 noisy distribution의 score를 학습하는 denoising score matching과 같다고 이해하시면 됩니다.

이제 여기까지 오면,

forward → ideal reverse → actual reverse → $\epsilon$-prediction → loss → score matching

이 한 줄로 다 연결된 상태입니다.