Score Function

log pdf에 대해 입력 $\bold{x}$에 관한 1차 gradient 함수 $\nabla_{\bold{x}}\log p_{\boldsymbol{\theta}}(\bold{x})$을 (Stein) score function이라고 부르며 다음과 같이 정의한다. (pmf의 경우 미분이 직접적으로 적용되지 않으므로 일반적으로 pdf에 대해 사용된다.)

$$ \bold{s}{\boldsymbol{\theta}}(\bold{x}) \triangleq \nabla{\bold{x}} \log p_{\boldsymbol{\theta}}(\bold{x}) $$

EBM의 식에 대해 입력에 관한 1차 gradient를 취하면 다음과 같이 정리된다.

$$ \nabla_{\bold{x}} \log p_{\boldsymbol{\theta}}(\bold{x}) = -\nabla_{\bold{x}} \mathcal{E}{\boldsymbol{\theta}}(\bold{x}) - \underbrace{\nabla{\bold{x}} \log Z_{\boldsymbol{\theta}}}{=0} = -\nabla{\bold{x}} \mathcal{E}_{\boldsymbol{\theta}}(\bold{x}) $$

결국 EBM에 대해 Score 함수를 이용하면 계산이 까다로운 정규화 상수를 계산하지 않고도 모델을 학습할 수 있다.

Score Matching(SM)

Score Matching은 데이터의 score 함수와 모델의 score 함수를 구하고 그 둘 사이의 차이를 최소화하는 방법을 말한다. Score Matching 목적은 Fisher Divergence라 부르는 두 분포의 score 함수 사이의 차이를 최소화한다.

$$ D_F(p_\mathcal{D}(\bold{x})\|p_{\boldsymbol{\theta}}(\bold{x})) = \mathbb{E}{p\mathcal{D}(\bold{x})}\left[ {1\over2}\|\nabla_\bold{x} \log p_\mathcal{D}(\bold{x}) - \nabla_\bold{x} \log p_{\boldsymbol{\theta}}(\bold{x})\|^2 \right] $$

참고로 Score Matching과 Contrastive Divergence는 매우 다른 접근으로 보이지만 서로 밀접하게 연결되어 있다. 실제로 Score Matching은 특정한 MCMC 샘플러의 극한에서 Contrastive Divergence의 special case로 볼 수 있다.

Basic Score Matching

특정한 regularity 조건 아래 Fisher divergence에서 $p_\mathcal{D}(\bold{x})$의 알려지지 않은 1차 도함수를 $\mathcal{E}_{\boldsymbol{\theta}}(\bold{x})$의 2차 도함수로 교체하여 부분별 적분을 사용할 수 있다.

$$ \begin{aligned} D_F(p_\mathcal{D}(\bold{x})\|p_{\boldsymbol{\theta}}(\bold{x})) &= \mathbb{E}{p\mathcal{D}(\bold{x})}\left[{1\over2}\sum_{i=1}^d\left({\partial \mathcal{E}{\boldsymbol{\theta}}(\bold{x}) \over \partial x_i} - {\partial^2 \mathcal{E}{\boldsymbol{\theta}}(\bold{x})\over \partial x_i^2} \right)^2 \right] + \text{const} \\ &= \mathbb{E}{p\mathcal{D}(\bold{x})}\left[{1\over2}\|\bold{s}{\boldsymbol{\theta}}(\bold{x})\|^2 + \text{tr}(\bold{J}\bold{x}\bold{s}_{\boldsymbol{\theta}}(\bold{x})) \right] + \text{const} \end{aligned} $$

여기서 $d$는 $\bold{x}$의 차원이고 $\bold{J}\bold{x}\bold{s}{\boldsymbol{\theta}}(\bold{x})$는 score 함수의 야코비안이다. 상수는 최적화에 영향이 없으므로 학습하는 동안 제거할 수 있다.

위 목적의 단점은 야코비안의 trace를 계산하는데 $O(d^2)$ 시간이 걸린다는 것이다. 이러한 이유로 위 방정식의 암시적 SM 공식은 2차 도함수 계산이 다루기 용이한 상대적으로 단순한 에너지 함수에만 적용되었다.

Denoising Score Matching(DSM)

기본 Score Matching의 목적은 어디에서나 연속이고 미분가능하고 유한해야 한다는 조건을 요구하는데, 이미지 데이터는 픽셀 값이 0-255로 이산이기 때문에 이것을 직접 사용할 수 없다. 이를 완화하기 위한 방법으로 데이터 포인트에 노이즈를 추가할 수 있는데, 이 방법을 Denoising Score Matching(DSM)이라 한다.

이를 위해 데이터 포인트에 다음과 같이 노이즈를 추가한다.

$$ \tilde{\bold{x}} = \bold{x} + \boldsymbol{\epsilon} $$