Search
Duplicate

AI/ Contrastive Divergence(CD)

Contrastive Divergence(CD)

EBM에서 샘플링을 위해 x\bold{x}에 관한 gradient를 취하면 ZθZ_{\boldsymbol{\theta}}는 x\bold{x}와 무관하므로 ∇xlog⁡Zθ=0\nabla_\bold{x} \log Z_{\boldsymbol{\theta}}=0이 되어 ∇xEθ(x)\nabla_{\bold{x}} \mathcal{E}_{\boldsymbol{\theta}}(\bold{x}) 만 계산하면 된다.
∇xlog⁡pθ(x)=−∇xEθ(x)−∇xlog⁡Zθ⏟=0=−∇xEθ(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 분포에서 초기 샘플 x0\bold{x}^0을 뽑고, 그 다음에 step size ϵ>0\epsilon >0를 사용하여 KK 단계에 대해 overdamped 랑주뱅 diffusion 프로세스를 시뮬레이션 한다.
xk+1←xk+ϵ22∇xlog⁡pθ(xk)⏟=−∇xEθ(x)+ϵzkk=0,1,...,K−1\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
여기서 zk∼N(0,I)\boldsymbol{z}^k \sim \mathcal{N}(\bold{0},\bold{I})은 가우시안 노이즈 항이다. ϵ→0\epsilon \to 0이고 K→∞K \to \infty 일 때 xK\bold{x}^K는 어떤 regularity 조건 아래 pθ(x)p_{\boldsymbol{\theta}}(\bold{x})로 분포하는 것이 보장된다.
그러나 샘플 x∼pθ(x)\bold{x} \sim p_{\boldsymbol{\theta}}(\bold{x})을 얻기 위해 수렴할 때까지 MCMC를 실행하는 것은 계산적으로 비싸기 때문에 체인이 목표에 수렴할 때까지 실행하지 않고 고정된 수의 단계만 MCMC 절차를 수행하는 대안 방법을 사용할 수 있는데 이것이 바로 Contrastive Divergence(CD)이다.
CD의 TT 단계는 다음 목적을 최소화하는 것과 동등하다.
CDT=DKL(p0∥p∞)−DKL(pT∥p∞)\text{CD}_T = D_\text{KL}(p_0\|p_\infty) - D_\text{KL}(p_T\|p_\infty)
여기서 pTp_T는 TT MCMC 업데이트 이후 x\bold{x}에 대한 분포이고 p0p_0는 데이터 분포이고 p∞p_\infty는 모델이 수렴했을 때에 대한 이론적 분포이다. (일반적으로 작은 TT 값을 사용하여 좋은 결과를 얻을 수 있고 때때로 T=1T=1을 설정하기도 한다.)
DKL(p0∥p∞)D_\text{KL}(p_0\|p_\infty)는 실제 데이터 분포 p0p_0와 모델이 수렴했을 때의 분포 p∞p_\infty 사이의 KL divergence로 모델이 실제 데이터를 얼마나 잘 표현하고 있는지를 나타낸다.
DKL(pT∥p∞)D_\text{KL}(p_T\|p_\infty)는 MCMC의 TT 단계 후에 얻은 분포 pTp_T와 모델이 수렴했을 때의 분포 p∞p_\infty 사이의 KL divergence로 pTp_T가 p∞p_\infty에 얼마나 근접한지를 나타낸다.
CD는 목적은 단순히 모델 분포와 데이터 분포 사이의 차이를 최소화하는게 아니라 두 KL divergence 사이의 차를 최소화하는 것으로, 모델 분포가 데이터 분포 뿐만 아니라 MCMC 수행 후의 분포에도 가까워야 한다는 목적이 된다. 이것은 데이터 분포와 모델 분포 사이의 ‘energy gab’을 최소화하기 위한 전략으로 설계 되었다.

Sample Code

Model

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))
Python
복사

Objective

입력에 대해 Langevin MCMC를 수행한 후에 CD 목적에 맞게 손실을 구하는 함수를 정의한다.
본래 CD의 이론을 따르면 입력과 입력에 대해 MCMC를 수행한 결과에 대해 모델과의 KL divergence를 구하고 그 차이를 손실로 사용해야 하지만, 두 데이터에 대해 동일한 함수를 사용하고 그 차리를 구하면 의미상 KL divergence를 구한 것과 결과가 동일하기 때문에, KL divergence를 직접 구하지 않고 단순히 model에 통과시켜 손실을 구한다. 이 방법이 KL divergence를 계산하는 것보다 더 효과적이다.
def contrastive_divergence(ebm, x, k=1, learning_rate = 0.01, noise = 0.01): original_x = x.clone().detach() x = x.detach().requires_grad_() # 랑주뱅 MCMC를 k번 수행 한다. for _ in range(k): # 입력을 에너지 함수(모델)에 넣어서 에너지 값을 얻고 backward() 통해 기울기를 구한다. 그 기울기는 x의 grad에 저장됨 energy = ebm(x) energy.backward(torch.ones_like(energy), retain_graph=True) # retain_graph=True는 gradient 계산 후에도 그래프를 보존 함. retain_graph=False이면 backward() 후에 graph가 삭제되어 다음번 gradient 계산이 안 된다. # 랑주뱅 다이나믹스를 이용한 업데이트 - 에너지 함수의 gradient에 학습률을 적용하고 노이즈를 더한 값을 기존 데이터에서 뺀다. x.data = x.data - learning_rate * x.grad.data + noise * torch.randn_like(x) x.grad.data.zero_() # 다음 반복에서 gradient가 누적되지 않도록 초기화 fake_data = x.detach() # fake_data가 이후 gradient에서 영향 받지 않도록 detach() # 이론적으로는 D_KL(p_0||model 수렴 분포) D_KL(p_T||model 수렴 분포)를 사용해야 하지만 # KL Divergence 대신 동일한 함수를 사용해도 의미 상 결과가 동일하기 때문에 그냥 model을 돌린다. origin_energy = ebm(original_x) fake_energy = ebm(fake_data) # real과 fake의 차이를 loss로 사용한다 return origin_energy.mean() - fake_energy.mean(), fake_data # fake_data 는 plot 용도
Python
복사

Train

모델 학습은 다음과 같이 수행할 수 있다.
import torch import torch.nn as nn from torch.autograd import grad from torch.optim import Adam from torchvision.datasets import MNIST from torchvision import transforms from torch.utils.data import DataLoader import numpy as np import matplotlib.pyplot as plt # gpu를 사용하는 경우 gpu에서 처리 device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"Using device: {device}") # 데이터 로더 설정 def get_data_loader(batch_size=128): transform = transforms.Compose([ transforms.ToTensor(), # 이미지를 PyTorch 텐서로 변환 transforms.Normalize((0.5,), (0.5,)) # 정규화: [0,1] -> [-1,1] ]) train_dataset = MNIST(root='./data', train=True, transform=transform, download=True) train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True) return train_loader data_loader = get_data_loader() # 모델, 옵티마이저 및 데이터 로더 설정 ebm_cd = EBM().to(device) optim_cd = torch.optim.Adam(ebm_cd.parameters(), lr=0.001) # 학습 루프 epochs = 10 for epoch in range(epochs): for batch_idx, (data, _) in enumerate(data_loader): real_data = data.to(torch.float32).to(device) optim_cd.zero_grad() # 모델을 통해 real과 fake 데이터를 얻는다. loss, fake_data = contrastive_divergence(ebm_cd, real_data) loss.backward() optim_cd.step() if batch_idx % 100 == 0: print(f"Epoch {epoch}, Batch {batch_idx}, Loss {loss.item()}")
Python
복사

참고

Made with