비평형 열역학

물에 잉크 방울을 떨어뜨리면 시간이 지남에 따라 잉크가 퍼져나가며 결국 완전히 흐려지는 현상(엔트로피 증가)을 의미한다. 디퓨전 모델은 이 자연계의 확산 과정을 수학적인 마르코프 연쇄로 모델링하여 원본데이터에 점진적으로 노이즈를 주입하는 forward process의 근간으로 삼았다.

새로운 연결 고리와 가중치가 부여된 변분 하한

통계학에서 VAE를 풀기 위해 사용한 증거 하한(ELBO, 변분 하한) 수식을 전개해보니, 그 결과가 물리학의 랑주뱅 동역학 기반의 스코어 매칭 목적함수와 동일한 형태가 됨을 밝혀낸 것이다. 여기서 ‘가중치가 부여된’이란 말은, 수학적으로 엄밀하게 유도된 원래의 ELBO 수식에서 특정 시간대(t)의 가중치 항들을 의도적으로 날리고 단순화한 목적 함수를 사용했더니 오히려 이미지 생성 품질이 훨씬 좋아졌다는 발견을 의미한다.

점진적 손실 압축 해제

순수한 노이즈( $\mathbf{x}_T$ )에서 시작해 원본 이미지( $\mathbf{x}_0$ )를 만들어가는 1000번의 디노이징 과정을 정보 이론의 관점에서 해석한다. 픽셀을 한 번에 하나씩 칠하는 기존의 자기회귀 모델과 달리, 디퓨전은 거대한 윤곽(저주파)부터 시작해 미세한 디테일(고주파)까지 점진적으로 복원해나간다. 이는 아주 강하게 압축된 파일(노이즈)의 압축을 서서히 풀면서 이미지를 선명하게 만들어가는 과정과 동치이다.

변분 추론을 사용하여 학습된 마르코프 연쇄

디퓨전 모델이 1000개의 계층을 가지는 마르코프 연쇄이며, 이를 학습할 때 VAE에서 사용했던 증거 하한 수식을 가져와 최적화한다는 것을 밝힌다.

신호가 파괴될 때까지 노이즈를 추가하는 확산 과정

forward process의 정의이다. 데이터에 노이즈를 계속 더해서 원점( $\mathcal{N}(\mathbf{0}, \mathbf{I})$ )으로 유도하여 원본 이미지의 정보(신호)를 완전히 없애버리는 과정을 의미한다.

연쇄의 전이는 확산 과정을 되돌리도록 학습된다

reverse process(디노이징)의 정의이다. 망가진 노이즈 상태에서 출발하여 원래의 데이터 분포로 되돌아오는 복원력(스코어)을 학습한다는 뜻이다.

단순한 신경망 매개변수화

매 스텝마다 추가되는 노이즈의 양이 매우 작다면, 그 방향(디노이징)으로 돌아가는 확률 분포 역시 수학적으로 가우시안 분포가 됨이 통계적으로 증명되어 있다. 역방향이 가우시안 분포라는 것은, 딥러닝 모델(u-net)이 복잡한 확률 분포 전체를 어렵게 예측할필요 없이 오직 가우시안의 평균(결국 노이즈 혹은 스코어) 단 하나만 예측하도록 단순하게 설계할 수 있다는 이점을 설명한다.

디퓨전 모델 자체는 2015년에 처음 제안되었지만, 당시에는 생성된 이미지의 품질이 좋지 않았다. 본 논문의 저자들은 u-net이 원본 이미지( $\mathbf{x}_0$ )나 평균( $\tilde{\mu}_t$ )을 예측하게 하지 않고, 추가된 가우시안 노이즈( $\epsilon$ )를 예측하도록 수식을 우회하였다(특정 매개변수화). 이 우회하는 과정이 랑주뱅 동역학의 복원력(스코어 함수)을 계산하는 식과 동치인 것이다.

우도(likelihood)란 이 모델이 실제 데이터의 분포를 얼마나 근사하는가를 측정하는 지표이다. VAE나 flow모델들은 이 점수가 아주 높지만, DDPM은 화질이 좋음에도 불구하고 이 점수는 상대적으로 낮았다.

무손실 코드 길이와 눈에 보이지 않는 디테일(이미지 세부 사항)

우도가 낮게 나온 이유를 분석해 보니, 모델의 정보 처리 용량의 절반 이상을 사람 눈에는 보이지도 않는 픽셀 단위의 미세한 노이즈(고주파 패턴)를 복원하는 데 낭비하고 있었기 때문이다. 때문에, 모델은 사람이 이미지를 인식하는 방식처럼 거대한 측징(저주파)부터 먼저 복원하고, 작은 디테일은 나중에 채워 넣는 디코더(압축 해제기)의 역할(점진적 손실 압축 해제)을 유도한다.

실제 데이터가 샘플링된 분포 $q(\mathbf{x}_0)$ 를 알 수 없으므로, 신경망( $p_\theta(\mathbf{x}_0) := \int p_\theta(\mathbf{x}_{0:T}) d\mathbf{x}_{1:T}$ )을 통해 근사한다. 하지만 이 분포는 복잡하여 한 번에 정의할 수 없으므로, $\mathbf{x}_1$ 부터 $\mathbf{x}_T$ 까지 총 $T$ 개(논문에서는 1000개)의 잠재 변수를 도입한다. 전체 결합 확률 $p_\theta(\mathbf{x}_{0:T})$ 에서 잠재 변수들( $\mathbf{x}_1 \sim \mathbf{x}_T$ )을 적분해서 없애면(경우의 수를 다 더하면), 진짜 데이터의 확률 $p_\theta(\mathbf{x}_0)$ 만 남는다는 뜻이다.

VAE 같은 기존 모델들은 원본 이미지( $\mathbf{x}_0$ )를 저차원의 잠재 벡터로 압축한다. 반면 디퓨전 모델은 $256 \times 256$ 해상도라면, $\mathbf{x}_1$ 도 $256 \times 256$ 이고, 1000번째 노이즈 $\mathbf{x}_{1000}$ 도 완벽하게 똑같은 $256 \times 256$ 해상도이다. 차원이 똑같다는 것은, 1000개의 단계가 모드 이미지 공간 그 자체라는 뜻이며, 이를 통해 이미지의 세밀한 구조가 파괴되지 않고 유지된다.

(디노이징, reverse) 출발점 ( $p(\mathbf{x}_T)$ ): 먼저 $p(\mathbf{x}_T) = \mathcal{N}(\mathbf{x}_T; \mathbf{0}, \mathbf{I})$ 라는 사전 분포(완전 노이즈)에서 출발한다.

마르코프 연쇄 ( $\prod$ ): 결합 분포 $p_\theta(\mathbf{x}_{0:T})$ 를 단순히 이전 단계에서 다음 단계를 조건부로 예측하는 확률( $p_\theta(\mathbf{x}_{t-1} \mid \mathbf{x}_t)$ )들의 계속된 곱셈( $\prod$ )으로 표현했다. 즉, $\mathbf{x}_{999}$ 를 만들 때는 오직 $\mathbf{x}_{1000}$ 만 보고 만들지, $\mathbf{x}_{1000}$ 이전의 상태는 고려하지 않겠다는 의미이다.

가우시안 전이 ( $\mathcal{N}$ ): 또한 각 스텝에서 노이즈( $\mathbf{x}_t$ )를 살짝 덜어낸 이전 상태( $\mathbf{x}_{t-1}$ )를 예측할 때, 그 확률 분포가 정규분포를 따른다고 강제한다.

신경망의 역할 ( $\boldsymbol{\mu}_\theta, \boldsymbol{\Sigma}_\theta$ ): $\theta$ 는 딥러닝 모델(U-Net)의 가중치를 의미한다. 즉, u-net이 해야할 일은 이미지를 직접 그리는 것이 아니라, 특정 시간 $t$ 에 노이즈 낀 이미지 $\mathbf{x}_t$ 가 들어왔을 때, 바로 직전 단계( $\mathbf{x}_{t-1}$ )의 가우시안 분포는 평균이 $\boldsymbol{\mu}$ 이고 분산이 $\boldsymbol{\Sigma}$ 일 것을 예측하는 회귀 문제를 푸는 것이다.

그 이전 상태를 고려하지 않는다는 것은 무슨의미인가? 이전의 정보들을 가지고 있으면 좋지않나?

디퓨전 모델은 입자가 무작위로 퍼져나가는 물리학의 브라운 운동에서 영감을 받았다.

당구대 위에서 굴러가는 당구공을 상상해보자. 당구공이 1초 뒤에 어디로 갈지 예측하기 위해, 우리는 공의 현재 위치와 현재 속도만 알면 된다. 이 공이 5초 전에 쿠션을 몇 번 맞고 굴러왔는지(과거의 정보)는 전혀 알 필요가 없다. 왜냐하면 그 과거의 모든 사건들이 누적된 결과물이 바로 현재 상태이기 때문이다.

디퓨전또한 $t$ 시점의 이미지 $\mathbf{x}_t$ 안에는 이미 1000번째부터 $t+1$ 번째까지 가해진 모든 디노이징의 결과가 담겨있다. 즉, $\mathbf{x}_t$ 자체가 과거 정보의 요약이므로, 굳이 방대한 과거 사진 스냅샷을 들고 다닐 필요가 없는 것이다. $\mathbf{x}_t$ 하나가 256x256 해상도의 컬러 이미지(약 20만 개의 숫자)일때, 만약 990번째 이미지를 예측하기 위해 1000번째부터 991번째까지 총 10장의 이미지를 모두 u-net에 입력으로 넣는다면, 연산량이 어마어마할 것이다.

또한 수식 (1)에서 마르코프 연쇄 덕분에 거대한 결합 확률을 단순한 조건부 확률들의 곱셈으로 쪼갤 수 있었다.

만약 마르코프 가정을 빼고 과거를 모두 기억하게 한다면 위 결합 확률은 다음과 같다.

뒤로 갈 수록 조건부가 무한히 길어질 것이며, ELBO를 분수 형태로 약분할 수도 없고, KL발산 형태로 수식을 정리하는 것 자체가 불가능하다.

마르코프 가정을 깨고(비(Non)-마르코프 과정), 과거의 정보(원본 이미지로 향하는 방향 등)를 우회해서 활용한 것이 DDIM(Denoising Diffusion Implicit Models)이다. 현재 Stable Diffusion 등의 상용화 모델들은 모두 이 DDIM의 비마르코프 방식을 채택하였다.

(노이징, forward) VAE는 원본 이미지를 잠재 벡터로 압축하는 인코더를 학습시켜야 하는 반면, 디퓨전은 forward에서 학습이라는 과정이 없다. 전체 노이즈 주입 과정은, 바로 직전 단계 $\mathbf{x}_{t-1}$ 에 노이즈를 한 번 더 추가해서 현재 단계 $\mathbf{x}_t$ 를 만드는 단순한 작업의 $T$ 번(=1,000) 반복( $\prod$ )일 뿐이다.

논문에서는 총 스텝 수를 $T=1000$ 으로 고정하고, 분산 스케줄은 $\beta_1 = 0.0001$ (매우 작은 노이즈)에서 시작하여, $\beta_{1000} = 0.02$ (상대적으로 큰 노이즈)까지 선형적으로 서서히 증가하도록 스케줄을 짜두었다.

수식 (2)의 의미는 $t=1$ 일 때( $q(\mathbf{x}_1 \mid \mathbf{x}_0) = \mathcal{N}(\mathbf{x}_1; \sqrt{1 - \beta_1} \mathbf{x}_0, \beta_1 \mathbf{I})$ ), $\mathbf{x}_1$ 가 평균이 $\sqrt{1 - \beta_1} \mathbf{x}_0$ 이고 (공)분산이 $\beta_1 \mathbf{I}$ 인 가우시안을 따른다는 뜻이다. 실제로 $\mathbf{x}_1$ 를 샘플링하려면 재매개변수화 트릭을 통해 다음 연산을 수행해야 한다.

(공분산이 대각행렬이라는 것은 변수들 간에 상관관계가 없다는 뜻이다. 즉, 이미지 왼쪽 위에 추가된 노이즈와, 오른쪽 아래에 추가된 노이즈는 서로 아무런 연관성없이 독립적으로 무작위하게 뽑힌다.)

원본 이미지와 같은 고차원 벡터인 노이즈 벡터 $\boldsymbol{\epsilon}$ 의 각 원소는 $\mathcal{N}(0, 1)$ 에서 무작위로 뽑힌 난수이다.

$\mathbf{x}_{t-1}$ 가 가우시안 분포를 따르는 이유는 가우시안 확률 변수( $\boldsymbol{\epsilon} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$ )에 선형 결합(또는 아핀 변환)의 결과는 가우시안이기 때문이다. 위 재매개변수화 트릭(아핀 변환, $Y = aX + c$ )으로 만들어진 $\mathbf{x}_{t-1}$ 의 평균과 기댓값은 다음과 같이 구할 수 있다.

$\boldsymbol{\epsilon}$ 의 평균이 $\mathbf{0}$ 이므로, $E[\mathbf{x}_t \mid \mathbf{x}_{t-1}] = \sqrt{1 - \beta_t} \mathbf{x}_{t-1}$ 이다.

$Var(\boldsymbol{\epsilon}) = \mathbf{I}$ 이므로 $Var(\mathbf{x}_t \mid \mathbf{x}_{t-1}) = \beta_t \mathbf{I}$ 이다.

Q. 분산 스케줄을 파라미터로 두면 어떨까?

“The forward process variances $\beta_t$ can be learned by reparameterization…“에서 “순방향 과정의 분산 $\beta_t$ 는 재매개변수화를 통해 학습되거나 하이퍼파라미터로 상수로 유지될 수 있다.”고 말한다. 그러나 사전에 선형 스케줄로 고정했음에도 불구하고 당시 좋은 품질을 가질 수 있었다.

추후 연구에서 코사인 스케줄을 제안해 샘플링 속도를 1,000번에서 50번으로 줄이거나(Improved DDPM, 2021), 노이즈 스케줄 $\beta_t$ 자체를 신경망으로 두고 학습하게 만들었다(Variational Diffusion models, 2021). 그 결과 DDPM의 낮았던 로그 우도 점수를 최고 수준으로 끌어올리게 된다.

목표는 $p_\theta(\mathbf{x}_0)$ 을 최대화하는 $\theta$ 를 찾는 것이다. 여기에 로그를 씌워도 최적해는 바뀌지 않으며(로그가 단조증가함수이므로), 음수를 취해 minimize문제로 바꾼다.

그러나 이 값을 구하려면 1000개의 모든 노이즈 단계( $\mathbf{x}_1 \dots \mathbf{x}_T$ )의 경우의 수를 전부 적분해야 한다(intractable).

우리가 알고 싶은 것은 $p_\theta(\mathbf{x}_0)$ 이지만, 디퓨전 모델은 $\mathbf{x}_1$ 부터 $\mathbf{x}_T$ 까지의 잠재 변수를 포함하는 결합 확률 분포 $p_\theta(\mathbf{x}_{0:T})$ 로 정의되어 있다. marginalization에 의해 $p_\theta(\mathbf{x}_0)$ 는 다음과 같다.

이때 순방향 과정을 나타내는 확률 밀도 함수 $q(\mathbf{x}_{1:T} \mid \mathbf{x}_0)$ 를 분모와 분자에 곱하는 트릭을 쓴다(1을 곱하는 것과 같으므로 등식은 유지됨.).

어떤 함수 $f(X)$ 에 대해 확률 밀도 함수 $q(X)$ 를 곱하여 적분한 것은, $q$ 분포에 대한 $f(X)$ 의 기댓값의 정의와 같다.

양변에 $-\log$ 를 취하면

$-\log(x)$ 는 볼록함수이므로 젠슨 부등식( $f(\mathbb{E}[x]) \le \mathbb{E}[f(x)]$ )에 의해 다음이 성립한다.

좌변의 상한(upper bound)가 우변이므로, 우변을 최소화하는 $\theta$ 를 찾는 것은 좌변의 상한을 최소화하는 $\theta$ 를 찾는것과 동일하다.

위 upper bound의 분수에서 분자(역과정)는 $p_\theta(\mathbf{x}_{0:T}) = p(\mathbf{x}_T) \prod_{t=1}^T p_\theta(\mathbf{x}_{t-1} \mid \mathbf{x}_t)$ , 분모(순방향) $q(\mathbf{x}_{1:T} \mid \mathbf{x}_0) = \prod_{t=1}^T q(\mathbf{x}_t \mid \mathbf{x}_{t-1})$ 이므로

로그의 기본 성질 $\log(AB) = \log A + \log B$ 와 $\log(A/B) = \log A - \log B$ , 그리고 $\log(\prod A_i) = \sum \log A_i$ 를 이용하여 식을 전개한다.

이 전체 항을 기댓값 $\mathbb{E}_q$ 로 묶으면 식 (3)이 도출된다.

스케줄은 학습될 수 있으며, 노이즈 주입량( $\beta_t$ , 순방향)이 충분히 작다면, 역방향 과정의 확률 분포 또한 가우시안이다.

Q. 노이즈 주입량이 작은 것이 역방향또한 가우시안임을 어떻게 보장하는가?

역과정 확률인 $q(\mathbf{x}_{t-1} \mid \mathbf{x}_t)$ 를 베이즈 정리로 전개하면 다음과 같다.

$q(\mathbf{x}_t \mid \mathbf{x}_{t-1})$ 는 순방향 과정에서 $\mathcal{N}(\sqrt{1-\beta_t}\mathbf{x}_{t-1}, \beta_t \mathbf{I})$ 로 정의한 가우시안 분포이며, $q(\mathbf{x}_{t-1})$ 및 $q(\mathbf{x}_t)$ 는 전체 데이터 매니폴드가 시간 $t$ 에 따라 어떻게 분포하는지를 나타내는 확률 밀도 함수이다. 데이터셋 전체의 확률 분포이므로, 매우 복잡한 다봉형의 비선형 함수이다.

만약 $\beta_t$ 가 크다면(큰 스텝) $\mathbf{x}_t$ 와 $\mathbf{x}_{t-1}$ 사이의 거리가 멀어진다. 가우시안 함수인 $q(\mathbf{x}_t \mid \mathbf{x}_{t-1})$ 에 매우 복잡하고 굴곡진 $q(\mathbf{x}_{t-1})$ 를 곱하게 되므로, 그 결과물인 좌변의 $q(\mathbf{x}_{t-1} \mid \mathbf{x}_t)$ 는 찌그러진 비대칭 형태가 되어 가우시안 분포를 유지할 수 없다.

만약 $\beta_t$ 가 극도로 작다면( $\beta_t \to 0$ ) $\mathbf{x}_{t-1}$ 은 $\mathbf{x}_t$ 의 극소 주변에만 존재하게 된다. 미적분학의 원리에 따라, 복잡한 비선형함수(데이터의 원래 확률 분포 $q$ )라도 극소 영역만 보면 선형 함수로 근사할 수 있다(1차 테일러 전개).

극소 영역에서 $\log q(\mathbf{x}_{t-1})$ 는 선형 근사에 의해 상수 벡터 $\mathbf{c}$ 에 대하여 $\mathbf{c}^T \mathbf{x}_{t-1}$ 로 표현될 수 있으며, 따라서 $q(\mathbf{x}_{t-1}) \approx e^{\mathbf{c}^T \mathbf{x}_{t-1}}$ 형태의 단순한 지수 함수로 근사된다.

가우시안 함수에 지수함수( $e^{\mathbf{c}^T \mathbf{x}}$ )를 곱하면, 그 결과는 분산은 그대로고 평균만 c방향으로 평행이동한 가우시안 함수가 된다. 따라서 $\beta_t$ 가 충분히 작을 경우

라는 연산이 성립한다.

(아직 이해가 부족. 테일러 전개, 선형 근사)

$\beta_t$ 를 $\alpha_t := 1 - \beta_t$ 로 표기함으로써 매 스텝 추가되는 가우시안 노이즈의 분산(노이즈 주입률, $\beta_t$ )을 전체 에너지 1 중에서 노이즈 $\beta_t$ 를 뺀 나머지, 즉 $\mathbf{x}_{t-1}$ 의 정보가 다음 스텝 $\mathbf{x}_t$ 로 얼마나 보존되었는지(비율, 신호 보존율)로 해석할 수 있다. 식 (4)에서 $\sqrt{\bar{\alpha}_t}$ 는 1단계부터 $t$ 단계까지 거치며 살아남은 원본 신호( $\mathbf{x}_0$ )의 총 비율이며, $(1 - \bar{\alpha}_t)$ 는 1단계부터 $t$ 단계까지의 누적된 노이즈의 총분산이다.

식 (5)의 유도(부록 A)는 아래에 있다.

$L_T$ : $D_{\text{KL}}(q(\mathbf{x}_T \mid \mathbf{x}_0) \parallel p(\mathbf{x}_T))$ , 순방향 과정이 끝나고 도달한 노이즈( $q$ )가, 사전에 가정한(prior) 가우시안 분포( $p$ )와 얼마나 똑같은지를 측정한다. 이는 $\theta$ 와 무관한 상수이다.

$L_{t-1}$ : $D_{\text{KL}}(q(\mathbf{x}_{t-1} \mid \mathbf{x}_t, \mathbf{x}_0) \parallel p_\theta(\mathbf{x}_{t-1} \mid \mathbf{x}_t))$ , 정답 분포 $q$ 와 모델이 예측한 가우시안 분포 $p_{\theta}$ 가 얼마나 비슷한지를 측정하며, 이 항이 모델이 학습할(오차를 줄일) 대상이다. 이때 $q$ 쪽에 원본 이미지 $\mathbf{x}_0$ 가 들어와 있는데, 이는 아래 식 (6)에서 설명한다.

$L_0$ : $-\log p_\theta(\mathbf{x}_0 \mid \mathbf{x}_1)$ , $t=1$ 에서 마지막으로 노이즈를 한 번 더 걷어내어 최종 원본 이미지( $\mathbf{x}_0$ )를 만들어낼 때의 오차(재구성 오차)이다. KL발산이 아닌 우도형태이며, 생성된 실수값(연속형 픽셀 값)을 0~255 사이의 정수값(이산형)으로 매핑하는 디코더를 거친다.

$L_0 = -\log p_\theta(\mathbf{x}_0 \mid \mathbf{x}_1)$ 이 재구성 오차인 이유.

디퓨전 모델은 계층이 1,000개인 특수한 계층적 VAE이다. VAE의 목적 함수를 떠올려보자.

VAE에서 잠재 변수 $z$ 를 입력받아 최종 원본 데이터 $x$ 를 복원해 내는 디코더의 확률값이 바로 $p_\theta(x \mid z)$ 입니다. 이 확률에 음의 로그( $-\log$ )를 취한 것은 정보 이론에서 크로스 엔트로피를 의미하며, 이는 생성된 결과물과 실제 데이터 사이의 오차(loss)를 나타낸다.

디퓨전 수식에서 $\mathbf{x}_1$ 은 원본 데이터 $\mathbf{x}_0$ 에 닿기 직전의 마지막 잠재 변수입니다. 즉, $\mathbf{x}_1$ 이 VAE의 $z$ 역할을 하고, $\mathbf{x}_0$ 가 $x$ 역할을 한다. 따라서 $-\log p_\theta(\mathbf{x}_0 \mid \mathbf{x}_1)$ 는 마지막 잠재 변수에서 원본 데이터를 재구성할 때 발생하는 오차를 측정하는 항이다.

다른 $L_{t-1}$ 항들과 구조적으로도 차이가 있는데, $t > 1$ 일 때의 $L_{t-1}$ 항들은 모두 가우시안 분포끼리의 비교이다. 잠재 공간에서 확률 분포 거리를 조정하는 작업이다. 반면 $t=1$ 에서 $t=0$ 으로 넘어가는 순간 $\mathbf{x}_0$ 은 관측된(결정론적인) 실제 데이터이다. 따라서 두 확률 분포 사이의 거리를 재는 KL 발산 대신 주어진 $\mathbf{x}_1$ 상태에서 진짜 데이터 $\mathbf{x}_0$ 가 등장할 확률(likelihood)을 측정하는 $-\log p$ 형태를 취해야 한다.

섹션 3.3에서 언급되지만, 모델이 예측한 연속적인 가우시안 분포와 관측된 이산형 데이터를 연결하기 위해 독립적인 이산 디코더를 도입하였다. 연속적인 가우시안 분포의 확률 밀도 함수를 특정 픽셀 값의 아주 좁은 구간에 대해 적분하여 이산적인 확률값으로 반환한다.

$\mathbf{x}_0$ 를 조건으로 한다는 것은, 편향된 추정이 아닌가? X

조건은 q에만 있음. 모델이 추정하는 $p_{\theta}$ 에는 없음.