[평범한 학부생이 하는 논문 리뷰] Generative Modeling by Estimating Gradients of the Data Distribution (NCSN)

이번에는 Diffusion에 제대로 도전해보자! 하는 마인드로 Diffusion 논문들도 블로그에 올리기로 다짐했다.

그래서 첫 논문으로 NCSN을 들고 왔다. 스터디해보면서 엄청 벽을 느낀 논문들이지만 다시 읽고 제대로 이해해서 넘어갈겸 포스팅에 도전했다.

https://arxiv.org/pdf/1907.05600.pdf

 

1. Introduction

1.1 기존 생성모델들의 단점

기존 생성모델들은 각 모델들마다 단점을 지닌다. 예를 들면 likelihood-based 모델은  autoregressive model이나 flow model과 같이 정규화된 확률 모델을 구축하기 위해 특수 아키텍처를 사용하거나 학습을 위해 VAE에서의 ELBO와 같은 surrogate loss를 사용해야 한다. 그리고 GAN 모델들은 적대적인 훈련 방식때문에 훈련이 불안정할 수 있다.

 

이 논문에서는 이러한 단점을 해결하기 위해 score matching을 이용한 새로운 생성모델을 소개한다.

NCSN의 구조를 설명하기 전에 본 논문의 저자들이 소개한 이 모델의 바람직한 특성들이 아래와 같다.

• objective는 특별한 제약 조건이나 아키텍처 없이도 score network의 거의 모든 parameterizations에 대해 tractable하며, 훈련 중에 적대적 훈련, MCMC 샘플링 또는 기타 근사치 없이도 최적화할 수 있다.
• objective는 동일한 데이터 세트에서 서로 다른 모델을 정량적으로 비교하는 데에도 사용할 수 있다.
• 샘플이 최신 likelihood-based 모델 및 GAN에서 생성된 샘플과 비슷하게 보인다.

2. Scored-based generative modeling

난 이 논문을 읽을 때 score가 뭐고 score matching이 뭘까하는 의문이 많이 들었고 이해를 아예 못했었다. 그래서 이 파트를 설명하기 전에 score가 뭔지를 알아보자.

Score의 의미를 살펴보면 모든 $x$에 대해 $x$에 대해 log likelihood의 gradient를 취하는 것으로 data space에서 $x$의 likelihood를 증가시키는 방향을 나타낸다고 볼 수 있다.

자세한 설명은 아래 논문에 나와있다.

https://arxiv.org/pdf/2208.11970.pdf

Energy-based 모델을 살펴보면 우리가 알기 쉬운 $f_{\theta}(x)$를 가지고 우리가 원하는 $p_{\theta}(x)$를  구해보자라는 아이디어에서 나왔다고 이해했다. 그래서 우리가 원하는 분포를 아래와 같이 정의 할 수 있다.

$$p_{\theta}(x) = \frac{e^{-f_{\theta}(x)}}{Z_{\theta}}, \quad Z_{\theta} = \int e^{- f_{\theta}(x)}$$

위에 정의된 분포를 MLE 추정법을 통해 $\theta$를 추정해서 분포를 찾을 수 있다. 하지만 위 분포 함수에서 $Z_{\theta}$를 계산하는 것이 어렵기에 $p_{\theta}(x)$대신 ${\rm log}p_{\theta}(x)$를 사용하면 $Z_{\theta}$ term이 분리되고 $x$에 무관하기 때문에 $x$에 대한 gradient가 0이 된다. 그렇기에 neural network $s_{\theta}(x)$를 $p_{\theta}(x)$를 학습하게 하기 보다는 score $\nabla_x{\rm log}p_{\theta}(x)$를 학습하게 하는 것을 목표로 한다.

 

Score network $s_θ$는 $\theta$로 매개변수화된 신경망으로, 이 신경망은 $p_{data}(x)$의 score를 근사하도록 훈련되고 최종적으로 $p_{data}(x)$에서 새로운 샘플을 생성하기 위한 모델을 학습하는 것이다. 

2.1 Score matching for score estimation

Score matching은 우리가 모르는 데이터 분포인 $p_{data}(x)$의 score인 ${\rm log}p_{data}(x)$를 추정하는 score network $s_{\theta}(x)$를 학습시키는 것이다. 즉, score network를 통해 $p_{\theta}(x)$의 score를 $p_{data}(x)$의 score로 근사시켜서 $p_{data}(x)$를 추정하는 것이라 볼 수 있다. 

Score-based model의 최종 objective function은 아래와 같다.

$$\frac{1}{2} \mathbb{E}_{p_{data}} [\Vert s_θ(x) - ∇_x log p_{data}(x)\Vert_2^2 ]$$​

이를 정리하면, 아래와 같이 나온다.

$$\mathbb{E}_{p_{data}(x)}[ tr(∇_xs_θ(x))+\frac{1}{2} \Vert s_\theta(x)\Vert _2^2]​, \quad (1)$$

손으로 풀어봤는데 악필인 점 양해 부탁드립니다...ㅎㅎ

Score matching은 $tr(∇_xs_θ(x))$를 계산해야 하기 때문에 deep network와 고차원 데이터로 확장할 수 없다. 그래서 다음으로 대규모 score matching을 위해 널리 사용되는 두 가지 방법에 대해 설명하겠다.

 

Denoising score matching

Denoising score matching은 $tr(∇_xs_θ(x))$를 완전히 우회하는 score matching의 변형이다. 먼저 데이터 포인트 $x$를 미리 지정된 노이즈 분포 $q_σ(\tilde{x} | x)$로 교란한 다음 score matching을 사용하여 perturbed된 데이터 분포 $q_σ(\tilde{x})\triangleq \int q_σ(\tilde{x} | x)p_{data}(x)dx$의 score를 추정한다. Objective는 다음과 같다.

$$\frac{1}{2} \mathbb{E}_{q_σ (\tilde{x}|x)p_{data}(x)} [\Vert s_\theta(\tilde{x})-\nabla_{\tilde{x}}log q_{\sigma}(\tilde{x}\vert x)\Vert_2^2]$$

위 식을 적절히 변형하면 아래의 식이 나온다. 증명은 아래와 같다.

$$\frac{1}{2}\mathbb{E}_{q_{\sigma}(\tilde{x})}[\Vert s_{\theta}(\tilde{x}) - \nabla_{\tilde{x}}{\rm log}q_{\sigma}(\tilde{x})\Vert^2_2]\quad (2)$$

식 (2)를 최소화하는 optimal score network($s_{θ^∗} (x)$)는 대부분 $s_{θ^∗} (x) = ∇_x log q_σ(x)$를 만족시킨다고 한다. 그러나 $s_{θ ^∗} (x) = ∇_x log q_σ(x) ≈ ∇_x log p_{data}(x)$는 노이즈$\sigma$가 충분히 작은 경우에만 참이다.

 

Sliced score matching

Sliced score matching은 random projection을 사용하여 점수 매칭에서 $tr(∇_xs_θ(x))$를 근사화한다. Objective는 아래와 같다.

$$\mathbb{E}_{p_v} \mathbb{E}_{p_{data}} [\mathbf{v^\top} \mathbf{∇_x s_θ(x) v} + \frac{1}{2}\Vert \mathbf{s_θ (x)}\Vert_2^2] , \quad (3)$$

여기서 $p_v$는 다변량 표준 정규분포와 같은 무작위 벡터의 단순 분포이다. $\mathbf{v^{\top}∇_x s_θ(x) v}$항은 forward mode auto-differentiation을 통해 효율적으로 계산할 수 있습니다. Sliced score matching은 perturbed되지 않은 데이터 분포에 대한 점수 추정을 제공하지만 forward mode auto-differentiation으로 인해 약 4배 더 많은 계산을 필요로 한다.

식 (3)이 식 (1)과 동일하게 나온다는 것은 아래와 같이 증명할 수 있다.

 

2.2 Sampling with Langevin dynamics

Langevin dynamics은 score function $∇_x log p(x)$만을 사용하여 확률 밀도 $p(x)$로부터 샘플을 생성할 수 있다. $\epsilon > 0$ 인 고정된 스텝과 초기값 $\tilde{x}_0 ∼ π(x)$가 주어지면 Langevin 방법은 다음과 같이 재귀적으로 계산된다. (여기서 $\pi (x)$는 uniform distribution과 같은 단순한 prior distribution이다.)

$$\tilde{x}_t = \tilde{x}_{t-1} + \frac{\epsilon}{2} ∇_x log p(\tilde{x}_{t-1}) + \sqrt{\epsilon} z_t, \quad (4)$$

여기서 $z_t ∼ \mathcal{N}(0, I)$입니다. $\tilde{x}_T$의 분포는 $\epsilon → 0$ 및 $T → ∞$일 때 $p(x)$와 같으며, 이 경우 $\tilde{x}_T$는 일부 regularity 조건 하에서 $p(x)$의 정확한 표본이 된다.

식 (4)에서 샘플링하려면 score function $∇_x log p(x)$ 만 필요한다. 따라서 $p_{data}(x)$로부터 샘플을 얻으려면 먼저 $s_θ(x) ≈ ∇_x log p_{data}(x)$가 되도록 Score network를 훈련한 다음 $s_θ(x)$를 사용하여 Langevin dynamics으로 샘플을 대략적으로 얻을 수 있다. 이것이 Score-based generative modeling 프레임워크의 핵심 아이디어다.

3. Challenges of score-based generative modeling

위에 소개된 방식을 그냥 사용하게되면 두 가지 문제점이 발생한다고 한다. 그 두 가지 문제점에 대해 알아보자.

3.1 The manifold hypothesis

The manifold hypothesis 하에서 score-based generative models는 두 가지 주요 어려움에 직면하게 된다.

1. Score $∇_x log p_{data}(x)$는 ambient 공간(주위, 에워싼 공간) 에서 취한 gradient이기 때문에 $x$가 저차원 manifold에 국한되면 정의되지 않는다.
2. Score matching objective 식 (1)은 데이터 분포의 support가 전체 공간일 때만 일관된 score estimator를 제공하며, 데이터가 저차원 manifold에 있을 때는 일관성이 떨어진다고 한다.

매니폴드 가설이 점수 추정에 미치는 부정적인 영향은 그림 1에서 명확하게 볼 수 있다. 그림 1(왼쪽)에서 볼 수 있듯이, 원본 CIFAR-10 이미지로 훈련하면 sliced score matching loss가 먼저 감소한 후 불규칙하게 변동한다. 반대로 교란된 데이터 분포가 $\mathbb{R}^D$에 대한 full support를 가지도록 작은 가우시안 노이즈로 데이터를 교란하면 손실 곡선이 수렴한다. 

3.2 Low data density regions

저밀도 영역의 데이터 부족은 score matching을 사용한 score 추정과 Langevin dynamics를 사용한 MCMC 샘플링 모두에 어려움을 초래할 수 있다.

 

3.2.1 Inaccurate score estimation with score matching

데이터 밀도가 낮은 지역에서는 데이터 샘플이 부족하기 때문에 score 매칭으로 score 함수를 정확하게 추정하기에 충분한 증거가 없을 수 있다. 

Score 매칭이 score 추정치의 예상 제곱 오차, 즉 $\frac{1} {2}\mathbb{E}_{p_{data}} [\Vert s_θ(x) - ∇_x log p_{data}(x)\Vert^2_2 ]$를 최소화한다는 점을 상기해보자. 실제로, 데이터 분포에 대한 기대값은 항상 i.i.d. 샘플 $\{x_i\}^N_{i =1} \overset{i.i.d}{\sim} p_{data}(x)$를 사용하여 추정된다. 그런데 $p_{data}(\mathcal{R}) ≈ 0$인 모든 영역 $\mathcal{R} ⊂ \mathbb{R} ^D$를 고려해보면, 대부분의 경우 $\{x_i\}^ N_{ i=1} ∩ \mathcal{R} = \emptyset$이고, 점수 매칭은 $x ∈ \mathcal{R}$에 대해 $∇_x log p_{data}(x)$를 정확하게 추정하기에 충분한 데이터 샘플을 갖지 못한다.

이러한 부정적인 효과를 입증하기 위해 toy experiments를 수행했다고 한다. 그 결과를 보면 아래 그림에서 볼 수 있듯이, 실제 데이터가 많이 분포해 있는 진한 부분의 score는 어느 정도 정확하지만 그 외 부분은 정확하지 않은 것을 볼 수 있다.

 

3.2.2 Slow mixing of Langevin dynamics

데이터 분포의 두 모드가 저밀도 영역으로 분리되어 있는 경우, Langevin dynamics은 합리적인 시간 내에 두 모드의 상대적 가중치를 올바르게 복구할 수 없으므로 실제 분포에 수렴하지 못할 수 있다.

혼합 분포 $p_{data}(x) = πp_1(x)+(1-π)p_2(x)$에서 $p_1(x)$와 $p_2(x)$는 서로 다른 support를 갖는 normalized distribution이고 $π∈ (0, 1)$이라고 가정하자.

$p_1(x)$의 support에서 $∇_x log p_{data}(x) = ∇_x(log π + log p_1(x)) = ∇_x log p_1(x)$이고,

$p_2(x)$의 support에서 $∇_x log p_{data}(x) = ∇_x(log(1 - π) + log p_2(x)) = ∇_x log p_2(x)$ 이다.

Langevin dynamcs은 $∇_x log p_{data}(x)$를 사용하여 $p_{data}(x)$에서 표본을 추출하므로, 얻은 표본은 $π$에 의존하지 않는다.

실제로 이 분석은 서로 다른 모드가 동일한 support를 공유하지만 데이터 밀도가 작은 영역으로 연결되는 등 거의 분리된 support를 갖는 경우에도 적용된다고 한다. 이 경우 Langevin dynamic은 이론적으로 올바른 샘플을 생성할 수 있지만, 매우 작은 스텝 크기와 매우 많은 수의 스텝을 혼합해야 할 수 있다고 한다. 아래 (a)와 (b)를 통해 결과를 확인 할 수 있는데, 두 모드 간의 상대적 가중치가 잘못 된 것을 볼 수 있다.

4. Noise Conditional Score Networks : learning and inference

무작위 가우시안 노이즈로 데이터를 교란하면 데이터 분포가 score-based 생성 모델링에 더 적합하다는 것을 관찰했다. 첫째, 가우스 노이즈 분포의 support가 전체 공간이기 때문에 교란된 데이터는 저차원 manifold에 국한되지 않으므로 manifold hypothesis의 어려움을 없애고 점수 추정을 잘 정의할 수 있다. 둘째, 큰 가우시안 노이즈는 원래의 방해받지 않은 데이터 분포에서 저밀도 영역을 채우는 효과가 있으므로 점수 매칭이 더 많은 훈련 신호를 얻어 점수 추정을 개선할 수 있다. 더 나아가 여러 노이즈 레벨을 사용하면 실제 데이터 분포에 수렴하는 일련의 noise-perturbed 분포들을 얻을 수 있다. Simulated annealing 및 annealed importance 샘플링에 따라 이러한 intermediate 분포를 활용하여 다중 모드 분포에서 Langevin dynamics의 mixing rate를 개선할 수 있다.( 여기서 multimodal 분포의 의미는 위 이미지처럼 mode가 여러 개인 분포를 의미하는 것으로 이해했다.)

1) 다양한 수준의 노이즈를 사용하여 데이터를 교란하고, 2) 단일 조건부 score 네트워크를 훈련하여 모든 노이즈 수준에 해당하는 score를 동시에 추정함으로써 점수 기반 생성 모델링을 개선할 것을 제안한다. 훈련 후 Langevin dynamics을 사용하여 샘플을 생성할 때 처음에는 큰 노이즈에 해당하는 점수를 사용하고 점차적으로 노이즈 수준을 낮춘다. 이렇게 하면 큰 노이즈 레벨의 이점을 교란된 데이터가 원래 데이터와 거의 구분할 수 없는 낮은 노이즈 레벨로 원활하게 전달할 수 있다. 

4.1 Noise Conditioning Score Networks

$\{σ_i\} ^L_ {i=1}$은 $\frac{σ_1} {σ_2} = \cdots = \frac{σ_{L-1}}{σ_L} > 1$을 만족하는 양의 기하학적 수열, $q_{\sigma}(x) \triangleq \int p_{data}(t)\mathcal{N} (x | t, \sigma^2 I)\,dt$는 perturbed 데이터 분포를 나타낸다. 잡음 수준$\sigma_1$이 섹션 3에서 언급한 어려움을 완화할 수 있을 만큼 충분히 크고, $σ_L$이 데이터에 미치는 영향을 최소화할 수 있을 만큼 충분히 작도록 잡음 수준 $\{σ_i\}^L_{i=1}$를 선택한다. 논문에서 조건부 점수 네트워크를 훈련하여 모든 교란된 데이터 분포의 점수를 공동으로 추정하는 것을 목표로 한다(즉, $∀σ ∈ \{σ_i\}^L_{i=1} : s_θ(x, σ) ≈ ∇_x log q_σ(x)$ ) . $x \in \mathbb{R}^D$일 때 $s_θ(x, σ) \in \mathbb{R}^ D$이다. 이를 Noise Conditional Score Network(NCSN)라고 부른다.

 

Architecture

NCSN $s_θ(x, σ)$의 출력은 입력 이미지 $x$와 동일한 형태를 갖기 때문에, sementic segmentation에서 영감을 받아 모델 $s_θ(x, σ)$는 U-Net의 아키텍처 구조와 dilated/atrous 컨볼루션을 결합한다. 또한, 일부 이미지 생성 작업에서의 우수한 성능을 보인 instance normalization를 채택했으며, 조건부 instance normalization의 수정된 버전을 사용하여 $σ_i$ 에 대한 conditioning을 제공한다. 

위에서 설명한 세가지 구조적 특징을 자세하게 알아보자.

 

Instance Normalization

---

논문에서 적용한 conditional instance normalization을 알아보기 전에 conditional instance normalization은 뭐고 instance normalization은 뭘까??

Instance normalization은 input $x \in \mathbb{R}^{B \times C \times H \times W}$가 주어질 때, 하나의 배치와 채널 내의 instance에서 평균($\mu_{nc}(x)$)과 표준편차($\sigma_{nc}(x)$)를 구해 그것으로 normalization하는 것으로 이해했다.

$$z = γ\frac{ x - µ}{\sigma} + β$$

Conditional normalization은 위의 가중치 파라미터를 동일하게 주는 것이 아니라 conditional하게 주는 것으로 이해했다.

$$z_k = γ[i, k]\frac{ x_k - µ_k}{ s_k} + β[i, k]$$

---

Conditional Instance Normalization을 사용하여 점수를 예측할 때 $s_θ(x, σ)$가 $σ$를 고려하도록 했다고 한다. Conditional Instance Normalization에서는 다른 set의 scale과 biases에 대해 다른 $σ ∈ \{σ_i\}^ L_{ i=1}$이 사용된다. 더 구체적으로, $x$가 $C$ 특징 맵을 가진 입력이라고 가정하자. $µ_k$와 $s_k$는 공간 축을 따라 취한 $x$의 $k$번째 특징 맵의 평균과 표준 편차를 나타냅니다. 조건부 인스턴스 정규화는

$z_k = γ[i, k]\frac{ x_k - µ_k}{ s_k} + β[i, k]$

로 이루어지며, 여기서 $γ ∈ \mathbb{R}^{ L×C}$ 및 $β ∈\mathbb{R}^{ L×C}$는 학습 가능한 파라미터이고, $k$는 특징 맵의 인덱스, $i$는 $\{σ_i\}^L_{i=1}$에서 $σ$의 인덱스를 나타낸다.

그러나 인스턴스 정규화의 한 가지 단점은 다른 특징 맵에 대한 $µ_k$의 정보를 완전히 제거한다는 것입니다. 이로 인해 생성된 이미지의 색상이 바뀔 수 있다. 이 문제를 해결하기 위해 Conditional Instance Normalization에 대한 간단한 수정을 제안한다. 먼저 $µ_k$의 평균 $m$과 표준 편차$v$를 계산한다. 그런 다음 학습 가능한 파라미터 $α ∈ \mathbb{R}^{ L×C}$ 를 추가한다. 수정된 Conditional Instance Normalization는

$$z_k = γ[i, k]\frac{ x_k - µ_k}{ s_k} + β[i, k] + α[i, k]\frac{ µ_k - m} {v}$$

로 정의된다. 위 Conditional Instance Normalization을 CondInstanceNorm++로 부른다. 논문의 아키텍처에서는 모든 컨볼루션 레이어와 풀링 레이어 앞에 CondInstanceNorm++를 추가한다.

 

Dilated 컨볼루션

Dilated 컨볼루션은 특징 맵의 해상도를 유지하면서 receptive 필드의 크기를 늘리는 데 사용할 수 있다. 더 큰 해상도의 특징 맵을 사용하여 위치 정보를 더 잘 보존하기 때문에 semantic segmentation에 매우 효과적인 것으로 나타났다. NCSN의 아키텍처 설계에서는 첫 번째 레이어를 제외한 모든 서브샘플링 레이어를 대체하는 데 사용한다.

 

U-Net 아키텍처

$s_θ(x, σ)$를 구축하기 위해, 이전 정보를 skip connection을 통해 전달하는 아키텍처를 지닌 U-Net과 ResNet 설계를 통합한 RefineNet의 아키텍처를 사용합니다.

실험에서는 4-cascaded RefineNet을 사용하고, pre-activation residual block을 사용한다.

좌측의 이미지가 기존의 resnet이라면 우측의 이미지는 활성화 함수의 순서를 바꾼 pre-activation resnet을 나타냄.

RefineNet 아키텍처에서 모든 배치 정규화를 제거하고 이를 CondInstanceNorm++로 대체한다. 스타일 전송과 같은 이미지 생성 작업에서 average pooling이 더 부드러운 이미지를 생성하는 것으로 보고됨에 따라 Refine 블록의 max pooling 레이어를 average pooling으로 대체한다. 또한 원래 RefineNet 블록에는 정규화가 사용되지 않지만 각 컨볼루션 및 average pooling 전에 CondInstanceNorm++을 추가한다. 모든 활성화 함수는 ELU로 선택했다. 일반적인 관행에 따라 다음 캐스케이드로 진행할 때 확장을 2배로 늘린다. 

4.2 Learning NCSNs via score matching

Sliced 및 denoising score matching 모두 NCSN을 훈련할 수 있다. 여기서는 denoising score matching이 약간 더 빠르고 노이즈가 있는 데이터 분포의 score를 추정하는 작업에 자연스럽게 적합하기 때문에 denoising score matching을 채택했다고 한다.

Noise 분포는 $q_σ(\tilde{x} | x) = N (\tilde{x} | x, σ^2 I)$ 이고, 따라서 $∇_{\tilde{x}} log q_σ(\tilde{x} | x) = -(\tilde{x}-x)/{σ^2}$ 이다. 주어진 $σ$에 대해 denoising score matching objective(방정식 (2))는 아래와 같다.

$$\ell(\theta;\sigma) \triangleq \frac{1}{2} \mathbb{E}_{p_{data}(x)}\mathbb{E}_{\tilde{x}\sim\mathcal{N}(x,\sigma^2I)}[\Vert s_{\theta}(\tilde{x},\sigma)+\frac{\tilde{x}-x}{\sigma^2}\Vert^2_2] \quad (5)$$

방정식 (5)를 결합하면 아래와 같은 하나의 통일된 objective function을 얻을 수 있다.

$$\mathcal{L}(\theta,\{\sigma_i\}^L_{i=1}) \triangleq \frac{1}{L}\underset{i=1}{\overset{L}\Sigma}\lambda(\sigma_i)\ell(\theta;\sigma_{i})\quad (6)$$

여기서 $\lambda(\sigma_i) > 0$은 $\sigma_i$에 따른 계수 함수다. $s_θ(x, σ)$의 capacity가 충분하다고 가정하면, 식 (6)은 모든 $i ∈\{1, 2, \cdots , L\}$에 대해 $L$개의 denoising score matching objective의 조합이기 때문에 $s_{θ^*} (x, σ_i) = ∇_x log q_{σ_i} (x)$ 인 경우에만 식 (6)을 최소화한다.

모든 $\{σ_i\}^L_{i=1}$에 대한 $λ(σ_i)\ell(\theta;\sigma_i)$ 의 값은 거의 같은 크기인 것이 이상적이라고 한다. $λ(σ) = σ^2$ 를 선택하면, 경험적으로 score 네트워크가 최적성으로 훈련될 때 대략 $\Vert s_θ(x, σ)\Vert^2 ∝ 1/σ$ 를 갖는 것으로 보여지기 때문에 $\lambda(\sigma)\ell(\theta;\sigma) = \sigma^2\ell(\theta;\sigma)=\frac{1}{2}\mathbb{E}[\Vert \sigma s_{\theta}(\tilde{x},\sigma) + \frac{\tilde{x}-x}{\sigma}\Vert^2_2]$가 된다. $\frac{\tilde{x}-x}{ σ} ∼ \mathcal{N} (0, I)$ 및 $\Vert σs_θ(x, σ)\Vert^2∝ 1$이므로, $λ(σ)\ell(θ; σ)$의 magnitude의 order는 $σ$에 의존하지 않는다는 결론을 쉽게 내릴 수 있다.

Objective 식 (6)은 적대적 훈련, 대리 손실, 훈련 중 점수 네트워크에서 샘플링이 필요하지 않다는 점을 논문에서 강조한다. 또한 tractable하기 위해 $s_θ(x, σ)$가 특별한 구조를 가질 필요도 없다. 또한, $λ(\cdot)$와 $\{σ_i\} ^L_ {i=1}$이 고정되어 있으면 서로 다른 NCSN을 정량적으로 비교하는 데 사용할 수 있다.

4.3 NCSN inference via annealed Langevin dynamics

NCSN $s_θ(x, σ)$ 가 훈련된 후에는 annealed Langevin dynamics을 샘플링 방식으로 제안한다. Alg. 1에서 볼 수 있듯이, uniform noise과 같은 고정된 prior 분포로부터의 샘플을 초기화하여 annealed Langevin dynamics을 시작한다. 그런 다음 Langevin dynamics을 실행하여 step size $α_1$로 $q_{σ_1} (x)$에서 샘플을 추출한다. 다음으로 Langevin dynamics을 실행하여 이전 시뮬레이션의 최종 샘플에서 시작하여 감소된 step size $\alpha_2$를 사용하여 $q_{σ_2}(x)$에서 샘플링한다. 이 방식으로 계속하여 $q_{σ_{i-1}} (x)$에 대한 Langevin dynamics의 최종 샘플을 $q_{σ_i} (x)$에 대한 Langevin dynamics의 초기 샘플로 사용하고 $α_i = \epsilon \cdot σ^2_i /σ^2_L$로 step size $α_i$를 점차적으로 낮춘다. 마지막으로 Langevin dynamics을 실행하여 $σ_L ≈ 0$일 때 $p_{data}(x)$에 가까운 $q_{σ_L}(x)$에서 샘플을 추출한다.

분포 $\{q_{σ_i}\}^ L_{ i=1}$은 모두 gaussian noise에 의해 교란되기 때문에 supports가 전체 공간에 걸쳐 있고 score가 잘 정의되어 있어 manifold hypothesis의 어려움을 피할 수 있다. $σ_1$이 충분히 크면 $q_{σ_1}(x)$의 저밀도 영역이 작아지고 모드가 덜 고립된다. 이렇게 하면 score 추정이 더 정확해지고 Langevin dynamics의 혼합이 더 빨라질 수 있다. 따라서 Langevin dynamics이 $q_{σ_1}(x)$에 대한 좋은 샘플을 생성한다고 가정할 수 있다. 이러한 샘플은 $q_{σ_1}(x)$의 고밀도 영역에서 나올 가능성이 높으며, 이는 $q_{σ_1}(x)$와 $q_{σ_2}(x)$가 서로 약간만 다르다는 점을 감안할 때 $q_{σ_2}(x)$의 고밀도 영역에도 존재할 가능성이 높다는 것을 의미한다. 점수 추정과 Langevin dynamics은 고밀도 영역에서 더 잘 수행되므로, $q_{σ_1}(x)$의 샘플은 $q_{σ_2}(x)$의 Langevin dynamics에 대한 좋은 초기 샘플이 될 것입니다. 마찬가지로, $q_{σ_{i-1}} (x)$는 $q_{σ_i} (x)$에 대한 좋은 초기 샘플을 제공하며, 마지막으로 $q_{σ_L} (x)$에서 좋은 품질의 샘플을 얻습니다.

Alg.1에서 $σ_i$에 따라 $α_i$를 조정하는 방법에는 여러 가지가 있을 수 있다. 논문에서는 Langevin dynamics에서 $\frac{α_is_θ(x,σ_i)} {2 \sqrt{α_i} z}$의 크기를 고정하기 위해서 $α_i ∝ σ ^2_ i$를 선택했다. 앞에서 언급했듯이 $\Vert s_θ(x, σ)\Vert^2 ∝ 1/σ$를 발견했으므로  $\mathbb{E}[\Vert \frac{α_is_θ(x,σ_i)}{2\sqrt{α_i} z}\Vert^2_2 ] ≈ \mathbb{E}[ \frac{α_i\Vert s_θ(x,σ_i)\Vert ^2_2} {4} ] ∝ \frac{1} {4} \mathbb{E}[\Vert σ_is_θ(x, σ_i)\Vert^ 2_2 ]∝ \frac{1} {4}$가 되고 이는 $σ_i$에 의존하지 않게 된다. 이렇게 되면 앞에서 의도한 크기의 고정이 된다!

Alg. 1을 적용하여 3.2절에서 사용한 gaussian mixture에서 샘플링해보면, 위 그림 3 (b)와 (c)를 비교하면 annealed Langevin dynamics은 두 모드 간의 상대적 가중치를 올바르게 복구하는 반면 표준 Langevin dynamics은 실패한다.

5. Experiments

Setup

실험에서는 MNIST, CelebA [34], CIFAR-10 [31] 데이터 세트를 사용한다. CelebA의 경우, 이미지를 먼저 140 × 140으로 중앙에서 잘라낸 다음 32 × 32로 크기를 조정한다. 모든 이미지의 픽셀 값이 $[0, 1]$이 되도록 크기를 조정한다. $L = 10$개의 $\{σ_i\}^L_{i=1},σ_1 = 1, σ_{10} = 0.01$ 인 기하학적 수열인 서로 다른 표준 편차를 선택한다. 이미지 생성에 annealed Langevin dynamics을 사용할 때는 $T = 100$ 및 $\epsilon = 2 × 10^{-5}$ 를 선택하고 uniform noise를 초기 샘플로 사용한다. 결과는 $T$ 선택에 관계없이 견고하며, 일반적으로$\epsilon$은 $5 × 10^{-6}$에서 $5 × 10^{-5}$ 사이가 잘 작동하는 것으로 나타났다.

Image generation

그림 5에 나와있는 생성한 이미지는 최신 확률 기반 모델 및 GAN의 이미지보다 더 높거나 비슷한 품질을 가지고 있다. Annealed Langevin dynamics의 절차를 직관적으로 이해할 수 있도록 그림 4에서 중간 샘플을 제공하며, 각 행은 샘플이 순수한 랜덤 노이즈에서 고품질 이미지로 진화하는 과정을 보여준다. 아래의 그림9, 10에는 훈련 데이터 세트에서 생성된 이미지의 가장 가까운 이웃을 보여줌으로써 모델이 단순히 훈련 이미지를 암기하는 것이 아님을 보여준다. 여러 노이즈 레벨에 대해 조건부 점수 네트워크를 공동으로 학습하고 annealed Langevin dynamics을 사용하는 것이 중요하다는 것을 보여주기 위해 하나의 노이즈 레벨 $\{σ_1 = 0.01\}$만 고려하고 vanilla Langevin dynamics 샘플링 방법을 사용하는 기준 접근 방식과 비교한다. 이 작은 노이즈 추가는 매니폴드 가설의 어려움을 피하는 데 도움이 되지만, 데이터 밀도가 낮은 영역의 점수에 대한 정보를 제공하기에는 충분히 크지 않다. 결과적으로 이 baseline은 그림7, 8에서 볼 수 있듯이 합리적인 이미지를 생성하지 못한다.

정량적 평가를 위해 CIFAR-10의 inception 및 FID score를 측정했고 아래와 같은 결과가 나왔다.

Image inpainting

그림 6에서는 score network가 일반화 가능하고 의미론적으로 의미 있는 이미지 표현을 학습하여 다양한 이미지 인페인팅을 생성할 수 있음을 보여준다.  annealed Langevin dynamics 절차를 아래의 알고리즘에 보이듯이 간단하게 수정하여 임의의 형태의 occlusions이 있는 이미지를 자연스럽게 처리할 수 있다고 한다. 

6. Related Work

이번 섹션은 자세한 내용은 생락하고 논문에서 소개한 NCSN이 이전 방식들과 어떤 점이 다른지만 소개하고 넘어가겠다.

1. 훈련 중에 마르코프 체인에서 샘플링할 필요가 없다.
2. 훈련과 샘플링 방법이 서로 분리되어 있다.
3. 에너지 기반 모델의 기울기를 점수 모델로 사용하여 에너지 기반 모델(EBM)을 훈련하는 데 사용할 수 있다.

 

이번 논문 리뷰를 마무리하며...

Diffusion논문을 리뷰하는 것은 이번이 처음이였다. Diffusion의 기초를 다루는 논문이고 기초는 항상 중요하다고 생각하기에 최대한 꼼꼼하게 공부하려 했다. 물론 완전히 이해했냐고 묻는다면 아니라고 대답할 것 같다... 그래도 이전에 처음 읽었을 때 보다는 더 많은 것을 이해한 것 같다. 조만간 코드와 이 포스팅을 다시보면서 공부를 해야겠다는 생각이 들었다.

첫 diffusion논문 리뷰를 그래도 나름 성공적으로 마친 것 같아서 뿌듯하다 ㅎㅎ 

 

(스터디 진행 후, 다시 읽어보면서 수정예정입니다.)