역방향 SDE: 디퓨전 모델을 거꾸로 돌리기

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

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

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

노이즈 없이 되돌리기: 확률 플로우 ODE

포커–플랑크 방정식에서 분포의 변화를 정하는 것은 자리마다의 흐름 속도 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를 거꾸로 되짚어 데이터 분포로 돌아간다. 노이즈 없이 결정론적으로 분포의 변화를 되짚는 식 dx/dτ = x + ∇ln p_t를 확률 플로우 ODE(ODE는 상미분 방정식)라 부른다. 필요한 재료는 시간마다의 스코어 ∇ln p_t뿐이다.

거꾸로 가는 길에 노이즈 쌍 얹기

그런데 디퓨전 모델에서 흔히 쓰는 것은 거꾸로 되짚는 과정에도 노이즈를 넣는 방법이다. 확률 플로우 ODE에 랑주뱅 쌍 ∇ln p_t dτ + √2 dB를 더해 보자. 이 쌍은 U = −ln p_t, kT = 1인 랑주뱅 방정식이어서 p_t를 평형 분포로 가지므로, 더해도 분포는 바뀌지 않는다. 거꾸로 되짚는 과정에서도 끌려가는 항 ∇ln p_t dτ와 노이즈 √2 dB가 요동-소산 정리대로 맞춘 쌍으로 나오는 것이다. 둘을 더한 역방향 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는 확률 플로우 ODE에, 끌려가는 항과 노이즈의 크기를 요동-소산 정리대로 맞춘 랑주뱅을 얹은 것이고, 두 부분 모두 스코어만 있으면 쓸 수 있다.

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예요.

결국 디퓨전 모델의 샘플러도 끌려가는 항과 노이즈의 크기를 서로 맞추는 문제였네요.