확산 모델(Diffusion model)

평면에 점 480개가 여덟 무더기로 고리 모양을 이루고 있습니다. 이 점들이 우리가 가진 자료라고 합시다. 사진 한 장을 수십만 차원 공간의 점 하나로 보면 사진 모음도 이런 점구름입니다. 이제 걸음마다 모든 점을 원점 쪽으로 아주 조금 당기고, 작은 정규분포⁠(normal distribution)⁠ 잡음을 더합니다.

$$x_k = \sqrt{1-\beta_k}\; x_{k-1} + \sqrt{\beta_k}\; z_k, \qquad z_k \sim N(0, I)$$

$\beta_k$는 걸음마다 미리 정해 둔 작은 수로, 여기서는 0.0005에서 0.12까지 커지는 100개입니다. 점 하나는 원점으로 끌리는 무작위 걸음을 하고, 다음 위치가 지금 위치에만 달려 있으니 마르코프 연쇄⁠(Markov chain)⁠입니다. 당기는 비율 $\sqrt{1-\beta_k}$는 분산⁠(variance)⁠이 1인 분포가 한 걸음 뒤에도 분산 1이 되도록 고른 것입니다. 독립⁠(independence)⁠인 두 항의 분산은 더해지므로 $(1-\beta_k)\cdot 1 + \beta_k = 1$입니다. 아래에서 시간 t는 걸음 수 k를 100으로 나눈 것(t = k/100)이고, $x_t$는 그 걸음의 위치입니다. 방향: · 시간 t = · 점수 화살표: · 흘려 보기

t = 1에 이르면 고리는 거의 알아볼 수 없고, 점들은 원점 둘레에 둥글게 퍼진 표준정규분포에 가까운 모양이 됩니다. 이유는 한 줄입니다. 출발점 $x_0$를 고정하면 걸음마다 더하는 잡음은 서로 독립인 정규분포이고, 독립인 정규분포를 더하면 다시 정규분포입니다. 그래서 여러 걸음을 한꺼번에 계산하면 다음과 같습니다.

$$x_t = \sqrt{\bar\alpha_t}\; x_0 + \sqrt{1-\bar\alpha_t}\; \varepsilon, \qquad \bar\alpha_t = \prod_{s \le t} (1-\beta_s)$$

$\sqrt{\bar\alpha_t}$는 처음 위치가 남은 비율로, 지금 입니다. 마지막 걸음에서는 0.043까지 줄어 처음 위치가 거의 지워지고, 잡음의 분산 $1-\bar\alpha_t$는 1에 가까워집니다. 처음 분포가 무엇이었든 끝은 거의 같은 표준정규분포라는 점이 중요합니다. 그래서 거꾸로 갈 때 출발점을 표준정규분포에서 뽑으면 됩니다. 잡음으로 정규분포를 쓰는 것은 계산이 편한 선택이지만 자의적이지는 않습니다. 작고 독립인 흔들림이 많이 쌓이면 원래 모양과 상관없이 정규분포에 가까워진다는 것이 중심극한정리⁠(central limit theorem)⁠이기 때문입니다.

거꾸로 가기. 앞으로 가는 과정은 누구나 할 수 있습니다. 확산 모델의 요점은 거꾸로 가는 법을 배우는 데 있습니다. 표준정규분포에서 새로 뽑은 점들을 걸음마다 조금씩 되돌리면, 끝에는 자료처럼 생긴 점들이 나와야 합니다. 방향을 '거꾸로'로 바꾸고 t를 1에서 0으로 줄여 보세요. β가 작으면 거꾸로 가는 한 걸음도 거의 정규분포 모양입니다. 분산은 β로 정해 두면 되니 평균⁠(mean)⁠만 알면 되고, 그 평균을 정하는 것이 점수 $\nabla_x \log p_t(x)$입니다. 여기서 $p_t$는 시간 t의 점구름의 확률밀도이고, 점수는 그 밀도가 가장 빨리 짙어지는 방향을 가리킵니다. 1982년 브라이언 앤더슨은 잡음을 더하는 확률⁠(probability)⁠ 과정을 시간 거꾸로 돌린 것도 같은 꼴의 확률 과정이며, 그 표류 항⁠(drift term)⁠에 바로 이 점수가 들어간다는 것을 보였습니다. 거꾸로 한 걸음은 다음과 같습니다.

$$x_{k-1} = \frac{1}{\sqrt{1-\beta_k}}\Bigl(x_k + \beta_k\, \nabla \log p_k(x_k)\Bigr) + \sqrt{\beta_k}\; z_k$$

그림의 화살표가 그 점수입니다. 잡음이 클 때(t가 1 가까이)는 화살표가 모두 원점 쪽을 가리켜 점들을 대충 모으고, 잡음이 줄수록 여덟 무더기 쪽으로 갈라져 점들을 제 무더기로 보냅니다. 점의 색은 t = 0에서 도착한 무더기의 색입니다. t가 클 때는 색이 뒤섞여 있다가 t가 줄면서 갈라지는 것을 보면, 어느 무더기로 갈지가 거꾸로 가는 도중에 조금씩 정해진다는 것을 알 수 있습니다.

무엇을 배우는가. 이 그림은 한 가지를 속였습니다. 여기 자료는 중심과 퍼짐을 아는 정규분포 여덟 개의 혼합이라서, 어느 시간의 점수든 식으로 정확히 계산할 수 있습니다. 사진 모음의 점수는 아무도 모릅니다. 그래서 신경망 $\varepsilon_\theta(x_t, t)$가 잡음 섞인 점 $x_t$를 보고 거기 더해진 잡음 ε을 맞히도록 학습시킵니다. 손실은 맞힌 잡음과 실제 잡음의 차이의 제곱 $\lVert \varepsilon_\theta(x_t, t) - \varepsilon \rVert^2$입니다. 제곱 오차를 가장 작게 하는 함수⁠(function)⁠는 조건부 기댓값⁠(expected value)⁠ $E[\varepsilon \mid x_t]$이고, 이것은 점수와 다음처럼 이어져 있습니다(트위디 공식⁠, Tweedie's formula⁠).

$$\nabla \log p_t(x_t) = -\frac{E[\varepsilon \mid x_t]}{\sqrt{1-\bar\alpha_t}}$$

잡음을 맞히는 최선의 함수를 알면 점수를 아는 것과 같습니다. 2020년 조너선 호, 아제이 자인, 피터 아빌의 DDPM이 이 단순한 손실로 좋은 그림을 만들어 보였습니다. 이 손실은 자료의 로그 가능도⁠(log-likelihood)⁠에 대한 하한(증거 하한⁠, evidence lower bound⁠)을 정규분포 사이의 쿨백–라이블러 발산⁠(Kullback–Leibler divergence)⁠ 항들로 풀어 쓴 뒤 항들의 가중치⁠(weight)⁠를 바꾼 것이기도 합니다.

열 방정식과의 관계. 당기지 않고 잡음만 더한다면($x_t = x_0 + \sqrt t\, z$) 점구름의 밀도는 처음 밀도를 점점 넓은 종 모양으로 번지게 한 것이고, 이것은 열 방정식 $\partial_t p = \tfrac12 \Delta p$의 해입니다(시간의 눈금을 바꾸면 $u_t = u_{xx}$ 꼴). 점 하나하나의 무작위 걸음이 모여 열의 확산이 되는 것입니다. 당기는 항이 있으면 방정식에 표류 항이 하나 붙지만(포커–플랑크 방정식⁠, Fokker–Planck equation⁠) 이야기는 같습니다. 열 방정식을 시간 거꾸로 푸는 것은 잘 알려진 불안정한 문제입니다. $u_t = u_{xx}$에서 진동수⁠(frequency)⁠ ω인 성분은 앞으로는 $e^{-\omega^2 t}$배로 식으니, 거꾸로는 $e^{\omega^2 t}$배로 불어납니다. 진동수가 높은 성분일수록 더 빨리 불어나므로 아주 작은 오차도 폭발합니다. 확산 모델은 밀도를 거꾸로 푸는 대신 점수라는 정보를 따로 배워서, 표본⁠(sample)⁠ 하나하나를 잡음과 함께 거꾸로 옮깁니다. 앞으로 가며 잃은 정보를 되살리는 데 필요한 것은 점수에 담겨 있고, 점수는 자료에서 배웁니다.

그림을 만들기까지. 실제 그림 모델은 x가 수십만 차원이라는 것 말고는 그림과 같은 일을 합니다. 순수한 잡음에서 출발해 신경망이 어림한 점수로 수십에서 천 걸음을 거꾸로 가면 새 그림이 나옵니다. 문장으로 원하는 그림을 지시하려면 신경망에 문장도 함께 넣고, 문장이 있을 때와 없을 때의 점수 차이를 부풀리는 방법(분류기 없는 안내⁠, classifier-free guidance⁠)이 흔히 쓰입니다. 2022년의 Stable Diffusion은 그림을 오토인코더⁠(autoencoder)⁠로 작은 잠재 공간⁠(latent space)⁠에 압축한 뒤 그 안에서 확산을 돌려 계산을 줄였습니다(잠재 확산⁠, latent diffusion⁠). 걸음마다 잡음을 더하지 않고 점수를 따라 결정적으로 움직이는 길(확률 흐름 상미분방정식⁠, probability flow ODE⁠)도 있습니다. 점 하나하나의 길은 다르지만 각 시간의 점구름 분포는 잡음을 더하는 과정과 같아서, 걸음 수를 크게 줄이는 여러 방법이 여기서 나왔습니다.

역사와 한계. 2015년 야샤 솔-딕스타인과 동료들이 비평형 통계 역학⁠(statistical mechanics)⁠에서 착안해 이 틀을 처음 내놓았고, 2019년 양 쑹과 스테파노 에르몬은 여러 잡음 수준의 점수를 배워 표본을 만드는 방법을 보였으며, 2020년 DDPM이 지금의 표준이 된 형태를 보였습니다. 2021년 쑹과 동료들은 둘을 확률 미분방정식⁠(differential equation)⁠ 하나로 묶었습니다. 2022년 DALL·E 2, Imagen, Stable Diffusion이 잇달아 나오면서 확산 모델은 그림 생성의 주류가 되었습니다. 한계도 분명합니다. 걸음이 많아 느립니다. 또 2023년 니컬러스 칼리니 등은 여러 확산 모델에서 학습에 쓴 그림을 천 장 넘게 거의 그대로 다시 뽑아낼 수 있음을 보였습니다. 학습 자료 전체에 견주면 아주 적은 몫이지만, 사람 얼굴 사진이나 상표도 들어 있었습니다. 모델이 배운 것은 자료 분포의 어림이라서, 자료에 없던 것을 만들어 내는 능력과 자료를 외워 되풀이할 위험이 한 동전의 양면입니다.

이어지는 곳. 앞으로 가는 과정의 뼈대는 무작위 걸음과 마르코프 연쇄이고, 끝이 왜 표준정규분포인지는 정규분포의 덧셈 성질에서, 잡음으로 왜 정규분포를 쓰는지는 중심극한정리에서 볼 수 있습니다. 밀도가 번지는 모양은 열 방정식이 그리고, 진동수마다 지수적으로 식는 그 해의 모양이 거꾸로 돌리기가 왜 불안정한지 알려 줍니다. 학습 목표가 가능도⁠(likelihood)⁠의 하한⁠(lower bound)⁠에서 나오는 과정은 쿨백–라이블러 발산과 최대가능도법⁠(maximum likelihood)⁠에, 잡음 수준마다 점수를 맞히는 신경망의 학습은 확률적 경사 하강법⁠(stochastic gradient descent)⁠에 있습니다. 그림을 작은 잠재 공간으로 줄여 계산을 아끼는 부분은 오토인코더가 맡고, 무작위로 뽑은 표본으로 분포를 다루는 생각은 몬테카를로 방법⁠(Monte Carlo method)⁠과 무작위성에서 이어집니다. 같은 생성 모델이지만 한 토큰씩 차례로 뽑는 방식은 언어 모델⁠(language model)⁠입니다.