EBM에서 샘플링을 위해 $\bold{x}$에 관한 gradient를 취하면 $Z_{\boldsymbol{\theta}}$는 $\bold{x}$와 무관하므로 $\nabla_\bold{x} \log Z_{\boldsymbol{\theta}}=0$이 되어 $\nabla_{\bold{x}} \mathcal{E}_{\boldsymbol{\theta}}(\bold{x})$ 만 계산하면 된다.
$$ \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}) $$
Langevin MCMC는 에너지 함수의 gradient에 노이즈를 더한 형태로 위의 식과 유사하기 때문에 Langevin MCMC를 사용하면 EBM에서 쉽게 샘플링할 수 있다. 이를 위해 우선 단순 prior 분포에서 초기 샘플 $\bold{x}^0$을 뽑고, 그 다음에 step size $\epsilon >0$를 사용하여 $K$ 단계에 대해 overdamped 랑주뱅 diffusion 프로세스를 시뮬레이션 한다.
$$ \bold{x}^{k+1} \leftarrow \bold{x}^k + {\epsilon^2 \over 2} \underbrace{\nabla_{\bold{x}}\log p_{\boldsymbol{\theta}}(\bold{x}^k)}{=-\nabla\bold{x}\mathcal{E}_{\boldsymbol{\theta}}(\bold{x})} + \epsilon \bold{z}^k \\ k=0,1,...,K-1 $$
여기서 $\boldsymbol{z}^k \sim \mathcal{N}(\bold{0},\bold{I})$은 가우시안 노이즈 항이다. $\epsilon \to 0$이고 $K \to \infty$ 일 때 $\bold{x}^K$는 어떤 regularity 조건 아래 $p_{\boldsymbol{\theta}}(\bold{x})$로 분포하는 것이 보장된다.
그러나 샘플 $\bold{x} \sim p_{\boldsymbol{\theta}}(\bold{x})$을 얻기 위해 수렴할 때까지 MCMC를 실행하는 것은 계산적으로 비싸기 때문에 체인이 목표에 수렴할 때까지 실행하지 않고 고정된 수의 단계만 MCMC 절차를 수행하는 대안 방법을 사용할 수 있는데 이것이 바로 Contrastive Divergence(CD)이다.
CD의 $T$ 단계는 다음 목적을 최소화하는 것과 동등하다.
$$ \text{CD}T = D\text{KL}(p_0\|p_\infty) - D_\text{KL}(p_T\|p_\infty) $$
여기서 $p_T$는 $T$ MCMC 업데이트 이후 $\bold{x}$에 대한 분포이고 $p_0$는 데이터 분포이고 $p_\infty$는 모델이 수렴했을 때에 대한 이론적 분포이다. (일반적으로 작은 $T$ 값을 사용하여 좋은 결과를 얻을 수 있고 때때로 $T=1$을 설정하기도 한다.)
$D_\text{KL}(p_0\|p_\infty)$는 실제 데이터 분포 $p_0$와 모델이 수렴했을 때의 분포 $p_\infty$ 사이의 KL divergence로 모델이 실제 데이터를 얼마나 잘 표현하고 있는지를 나타낸다.
$D_\text{KL}(p_T\|p_\infty)$는 MCMC의 $T$ 단계 후에 얻은 분포 $p_T$와 모델이 수렴했을 때의 분포 $p_\infty$ 사이의 KL divergence로 $p_T$가 $p_\infty$에 얼마나 근접한지를 나타낸다.
CD는 목적은 단순히 모델 분포와 데이터 분포 사이의 차이를 최소화하는게 아니라 두 KL divergence 사이의 차를 최소화하는 것으로, 모델 분포가 데이터 분포 뿐만 아니라 MCMC 수행 후의 분포에도 가까워야 한다는 목적이 된다. 이것은 데이터 분포와 모델 분포 사이의 ‘energy gab’을 최소화하기 위한 전략으로 설계 되었다.
Model은 간단히 구현한다. 신경망의 모델은 하나의 함수로 볼 수 있으며 이 경우 에너지 함수에 해당한다.
# EBM 모델 정의. 이것을 energy 함수라고 할 수 있다.
class EBM(nn.Module):
def __init__(self):
super(EBM, self).__init__()
self.model = nn.Sequential(
nn.Linear(28*28, 128),
nn.ReLU(),
nn.Linear(128, 1),
)
def forward(self, x):
return self.model(x.view(x.size(0), -1))
입력에 대해 Langevin MCMC를 수행한 후에 CD 목적에 맞게 손실을 구하는 함수를 정의한다.