Author: Y Song
Publication: https://arxiv.org/abs/2101.03288
Diffusion model의 입문용 논문으로 추천돼있어 아는 내용 갈무리 겸 읽어보게 되었다. 각 섹션의 개요만 간단하게 정리해본다.
Energy-Based Models (EBMs)
Energy function: Unnormalized negative log-probability
Density estimation은 nonlinear regression 문제로 귀납된다.
Probability density: $p_\theta(x) = {{-E_\theta(x)} \over {Z_\theta}}$
E는 parameter $\theta$에 대한 nonlinear regression 함수이다.
Normalizing constant (partition function): $ Z_\theta = \int \exp (-E_\theta(x) )dx $
Z는 $\theta$에 관한 함수이지만 x에 대해 상수이고, $\log p_{\theta}(x)$에 대한 미분은 intractable integral을 포함한다.
Maximum Likelihood Training with MCMC
Expected log-likelihood function:
$$ \mathit{E}_{x \sim p{data}(x)} [\log p_\theta(x)] $$
이것을 최댓값으로 만드는 것은 data와 predicted distribution 사이의 KL divergence를 최소화하는 것과 같다.
$$ -\mathit{E}_{x \sim p{data}(x)} [\log p_\theta(x)] = D_{KL}(p_{data}(x) || p_{\theta}(x)) - E_{x \sim p_{data}(x)} [\log p_{data}(x)] $$
두 번째 항은 parameter에 의존하지 않기 때문에 상수이다.
Z 때문에 likelihood를 직접적으로 계산할 수는 없다. 그래서 MCMC 같은 방법론으로 log-likelihood의 gradient를 추정한다.
$$ \nabla_\theta \log p_\theta(x) = -\nabla_\theta E_\theta(x) - \nabla_\theta \log Z_\theta $$
첫 번째 항은 automatic diffferenciation으로 손쉽게 계산할 수 있지만, 두 번째 항은 intractable하다. 복잡한 수식을 거치면 다음을 도출할 수 있다.
$$ \nabla_\theta\log Z_\theta \simeq -\nabla_\theta E_\theta(\tilde x), \tilde x \sim p_\theta(x) $$
즉, $\tilde x$는 EBM에서 주어진 확률에서 랜덤 샘플링을 해서 얻은 샘플을 말한다.
Log-likelihood gradient의 unbiased Monte Carlo 추정값을 얻을 수 있으므로 stochastic gradient ascent 같은 걸로 parameter를 최적화할 수 있다.
MCMC는 너무나 계산량이 많기 때문에 Langevin MCMC 같은 방법론들이 제안되었다.
다시 한 번, log-probability (score)의 gradient는 energy의 gradient에 음수를 취한 값이라는 것을 상기한다.
$$ \nabla_x \log p_\theta(x) -\nabla_x E_\theta(x) - \nabla_x \log Z_\theta = -\nabla_x E_\theta(x) $$
Gaussian distribution (simple prior distribution)에서 initial sample $x^0$를 샘플링하면, Langevin diffusion process를 시뮬레이션할 수 있다. ($K$: step, $\epsilon$: step size)
$$ x^{k+1} \leftarrow x^k + {{\epsilon^2} \over {2}} \nabla_x \log p_\theta(x^k) + \epsilon z^k $$
두 번째 항은 Energy의 gradient에 음수를 취한 것과 같다.
Step size가 충분히 작고, step 수가 충분히 많다면 discretization error는 무시 가능하다.
Score matching
가장 쉬운 score matching objective는 data와 model distribution의 Fisher divergence를 최소화하는 것이다.
$$ D_F(p_{data}(x)||p_\theta(x)) = E_{p_{data}(x)} [{{1}\over{2}} || \nabla_x \log p_{data}(x) - \nabla_x \log p_\theta (x)||^2] $$
두 번째 항 계산은 $\nabla_x \log p_{data}(x)$를 알아야 쓸 수 있어 실용적이지 못하다. Energy function의 second derivative를 이용해 다시 쓰면,
$$ \mathbb{E}_{p{data}(x)} [{{1}\over{2}} \sum^d_{i=1} ({{\partial^2E_\theta(x)} \over {(\partial x_i)^2}})^2] + C $$
위 식의 문제점은 second derivative 계산이 data dimensionality $d$에 quardratic하다. 더 높은 차원이 데이터에 대해선 계산이 어려워진다. (위 식은 간단한 energy function에 대해서만 사용한다.)
Score matching은 continuous data distribution을 가정하지만, discrete / bounded data distribution에도 일반화될 수 있다. (image pixel 값이 {0, … 255}인 경우)
Denoising score matching (DSM)
각 데이터 포인트에 noise를 더해보자.
$$ \tilde x = x + \epsilon $$
Score matching objective는 $D_F(q(\tilde x)|| p_\theta(\tilde x)$가 된다. (noise distribution $q(\tilde x)$는 smooth해야 한다.)
DSM: Sample들에 대한 empirical average에 의해 기댓값을 근사하면 $p_{data}(x)$와 second derivative 계산을 둘다 피할 수 있다.
$$ D_F(q(\tilde x)|| p_\theta(\tilde x) = E_{q(x, \tilde x)} [{{1}\over{2}} || \nabla_x \log q(\tilde x|x) - \nabla_x \log p_\theta (\tilde x)||^2_2] + C $$
Noise를 더하는 것은 $p_{data}(x)$가 이미 score matching에 필요한 조건을 잘 만족하는 분포일 때 생긴다. $D_F(q(\tilde x)|| p_\theta(\tilde x) \neq D_F(p_{data}(x)|| p_{\theta}(\tilde x))$이기 때문에 DSM은 좋은 objective가 아니다. Optimal EBM은 data 분포가 아닌 noisy 분포 $q(\tilde x)$로 매칭되기 때문이다. 두 분포가 많이 다르면 문제가 된다.
해결 방법은 noise를 아주 조금 첨가하는 것이다. $q(\tilde x|x) = \mathit{N} (\tilde x| x, \sigma^2I), \sigma \sim 0$. 결과적으론 DSM의 분산이 커지고 최적화가 어려워진다.
Wang (2020)에서는 분산이 0에 가까울 때 DSM의 분산을 줄일 방법을 찾았다. Variance-reducing 변수, control variates $c_\theta (x,z) = {{2} \over {\sigma}} z^T \nabla _x \log p_\theta (x) + {{||z||^2_2} \over {\sigma^2}} - {{d} \over {\sigma^2}}$를 원래 식에서 빼주는 것이다.
Score-based generative models
Score matching의 문제점은 large low-density region의 두 가지 분리된 modes의 relative weight를 예측하는 데 있다.
Song (2021)을 보면 training data를 다른 scale의 분산으로 perturbing함으로써 해결했다.
- 작은 분산일 때는 두 modes가 더 끊겨 있지만
- 큰 분산일 때는 두 modes가 더 연결되어 있다 (길을 잃지 않도록 방향을 잘 찾아가게 해줌)
Langevin dynamics나 reverse diffusion process를 사용하면 noise-perturbed distribution을 먼저 샘플링하도록 해준다. 그 다음에 천천히 noise scale의 양을 줄이면서 작은 것에 가까워지도록 한다.
Noise contrastive estimation (NCE)
Gutmann and Hyvarinen (2010)에 의해 제안된 EBM을 알려진 density를 가지는 다른 분포와 대조(contrast)하여 학습하는 아이디어이다.
- $p_{data}(x)$: data distribution
- $p_{n}(x)$: chosen distribution with known density (noise distribution)
- Simple and has a tractable PDF
- 예) Gaussian distribution $N(0,I)$
- $p_\theta$: EBM
Binary classification 문제를 푼다고 가정하고 $y$는 Bernoulli 분포를 따르는 {0, 1} 중 하나인 변수이다.
- Data distribution
$$ p_{n,data}(x) = p(y=0)p_n(x) + p(y=1)p_{data}(x)$$ - y=0의 posterior 확률
$$ p_{n,data}(y=0|x) = {{p_{n,data}(x|y=0)p(y=0)}\over{p_{n,data}(x)}} = {{p_n(x)}\over{p_n(x) + \nu p_{data}(x)}} \text{, } \nu=p(y=1)/p(y=0)$$ - EBM 정의
$$ p_\theta (x) = \exp(-E_\theta (x)) / Z_\theta$$- $Z_\theta$는 학습 가능한 scalar parameter
- Mixture of noise and model distribution
$$ p_{n,\theta}(y=0|x) = {{p_n(x)}\over{p_n(x)+\nu p_\theta(x)}}$$
$p_{n,\theta}(y|x)$를 $p_{n,data}(y|x)$에 피팅시키는 식으로, 간접적으로 $p_\theta$를 $p_{data}$에 피팅시킬 수 있다.
이 때 사용하는 conditional MLE objective는 아래와 같고, stochastic gradient ascent로 풀 수 있다.
$$ \theta^* = \text{argmin}_\theta \mathbb{E}_{p_{n,data}(x)} [D_{KL} (p_{n,data}(y|x)||p_{n,\theta}(y|x))] = \text{argmax}_\theta \mathbb{E}_{p_{n,data}(x,y)}[\log p_{n,\theta}(y|x)]$$
모델이 충분히 깊으면 optimum에서 다음이 만족한다.
$$ p_{n,\theta^*} (y=0|x) \equiv p_{n,data}(y=0|x) \\ \iff {{p_{n}}\over{p_n(x)+\nu p_{\theta^*}(x)}} \equiv {{p_{n}}\over{p_n(x)+\nu p_{data}(x)}} \\ \iff p_{\theta^*}(x) \equiv p_{data}(x) $$
결과적으로 optimum에서 E는 $p_{data}(x)$에 일치하는 unnormalized energy function이 된다.
NCE는 학습의 부산물로서 EBM의 normalizing constant를 준다.
실사용에서 NCE의 성공을 결정하는 건 noise distribution을 잘 고르는 것이다. (structured 고차원 데이터를 사용할 때 특히) NCE는 noise와 data distribution이 가까울 때 성공한다. Noise distribution을 알아서 튜닝하는 방법들도 나와있다.
- Score matching과의 연관성: NCE는 noise distribution과 $\nu$에 따라 objecctive가 달라진다. 이런 유동성 때문에 NCE로부터 score matching을 끌어낼 수 있다.
References
'Papers > Generative Models' 카테고리의 다른 글
| [논문 리뷰] Protein Design with Guided Discrete Diffusion (NIPS 2023) (0) | 2023.10.22 |
|---|
