역방향 SDE: 확산 모델을 거꾸로 돌리기

이 장 첫머리에서 본 확산 모델 DDPM은 한 단계마다 샘플을 원점 쪽으로 조금 당기고 잡음을 더해, 데이터를 표준정규분포로 바꿔 놓았다. 이미지를 만들 때는 이 길을 거꾸로 가야 한다. 거꾸로 가는 길에도 잡음이 필요할까, 필요하다면 얼마나 넣어야 할까?

앞으로 가는 과정: 원점으로 당기는 랑주뱅

DDPM의 한 단계 x ← √(1 − ε)x + √ε z는 U = |x|²/2, kT = 1인 랑주뱅 방정식의 한 걸음이었다. 연속 시간으로 쓰면 송과 동료들(2021)이 분산 보존형(variance preserving) SDE라 부른 dx = −½ε(t)x dt + √ε(t) dB가 되고, 이 식의 평형은 볼츠만 분포 e^(−|x|²/2), 곧 표준정규분포다. 끌려가는 흐름 없이 퍼지기만 하는 분산 폭발형과 달리, 원점으로 당기는 흐름과 퍼짐이 비기는 곳이 있어서 분산이 1 근처에 묶인다. 이렇게 위치에너지가 x²에 비례하는 골짜기에서 움직이는 랑주뱅 방정식의 길을 오른스타인–울렌벡 과정이라 한다. 확산 모델이 생성을 시작하는 표준정규분포는 앞으로 가는 과정(전방 과정, forward process)의 평형 분포다.

잡음 없이 되돌리기: 확률 흐름

포커–플랑크 방정식에서 분포의 변화를 정하는 것은 자리마다의 흐름 속도 v였고, 입자를 잡음 없이 v로만 옮겨도 분포는 랑주뱅 방정식으로 옮길 때와 똑같이 흘렀다(끌려가는 속도 1, 확산 계수 0.5에서 표준정규분포를 시간 3까지 옮긴 입자 20만 개는 두 방법 모두 평균 3.00, 분산 3.99였다). 시간을 바꿔 전방 과정을 dx = −x dt + √2 dB로 두면 U = |x|²/2, ζ = 1, D = 1인 랑주뱅 방정식이고, 시간 t의 분포를 p_t라 하면 흐름 속도는 v = −x − ∇ln p_t다. 거꾸로 가는 시간 τ(t가 줄어드는 쪽으로 늘어나는 시간)로 보면 흐름의 방향이 뒤집히므로, 샘플을 dx/dτ = x + ∇ln p_t로 옮기면 분포는 p_t를 거꾸로 되짚어 데이터 분포로 돌아간다. 잡음 없이 결정론적으로 분포의 변화를 되짚는 이 식을 확률 흐름 ODE(ODE는 상미분 방정식)라 부른다. 필요한 재료는 시간마다의 스코어 ∇ln p_t뿐이다.

거꾸로 가는 길에 잡음 쌍 얹기

그런데 확산 모델에서 흔히 쓰는 것은 거꾸로 가는 길에도 잡음을 넣는 방법이다. 확률 흐름 ODE에 랑주뱅 쌍 ∇ln p_t dτ + √2 dB를 더해 보자. 이 쌍은 U = −ln p_t, kT = 1인 랑주뱅 방정식이어서 p_t를 평형 분포로 가지므로, 더해도 분포는 바뀌지 않는다. 거꾸로 가는 길에서도 마찰과 잡음의 짝이 나오는 것이다. 둘을 더한 역방향 SDE (전방 과정이 분포를 바꾼 길을 시간을 거꾸로 되짚는 확률 미분 방정식, reverse-time SDE)는 다음과 같다.

d\textcolor{#1b9e77}{\mathbf{x}} = \big[\textcolor{#1b9e77}{\mathbf{x}} + 2\nabla \ln \textcolor{#e377c2}{p_t}(\textcolor{#1b9e77}{\mathbf{x}})\big]\,d\tau + \sqrt{2}\,d\textcolor{#1f6066}{\mathbf{B}} = \underbrace{\big[\textcolor{#1b9e77}{\mathbf{x}} + \nabla \ln \textcolor{#e377c2}{p_t}\big]\,d\tau}_{\text{확률 흐름 ODE}} + \underbrace{\nabla \ln \textcolor{#e377c2}{p_t}\,d\tau + \sqrt{2}\,d\textcolor{#1f6066}{\mathbf{B}}}_{\textcolor{#e377c2}{p_t}\text{ 를 바꾸지 않는 랑주뱅}}

\begin{array}{ll} \textcolor{#e377c2}{p_t} & \text{전방 과정에서 시간 t의 분포 (데이터에 잡음이 섞인 분포)} \\ \nabla \ln \textcolor{#e377c2}{p_t} & \text{그 스코어 (신경망이 배우는 것)} \\ \textcolor{#a67a63}{\tau} & \text{거꾸로 가는 시간 (t가 줄어드는 쪽으로 늘어남)} \\ \textcolor{#1b9e77}{\mathbf{x}} + \nabla \ln \textcolor{#e377c2}{p_t} & \text{확률 흐름 ODE의 속도: 분포를 결정론적으로 되돌린다} \\ \textcolor{#1b9e77}{\mathbf{x}} & \text{샘플 (데이터 공간의 한 점)} \\ d\textcolor{#1f6066}{\mathbf{B}} & \text{거꾸로 가는 시간 dτ 동안 표준 브라운 운동의 변화} \end{array}

스코어 앞의 2는 둘로 나뉜다. 하나는 확률 흐름 ODE의 몫이고, 다른 하나는 U = −ln p_t, kT = 1인 랑주뱅 방정식의 끌려가는 흐름이다. 이 랑주뱅 부분은 잡음 √2 dB와 짝을 이뤄 p_t를 평형 분포로 가지므로 분포를 바꾸지 않는다. 역방향 SDE는 결정론적인 흐름으로 분포를 되돌리면서, 동시에 그 순간의 분포 안에서 샘플을 흔들어 섞는 것이다.

두 봉우리 데이터(−2에 0.2, +2에 0.8, 표준편차 0.5)에서 전방 과정 dx = −x dt + √2 dB가 샘플을 표준정규분포로 흩뜨리는 길(왼쪽)과, 스코어를 정확히 아는 역방향 SDE가 표준정규분포의 샘플을 두 봉우리로 되돌리는 길(오른쪽). 배경의 짙기는 시간마다의 분포 pₜ다.

ML에서: 담금질 랑주뱅과 예측–교정 샘플러

송과 에르몬(2019)의 담금질 랑주뱅이 잡음 수준마다 돌린 것도, 송과 동료들(2021)의 예측–교정 샘플러가 교정 단계에서 돌린 것도 바로 이 p_t를 바꾸지 않는 랑주뱅이다. 논문 요약의 「결정적으로, 역방향 SDE는 잡음 섞인 데이터 분포의 시간에 따른 기울기 장(자리마다 기울기가 하나씩 정해진 것, 곧 스코어)에만 의존한다」를, 이 장을 마친 독자는 「역방향 SDE는 확률 흐름에 요동-소산의 짝을 맞춘 랑주뱅을 얹은 것이고, 두 부분 모두 스코어만 있으면 쓸 수 있다」로 읽는다.

ML에서: 확산 모델의 CLD

확산 모델에도 같은 생각이 들어왔다. 도크혼·바흐다트·크라이스(2022)의 CLD(critically-damped Langevin diffusion)는 데이터 x마다 속도 v를 붙여, 전방 과정을 (x, v) 공간의 언더댐프드 랑주뱅 방정식으로 만들었다. 잡음은 속도에만 들어가고 데이터는 속도를 따라 매끄럽게 움직인다. 논문 요약은 이 보조 변수를 「해밀턴 역학에서처럼 데이터 변수와 짝지어진 〈속도〉」로 볼 수 있고, 모델은 「데이터가 주어졌을 때 속도의 조건부 분포의 스코어만 배우면 된다」고 적었다. 해밀턴 역학 + 마찰 + 잡음이라는 식이 샘플러에서 생성 모델의 전방 과정으로 자리를 옮긴 것이다.

문제 18. 확산 모델을 거꾸로 돌리기

데이터는 문제 15와 같은 두 봉우리(−2에 0.2, +2에 0.8, 표준편차 0.5)다. 전방 과정 dx = −x dt + √2 dB를 시간 5까지 돌린 분포 p_t의 스코어를 정확히 안다. 표준정규분포에서 1만 개를 뽑아 시간 5에서 0까지 거꾸로 간다. (가) 전방 과정은 어디에 이르는가? (나) 역방향 SDE dx = [x + 2∇ln p_t]dτ + √2 dB와 확률 흐름 ODE dx/dτ = x + ∇ln p_t의 결과를 비교하라. (다) 스코어의 계수와 잡음을 바꾸면 어떻게 되는가?

전방 과정은 U = x²/2인 랑주뱅 방정식이니까 평형이 표준정규분포예요. 시간 5면 처음 위치가 e^(−5) = 0.0067배만 남으니 표준정규분포에서 출발해도 되고요. DDPM이 샘플을 줄이면서 잡음을 더하는 것도 이 평형을 만들려는 거네요.

(나)는 둘 다 돌려 봤어요. 오른쪽 비율이 SDE 0.796, ODE 0.803, |x|의 평균은 둘 다 1.995, |x|의 표준편차는 0.500과 0.505예요. 둘 다 맞아요. 그런데 SDE는 왜 스코어가 두 배인지 모르겠어서, ODE에 잡음만 붙여 봤어요. 확률적인 버전이니까요. 그랬더니 오른쪽 비율 0.589, |x|의 평균 3.02, 표준편차 1.31… 봉우리가 뭉개지고 바깥으로 퍼졌어요.

잡음만 넣고 짝이 되는 마찰이 없잖아. 넣은 잡음만큼 계가 데워진 거야.

그러는 너는?

나는 반대로 SDE에서 잡음을 뺐어. 스코어 두 배만 두면 더 깔끔하게 봉우리로 가겠지 싶어서. 그랬더니… 모든 점이 오른쪽 봉우리 근처 한 점으로 모였어. 오른쪽 비율 1.000, |x|의 표준편차 0.003.

너는 너무 식혀서 한 점에 몰아넣었네. 왼쪽 봉우리는 아예 사라졌고.

둘 다 요동-소산의 짝을 깬 거예요. 스코어 두 배 가운데 하나는 확률 흐름 ODE의 몫이고, 나머지 하나와 잡음 √2 dB가 한 쌍이에요. 그 쌍은 p_t를 평형으로 갖는 랑주뱅이라 분포를 바꾸지 않아요. 쌍에서 잡음만 남기면 데워지고, 스코어만 남기면 식어요.

그러면 그 쌍의 비율만 지키면 몇 배로 넣어도 되겠네요. 스코어를 한 번 더 넣을 때 잡음도 √2 dB만큼 함께 넣으면요.

맞아요. 쌍을 얼마나 넣느냐에 따라 거꾸로 가는 길에서 샘플을 얼마나 섞을지를 조절할 수 있어요. 쌍을 하나도 넣지 않은 것이 확률 흐름 ODE이고, 하나 넣은 것이 역방향 SDE예요.

결국 확산 모델의 샘플러도 마찰과 잡음의 짝을 맞추는 문제였네요.