Search
Duplicate

AI/ Noise Contrastive Estimation(NCE)

Noise Contrastive Estimation(NCE)

Noise Contrastive Estimation(NCE)는 실제 데이터 분포를 알려진 밀도를 사용하는 다른 분포와 contrasting 하여 모델을 학습하는 방법으로 정규화 상수를 쉽게 구할 수 있기 때문에 Energy-Based Models(EBMs) 학습에 유용하게 사용될 수 있다.
pD(x)p_\mathcal{D}(\bold{x})를 데이터 분포라 하고 pn(x)p_n(\bold{x})을 알려진 밀도를 사용하는 선택된 분포라 하고 noise 분포라 하자. 이 노이즈 분포는 일반적으로 간단하고 N(0,I)\mathcal{N}(\bold{0},\bold{I})와 같은 다루기 용이한 pdf를 갖기 때문에 pdf를 계산하고 효과적으로 샘플을 생성할 수 있다.
베르누이 분포를 가진 이항 변수 yy를 이용하여 다음과 같이 노이즈와 데이터의 혼합 분포를 정의할 수 있다.
pn,data(x)=p(y=0)pn(x)+p(y=1)pD(x)p_\text{n,data}(\bold{x}) = p(y=0)p_n(\bold{x}) + p(y=1)p_\mathcal{D}(\bold{x})
베이즈 규칙에 따라 이 혼합에서 x\bold{x}가 주어지면 posterior 확률은 다음과 같다.
pn,data(y=0∣x)=pn,data(x∣y=0)p(y=0)pn,data(x)=pn(x)pn(x)+νpD(x)pn,data(y=1∣x)=pn,data(x∣y=1)p(y=1)pn,data(x)=νpD(x)pn(x)+νpD(x)\begin{aligned} p_\text{n,data}(y=0|\bold{x}) &= {p_\text{n,data}(\bold{x}|y=0)p(y=0) \over p_\text{n,data}(\bold{x})} = {p_n(\bold{x}) \over p_n(\bold{x}) + \nu p_\mathcal{D}(\bold{x})} \\ p_\text{n,data}(y=1|\bold{x}) &= {p_\text{n,data}(\bold{x}|y=1)p(y=1) \over p_\text{n,data}(\bold{x})} = {\nu p_\mathcal{D}(\bold{x}) \over p_n(\bold{x}) + \nu p_\mathcal{D}(\bold{x})} \end{aligned}
여기서 ν=p(y=1)/p(y=0)\nu = p(y=1)/p(y=0)
EBM pθ(x)p_{\boldsymbol{\theta}}(\bold{x})를 다음과 같이 정의하자.
pθ(x)=exp⁡(−Eθ(x))/Zθp_{\boldsymbol{\theta}}(\bold{x}) = \exp(-\mathcal{E}_{\boldsymbol{\theta}}(\bold{x}))/Z_{\boldsymbol{\theta}}
NCE에서 ZθZ_{\boldsymbol{\theta}}는 대부분의 다른 EBM과 달리 학습 가능한 (스칼라) 파라미터로 취급된다.
노이즈-데이터의 혼합과 유사하게 노이즈-모델 분포의 혼합을 다음처럼 정의할 수 있다.
pn,θ(x)=p(y=0)pn(x)+p(y=1)pθ(x)p_{n,\boldsymbol{\theta}}(\bold{x}) = p(y=0)p_n(\bold{x}) + p(y=1)p_{\boldsymbol{\theta}}(\bold{x})
이 노이즈-모델 혼합에 대해 y=0y=0의 posterior 확률은 다음과 같다.
pn,θ(y=0∣x)=pn(x)pn(x)+νpθ(x)pn,θ(y=1∣x)=νpθ(x)pn(x)+νpθ(x)\begin{aligned} p_{n,\boldsymbol{\theta}}(y=0|\bold{x}) &= {p_n(\bold{x}) \over p_n(\bold{x})+\nu p_{\boldsymbol{\theta}}(\bold{x})} \\ p_{n,\boldsymbol{\theta}}(y=1|\bold{x}) &= {\nu p_{\boldsymbol{\theta}}(\bold{x}) \over p_n(\bold{x})+\nu p_{\boldsymbol{\theta}}(\bold{x})} \end{aligned}
NCE에서 표준 조건부 maximum likelihood 목적을 통해 pn,θ(y∣x)p_{n,\boldsymbol{\theta}}(y|\bold{x})를 pn,data(y∣x)p_\text{n,data}(y|\bold{x})에 맞추어서 pθ(x)p_{\boldsymbol{\theta}}(\bold{x})를 pD(x)p_\mathcal{D}(\bold{x})에 간접적으로 맞춘다.
θ∗=arg min⁡θEpn,data(x)[DKL(pn,data(y∣x)∥pn,θ(y∣x))]=arg max⁡θEpn,data(x,y)[log⁡pn,θ(y∣x)]\begin{aligned} \boldsymbol{\theta}^* &= \argmin_{\boldsymbol{\theta}} \mathbb{E}_{p_\text{n,data}(\bold{x})}[D_\text{KL}(p_\text{n,data}(y|\bold{x})\| p_{n, \boldsymbol{\theta}}(y|\bold{x}))] \\ &= \argmax_{\boldsymbol{\theta}} \mathbb{E}_{p_\text{n,data}(\bold{x},y)}[\log p_{n,\boldsymbol{\theta}}(y|\bold{x})] \end{aligned}
이 식은 pn,data(x,y)pn,θ(y∣x)p_\text{n,data}(\bold{x},y) p_{n,\boldsymbol{\theta}}(y|\bold{x})의 기대값 형식으로 노이즈-데이터 분포 pn,data(x,y)p_\text{n,data}(\bold{x},y)에서 뽑은 데이터를 노이즈-모델 분포 pn,θ(y∣x)p_{n,\boldsymbol{\theta}}(y|\bold{x})에 넣어 해당 데이터가 실제인지 노이즈인지에 대한 기대값을 최대화하는 것으로 볼 수 있다.
yy가 베르누이 분포를 따르므로 이에 대한 식을 다음과 같이 정의할 수 있다.
log⁡pn,θ(y∣x)=log⁡pn,θ(y=1∣xd)+log⁡pn,θ(y=0∣xn)=log⁡pn,θ(y=1∣xd)+log⁡(1−pn,θ(y=1∣xn))\begin{aligned} \mathbb \log p_{n,\boldsymbol{\theta}}(y|\bold{x}) &= \log p_{n,\boldsymbol{\theta}}(y=1|\bold{x}_d) + \log p_{n,\boldsymbol{\theta}}(y=0|\bold{x}_n) \\ &= \log p_{n,\boldsymbol{\theta}}(y=1|\bold{x}_d) + \log (1-p_{n,\boldsymbol{\theta}}(y=1|\bold{x}_n)) \end{aligned}
여기서 xd\bold{x}_d는 데이터이고 xn\bold{x}_n는 노이즈이다. 모델 pn,θ(y∣x)p_{n,\boldsymbol{\theta}}(y|\bold{x})은 입력이 1일 확률을 예측하도록 학습된다.
이것은 stochastic gradient ascent를 사용하여 해결될 수 있다. 임의의 다른 deep 분류기와 같이 모델이 충분히 강력하면 pn,θ∗(y∣x)p_{n,\boldsymbol{\theta}^*}(y|\bold{x})는 최적에서 pn,data(y∣x)p_\text{n,data}(y|\bold{x})와 일치한다. 이 경우에 다음이 성립한다.
pn,θ∗(y=0∣x)≡pn,data(y=0∣x)⇔pn(x)pn(x)+νpθ∗(x)≡pn(x)pn(x)+νpD(x)⇔pθ∗(x)≡pD(x)\begin{aligned} p_{n,\boldsymbol{\theta}^*}(y=0|\bold{x}) &\equiv p_\text{n,data}(y=0|\bold{x}) \\ \Leftrightarrow {p_n(\bold{x}) \over p_n(\bold{x}) + \nu p_{\boldsymbol{\theta}^*}(\bold{x})} &\equiv {p_n(\bold{x}) \over p_n(\bold{x}) + \nu p_\mathcal{D}(\bold{x})} \\ \Leftrightarrow p_{\boldsymbol{\theta}^*}(\bold{x}) &\equiv p_\mathcal{D}(\bold{x}) \end{aligned}
결론적으로 Eθ∗(x)\mathcal{E}_{\boldsymbol{\theta}^*}(\bold{x})는 데이터 분포 pD(x)p_\mathcal{D}(\bold{x})와 일치하는 비정규화된 에너지 함수이고 Zθ∗Z_{\boldsymbol{\theta}^*}는 해당 정규화 상수이다. NCE는 학습 절차의 부산물로 에너지 기반 모델의 정규화 상수를 제공한다.
EBM이 매우 표현력 있는 경우(예컨대 많은 파라미터를 갖는 심층 신경망) 정규화된 확률 밀도를 근사하여 ZθZ_{\boldsymbol{\theta}}를 Eθ(x)\mathcal{E}_{\boldsymbol{\theta}}(\bold{x})의 파라미터로 흡수하거나 또는 동등하게 Zθ=1Z_{\boldsymbol{\theta}}=1로 고정할 수 있다고 가정할 수 있다. NCE를 사용하여 학습된 EBM 결과는 self-normalized 즉, 1에 가까운 정규화 상수를 갖는다.
NCE를 multi-class로 확장하면 InfoNCE 모델이 된다. 다만 이 경우 softmax 함수를 사용하여 정규화 상수를 계산하지 않으므로 EBM에는 사용할 수 없다. 자세한 내용은 다음 페이지 참조.

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

원래 NCE의 목적 함수는 최대화하는 것이지만, 실제 구현에서는 최소화를 해결하는 것이 더 편리하므로 다음과 같이 기대값에 음수를 취해 최소화 문제로 바꾼다.
arg min⁡θ−E[log⁡pn,θ(y=1∣xd)+log⁡(1−pn,θ(y=1∣xn))]\argmin_{\boldsymbol{\theta}}-\mathbb{E}[\log p_{n,\boldsymbol{\theta}}(y=1|\bold{x}_d) + \log (1-p_{n,\boldsymbol{\theta}}(y=1|\bold{x}_n))]
추가로 NCE에서는 학습 과정에서 noise 확률을 이용하여 정규화 상수 ZZ를 추정할 수 있다. 아래 참조
# 노이즈 데이터 생성 def generate_noise_data(batch_size, noise_dim=784): return torch.rand(batch_size, noise_dim) * 2 - 1 # [-1, 1] 범위의 균등 분포 # NCE Loss 함수 def nce_loss(model, real_data, noise_data, k=1): # 실제와 노이즈 데이터에 대해 각각 값을 구한다. real_scores = model(real_data) noise_scores = model(noise_data) # 분류 문제의 결과는 logit이므로 sigmoid()를 통해 확률로 변환한다. real_prob = torch.sigmoid(real_scores) noise_prob = torch.sigmoid(noise_scores) # real_score는 real_data의 결과이므로 점수가 높을수록 loss에 높게 반영되어야 함 real_loss = torch.log(real_prob + 1e-10) # noise_score는 noise_data의 결과이므로 점수가 낮을수록 loss에 높게 반영되어야 하므로 1 - p로 더함 noise_loss = torch.log(1 - noise_prob + 1e-10) # 원래 목적은 real_score와 noise_score의 합을 최대화 해야 하지만 계산상의 편의를 위해 음을 취해 최소화 문제로 바꾼다. loss = -torch.mean(real_loss + noise_loss) # optional) 정규화 상수 Z의 추정 Z = k * torch.mean(noise_prob / (1 - noise_prob + 1e-10)).detach() return loss, Z
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_nce = EBM().to(device) optim_nce = torch.optim.Adam(ebm_nce.parameters(), lr=0.001) data_loader = get_data_loader() # 학습 루프 epochs = 10 for epoch in range(epochs): for real_data, _ in data_loader: real_data = real_data.view(real_data.size(0), -1).to(device) # 데이터를 1차원으로 변환 noise_data = generate_noise_data(real_data.size(0)).to(device) # 노이즈 데이터 생성 # nce 손실을 구한 후 역전파 loss, estimated_Z = nce_loss(ebm_nce, real_data, noise_data) optim_nce.zero_grad() loss.backward() optim_nce.step() print(f"Epoch {epoch+1}, Loss: {loss.item()}, Estimated Z: {estimated_Z}")
Python
복사

참고

Made with