score based model 정리

확산모델의 수학 책을 읽고 정리한 내용

공부한 내용 정리
생성형 모델은 데이터의 분포 $p(x)$를 따라한 $q_\theta(x)$를 만드는 것이 목적이다.
데이터의 분포를 따라한 분포에서 sampling을 할 수 있으면 데이터를 생성하는 것이 가능하기 때문.

에너지 기반 모델

이때 일반적인 확률 분포 $q_\theta(\bold x)$를 확률 밀도 함수로 표현이 가능한데
$q_\theta(\bold x)=\gamma_\theta(\bold x)/Z(\theta)$라고 하는데
여기에서 $\gamma_\theta(\bold x)\ge0$로 확률밀도 함수이다.
아래 $Z(\theta)=\int_{x'\in X}\gamma_\theta(\bold x')d\bold x'$로 확률분포 $q_\theta(\bold x)$를 확률의 조건인 합 1을 맞추기 위해서 전체 확률밀도 함수의 가능한 모든 output 합을 나눠주는 것이다.

그런데 이때 $\gamma_\theta$는 $\ge0$의 조건이 달려 있기에 다루기 까다롭다. 그렇기에 에너지 함수라는 것이 등장하는데
에너지 함수는 다음과 같다.
$\gamma_\theta(\bold x)=\exp(-f_\theta(\bold x))$로 e의 내부에 임의의 함수 $f_\theta$를 가지고 기존 $\gamma_\theta$의 조건인 0보다 크게만드는 것을 만족시킨다..

에너지 함수를 이용해서 표현하면 다음과 같다.
$q_\theta(x)=\exp(-f_\theta(\bold x))/Z(\theta)$

이제 에너지 함수를 자유롭게 구성해서 확률분포의 모델링이 가능해졌다.

이때 energy 함수를 보면 알 수 있겠지만 에너지가 낮으면 sampling 확률이 높고 에너지가 높으면 sampling 확률이 낮다.

우도 증가?

에너지 함수를 이용한 분포에서 sampling을 어떻게 할 수 있을까?
sampling을 위해서는 학습 데이터 $D=\{x^{(1)},x^{(2)},...,x^{(N)}\}$인 데이터를 잘 모델링 하는 분포여야 하는데
이를 계산하기 위해서 우도(likelihood)로 표현을 한다.
이때 데이터 전체의 우도는 각 학습 데이터는 독립적이기에 각 데이터 우도의 곱으로 표현
$q_\theta(D)=\prod_iq_\theta(x^{(i)})$

곱을 다루기 쉽게 log 우도로 표현하면
$L(\theta)=\frac{1}{N}\log q_\theta(D)=\frac{1}{N}\sum_i\log q_\theta(x^{(i)})$로 표현이 된다.

이를 높이는 param이 우도를 최대로 하는 param이 되고 sampling을 하였을 때 학습 데이터의 분포에서 잘 생성할 수 있을 것이다.

위 우도 식에 처음 구한 에너지 함수 식을 넣으면
$L(\theta)=-\frac{1}{N}\sum^N_{i=1}[f_\theta(\bold x^{(i)})]-\log\int_{\bold x'\in X}\exp(-f_\theta(\bold x'))d\bold x'$이 된다.

위 식을최대화 하려면
첫번째 항의 의미는 학습 데이터 주변 energy를 낮춰서 sampling이 잘 이루어지게 되어야 하고
두번째 항의 의미는 전체 data의 분포에서는 energy를 늘려야 한다는 것이다.

경사 하강법
보통 param의 학습이 경사 하강법으로 이루어지기에 이를 $\theta$에 대해서 미분하면 다음과 같다.
$\frac{\partial L(\theta)}{\partial \theta} = - \frac{1}{N} \sum_{i=1}^N \left[ \frac{\partial f_\theta(\mathbf{x}^{(i)})}{\partial \theta} \right] - \frac{\partial}{\partial \theta} \log Z(\theta) \\ = - \frac{1}{N} \sum_{i=1}^N \left[ \frac{\partial f_\theta(\mathbf{x}^{(i)})}{\partial \theta} \right] - \frac{1}{Z(\theta)} \int- \frac{\partial f_\theta(\mathbf{x})}{\partial \theta} \exp(-f_\theta(\mathbf{x})) \, d\mathbf{x} \\= - \frac{1}{N} \sum_{i=1}^N \left[ \frac{\partial f_\theta(\mathbf{x}^{(i)})}{\partial \theta} \right]+ \mathbb{E}_{\mathbf{x} \sim q_\theta(\mathbf{x})} \left[ \frac{\partial f_\theta(\mathbf{x})}{\partial \theta} \right]$
위에서 첫번는 훈련 데이터 에너지의 기울기이고 두번째는 sampling 기댓값 $\bold x \sim q_\theta(\bold x)$의 에너지 기울기인데
두번째 부분에서 sampling을 해야하기에 매우 계산량이 많아지고 복잡해진다.
그렇기에 score를 사용해서 진행한다.

score

log 우도 $\log p(\bold x)$에 대한 기울기를 score 점수라고 가정하고 $\bold s(\bold x)$를 score 함수라고 한다.
$\bold s(\bold x):=\nabla_\bold x\log p(\bold x)$라고 한다.

이 score를 통해서 우도가 증가하는 방향으로 data를 이동해서 생성할 수 있다.
랑주뱅 몬테카를로를 적용하면
$\bold x_k:=\bold x_{k-1}+\alpha \nabla_\bold x\log p(\bold x_{k-1})+\sqrt{2\alpha}\bold u_k, \quad \bold u_k\sim\mathcal N(0,I)$
위처럼 noise $\bold u_k$를 추가하는 동시에 score가 증가하는 방향으로 데이터를 옮겨서 우도가 높은 즉, 데이터를 그럴듯한 분포로 맞출 수 있고 생성할 수 있게 된다.

이렇게 확률 분포 대신에 점수를 학습해서 생성 모델을 구현하는 방법이 Score based model이다.

Noise를 넣어주는 이유는 너무 극점으로 가는 것을 막기 위함과 동시에 다른 분포를 탐색할 확률을 주기 위함이다.
(한 꼭짓점으로 모이면 sampling이 수렴하는 것도 있다)

score matching

위에서 score를 통해서 데이터를 생성할 수 있음을 보였는데 그럼 score는 어떻게 학습할 수 있을까?

Explicit score matching

우선 바로 학습하는 명시적 점수 매칭을 봐보면
$J_{\text{ESM}_p}(\theta) = \frac{1}{2} \mathbb{E}_{p(\mathbf{x})} \left[ \left\lVert \nabla_{\mathbf{x}} \log p(\mathbf{x}) - s_\theta(\mathbf{x}) \right\rVert^2 \right]$이렇게 MSE로 바로 score model을 score와 같은 값을 가질 수 있게 학습한다.

이는 별로 설명할 것이 없다.

Implicit score matching

하지만 위에서 우리는 score를 알기가 매우 어렵다. 왜냐하면 data의 분포 $\nabla_\bold x \log p(x)$를 알아야 하는데 이를 모르기 때문...

그렇기에 $\nabla_\bold x \log p(x)$를 사용하지 않고 학습 목표를 설정하는데
이는 암묵적 점수 매칭으로 다음과 같이 정의가 된다.

$J_{\text{ISM}_p}(\theta) = \mathbb{E}_{p(\mathbf{x})} \left[ \frac{1}{2} \lVert s_\theta(\mathbf{x}) \rVert^2 + \text{tr}(\nabla_{\mathbf{x}} s_\theta(\mathbf{x})) \right]$
위 식인데 tr은 trace로 대각합이다.

$\nabla_\bold xs_\theta(\bold x)$는 헤세 matrix로 2계도함수 matrix.

이걸 학습하면 explit을 학습하는 것과 동일한데 왜 그런걸까?
두 식의 유도는 매우 복잡한데
가정이 4가지 필요하다

• $p(x)$가 미분 가능
• $\mathbb E_{p(x)}[||\nabla_\bold x \log p(\bold x)||^2]$이 유한
• 임의의 $\theta$에 대하여 $\mathbb E_{p(x)}[||s_\theta(\bold x)||^2]$이 유한
• $\lim_{||\bold x||\rightarrow \infin}[p(\bold x)s_\theta(\bold x)]=0$

이면 $J_{\text {ESM}}=J_{\text {ISM}}(\theta)+C_1$을 만족

이때 위 두 식이 만족하면 학습하는 것이 동일한데 왜냐하면 상수가 달라져도 두 함수의 최적 값은 만족하는 param은 동일할 것이기 때문

우선 explicit의 내부를 전개하면 다음과 같다.
$J_{\text{ESM}}(\theta) = \frac{1}{2} \mathbb{E}_{p(\mathbf{x})} \left[ \lVert \nabla_{\mathbf{x}} \log p(\mathbf{x}) - s_\theta(\mathbf{x}) \rVert^2 \right] \\= \int_{\mathbf{x} \in \mathbb{R}^d} p(\mathbf{x}) \left[ \frac{1}{2} \lVert \nabla_{\mathbf{x}} \log p(\mathbf{x}) \rVert^2 + \frac{1}{2} \lVert s_\theta(\mathbf{x}) \rVert^2 - \nabla_{\mathbf{x}} \log p(\mathbf{x})^\top s_\theta(\mathbf{x}) \right] d\mathbf{x}.$
여기에서 첫번째 항은 상수이기에 $C_1$이고 두번째 항은 implicit에도 동일하기에 그대로 두고 마지막 항이 $\text{tr}(\nabla_{\mathbf{x}} s_\theta(\mathbf{x}))$와 동일한지만 보면 된다.

마지막항은 내적이기 때문에 다음과 같이 표기가 가능하다.
$\int_{\bold x\in \mathbb R^d} -p(\bold x) \nabla_\bold x\log p(\bold x)^\top s_\theta(\bold x)d\bold x=-\sum_i \int_{\bold x\in \mathbb R^d} p(\bold x)\frac{\partial \log p(\bold x)}{\partial x_i} s_\theta(\bold x)_id\bold x$
여기에서 $s_\theta(\bold x)_i$는 i번째 항을 의미한다.

이후 ISM의 항인 $\text{tr}(\nabla_{\mathbf{x}} s_\theta(\mathbf{x}))$는 $\sum_i\int_{\bold x\in \mathbb R^d}\frac{\partial s_\theta(\bold x)_i}{\partial x_i}p(\bold x)d\bold x$이기 때문에

$- \int_{\mathbf{x} \in \mathbb{R}^d} \frac{\partial p(\mathbf{x})}{\partial x_i} s_\theta(\mathbf{x})_i \, d\mathbf{x} = \int_{\mathbf{x} \in \mathbb{R}^d} \frac{\partial s_\theta(\mathbf{x})_i}{\partial x_i} p(\mathbf{x}) \, d\mathbf{x}$가 만족하는지만 보이면 된다.

i=1일때 증명은 같이 된다.
$- \int_{\mathbf{x} \in \mathbb{R}^d} \frac{\partial p(\mathbf{x})}{\partial x_1} s_\theta(\mathbf{x}; \theta)_1 d\mathbf{x} \\ = - \int_{\mathbf{x}_{2:d} \in \mathbb{R}^{d-1}} \left[ \int_{x_1 \in \mathbb{R}} \frac{\partial p(\mathbf{x})}{\partial x_1} s_\theta(\mathbf{x}; \theta)_1 dx_1 \right] d\mathbf{x}_{2:d} \\ = - \int_{\mathbf{x}_{2:d} \in \mathbb{R}^{d-1}} \left[ \lim_{a \to \infty, b \to -\infty} \left[p(a, \mathbf{x}_{2:d}) s_\theta(a, \mathbf{x}_{2:d})_1 - p(b, \mathbf{x}_{2:d}) s_\theta(b, \mathbf{x}_{2:d})_1 \right] \right. \\ \left. - \int_{x_1 \in \mathbb{R}} \frac{\partial s_\theta(\mathbf{x}; \theta)_1}{\partial x_1} p(\mathbf{x}) dx_1 \right] d\mathbf{x}_{2:d} \\ = - \int_{\mathbf{x}_{2:d} \in \mathbb{R}^{d-1}} \left[ 0 - 0 - \int_{x_1 \in \mathbb{R}} \frac{\partial s_\theta(\mathbf{x}; \theta)_1}{\partial x_1} p(\mathbf{x}) dx_1 \right] d\mathbf{x}_{2:d} \\ = \int_{\mathbf{x} \in \mathbb{R}^d} \frac{\partial s_\theta(\mathbf{x}; \theta)_1}{\partial x_1} p(\mathbf{x}) d\mathbf{x}.$

내부 증명은 부분적분으로 이루어진다.

증명 끝.

denoising score matching

디노이징은
$\tilde\bold x=\bold x+\epsilon, \quad \epsilon\sim\mathcal N(0,\sigma^2I)$인 $\tilde \bold x$의 분포를 가지고 score를 학습하고 matching을 진행하는 것.

위에서 정리하면 $p_\sigma(\tilde \bold x|\bold x)=\mathcal N(\tilde \bold x;\bold x,\sigma^2I)$이다.
이때의 score은 다음과 같이 학습이 되는데
$J_{\text{ESM}_{p_\sigma}}(\theta) = \frac{1}{2} \mathbb{E}_{p_\sigma(\tilde{\mathbf{x}})} \left[ \left\lVert \nabla_{\tilde{\mathbf{x}}} \log p_\sigma(\tilde{\mathbf{x}}) - s_\theta(\tilde{\mathbf{x}}, \sigma) \right\rVert^2 \right]$

역시 ESM은 복잡하다.
ISM은 이전과 동일하게 유도가 됨으로 생략

DSM은 다음과 같이 나온다.
$J_{\text{DSM}_{p_\sigma}}(\theta) = \frac{1}{2} \mathbb{E}_{p_\sigma(\tilde{\mathbf{x}} \mid \mathbf{x}) p(\mathbf{x})} \left[ \left\lVert \nabla_{\tilde{\mathbf{x}}} \log p_\sigma(\tilde{\mathbf{x}} \mid \mathbf{x}) - s_\theta(\tilde{\mathbf{x}}, \sigma) \right\rVert^2 \right]$
여기에서
$\nabla_{\tilde{\mathbf{x}}} \log p_\sigma(\tilde{\mathbf{x}} \mid \mathbf{x}) = \nabla_{\tilde{\mathbf{x}}} \log \left( \frac{1}{(2\pi)^{d/2} \sigma^d} \exp\left(-\frac{1}{2\sigma^2} \lVert \tilde{\mathbf{x}} - \mathbf{x} \rVert^2 \right) \right) \\ = \nabla_{\tilde{\mathbf{x}}} \log \frac{1}{(2\pi)^{d/2} \sigma^d} + \nabla_{\tilde{\mathbf{x}}} \left( -\frac{1}{2\sigma^2} \lVert \tilde{\mathbf{x}} - \mathbf{x} \rVert^2 \right) \\ = 0 - \frac{1}{\sigma^2} (\tilde{\mathbf{x}} - \mathbf{x}) \\ = -\frac{1}{\sigma^2} \boldsymbol{\epsilon}$
위와 같이 되어서
$J_{\text{DSM}_{p_\sigma}}(\theta) = \frac{1}{2} \mathbb{E}_{\boldsymbol{\epsilon} \sim \mathcal{N}(0, \sigma^2 \mathbf{I}), \mathbf{x} \sim p(\mathbf{x})} \left[ \left\lVert -\frac{1}{\sigma^2} \boldsymbol{\epsilon} - s_\theta(\mathbf{x} + \boldsymbol{\epsilon}, \sigma) \right\rVert^2 \right]$
이렇게 DSM은 noise를 맞추는 문제가 된다.

DSM 증명

역시 ESM을 풀면 다음과 같고
$J_{\text{ESM}_{p_\sigma}}(\theta) = \frac{1}{2} \mathbb{E}_{p_\sigma(\tilde{\mathbf{x}})} \left[ \left\lVert \nabla_{\tilde{\mathbf{x}}} \log p_\sigma(\tilde{\mathbf{x}}) - s_\theta(\tilde{\mathbf{x}}, \sigma) \right\rVert^2 \right] \\ = \frac{1}{2} \mathbb{E}_{p_\sigma(\tilde{\mathbf{x}})} \left[ \lVert s_\theta(\tilde{\mathbf{x}}, \sigma) \rVert^2 \right] - \mathbb{E}_{p_\sigma(\tilde{\mathbf{x}})} \left[ \langle s_\theta(\tilde{\mathbf{x}}, \sigma), \nabla_{\tilde{\mathbf{x}}} \log p_\sigma(\tilde{\mathbf{x}}) \rangle \right] + C_2$
(이떄 <>기호는 내적)
DSM을 풀면 다음과 같다.
$J_{\text{DSM}_{p_\sigma}}(\theta) = \frac{1}{2} \mathbb{E}_{p_\sigma(\tilde{\mathbf{x}} \mid \mathbf{x}) p(\mathbf{x})} \left[ \left\lVert \nabla_{\tilde{\mathbf{x}}} \log p_\sigma(\tilde{\mathbf{x}} \mid \mathbf{x}) - s_\theta(\tilde{\mathbf{x}}, \sigma) \right\rVert^2 \right] \\ = \frac{1}{2} \mathbb{E}_{p_\sigma(\tilde{\mathbf{x}})} \left[ \lVert s_\theta(\tilde{\mathbf{x}}, \sigma) \rVert^2 \right] - \mathbb{E}_{p_\sigma(\tilde{\mathbf{x}}, \mathbf{x})} \left[ \langle s_\theta(\tilde{\mathbf{x}}, \sigma), \nabla_{\tilde{\mathbf{x}}} \log p_\sigma(\tilde{\mathbf{x}} \mid \mathbf{x}) \rangle \right] + C_3$
여기에서 첫항은 동일하니 제외하고 두번째 항이 동일한지만 보면 된다.

이때 $\mathbb{E}_{p_\sigma(\tilde{\mathbf{x}})} \left[ \langle s_\theta(\tilde{\mathbf{x}}, \sigma), \nabla_{\tilde{\mathbf{x}}} \log p_\sigma(\tilde{\mathbf{x}}) \rangle \right]$를 풀면 다음과 같은데
$=\int_{\tilde\bold x}p_\sigma(\tilde \bold x)\langle s_\theta(\tilde\bold x,\sigma),\nabla_{\tilde \bold x}\log p_\sigma(\tilde \bold x) \rangle d\tilde \bold x$
의 내부 $\nabla_{\tilde \bold x} \log p_\sigma(\tilde \bold x)=\frac{\nabla_{\tilde \bold x}p_\sigma(\tilde \bold x)}{p_\sigma(\tilde \bold x)}$그렇기에

$\mathbb{E}_{p_\sigma(\tilde{\mathbf{x}})} \left[ \langle s_\theta(\tilde{\mathbf{x}}, \sigma), \nabla_{\tilde{\mathbf{x}}} \log p_\sigma(\tilde{\mathbf{x}}) \rangle \right]=\int_{\tilde\bold x}\langle s_\theta(\tilde \bold x,\sigma),\nabla_{\tilde \bold x}p_\sigma(\tilde \bold x) \rangle d\tilde \bold x$가 된다.

이후 내부 미분이
$\nabla_{\tilde{\mathbf{x}}} p_\sigma(\tilde{\mathbf{x}}) = \nabla_{\tilde{\mathbf{x}}} \int_{\mathbf{x}} p(\mathbf{x}) p_\sigma(\tilde{\mathbf{x}} \mid \mathbf{x}) d\mathbf{x} \\ = \int_{\mathbf{x}} p(\mathbf{x}) \nabla_{\tilde{\mathbf{x}}} p_\sigma(\tilde{\mathbf{x}} \mid \mathbf{x}) d\mathbf{x} \\ = \int_{\mathbf{x}} p(\mathbf{x}) p_\sigma(\tilde{\mathbf{x}} \mid \mathbf{x}) \nabla_{\tilde{\mathbf{x}}} \log p_\sigma(\tilde{\mathbf{x}} \mid \mathbf{x}) d\mathbf{x}$
를 만족하게 되어 이를 대입하면

$S(\theta) = \int_{\tilde{\mathbf{x}}} \left\langle s_\theta(\tilde{\mathbf{x}}, \sigma), \int_{\mathbf{x}} p(\mathbf{x}) p_\sigma(\tilde{\mathbf{x}} \mid \mathbf{x}) \nabla_{\tilde{\mathbf{x}}} \log p_\sigma(\tilde{\mathbf{x}} \mid \mathbf{x}) d\mathbf{x} \right\rangle d\tilde{\mathbf{x}} \\ = \int_{\tilde{\mathbf{x}}} \int_{\mathbf{x}} p(\mathbf{x}) p_\sigma(\tilde{\mathbf{x}} \mid \mathbf{x}) \left\langle s_\theta(\tilde{\mathbf{x}}, \sigma), \nabla_{\tilde{\mathbf{x}}} \log p_\sigma(\tilde{\mathbf{x}} \mid \mathbf{x}) \right\rangle d\mathbf{x} d\tilde{\mathbf{x}} \\ = \int_{\tilde{\mathbf{x}}} \int_{\mathbf{x}} p_\sigma(\tilde{\mathbf{x}}, \mathbf{x}) \left\langle s_\theta(\tilde{\mathbf{x}}, \sigma), \nabla_{\tilde{\mathbf{x}}} \log p_\sigma(\tilde{\mathbf{x}} \mid \mathbf{x}) \right\rangle d\mathbf{x} d\tilde{\mathbf{x}} \\ = \mathbb{E}_{p_\sigma(\tilde{\mathbf{x}}, \mathbf{x})} \left[ \left\langle s_\theta(\tilde{\mathbf{x}}, \sigma), \nabla_{\tilde{\mathbf{x}}} \log p_\sigma(\tilde{\mathbf{x}} \mid \mathbf{x}) \right\rangle \right]$가 나오게 된다.
증명 끝.

정리

결국 score based model은 우도의 미분 값 즉,score를 학습을 하고 이를 통해서 data를 우도가 높은 곳으로 이동시켜 생성을 한다.
이를 위해서 score를 학습하는 방법으 explicit, implicit이 있으며 noise를 넣은 data로부터 noise를 예측하는 식으로 진행을 하는 denoising score matching을 알아보았다.