논문: https://arxiv.org/pdf/1703.04977.pdf
Epistemic uncertainty와 aleatoric uncertainty를 동시에 측정할 수 있게 해주는 방법을 소개한 논문이다. 논문을 간단히 정리하고 PyTorch 코드를 함께 소개하고자 한다.
Uncertainty의 종류
1. Epistemic uncertainty (=Model uncertainty)
모델구조나 학습과정에서 발생하는 uncertainty이다. 모델이 충분히 학습되지 않았을 수도 있고, 전체 데이터 분포를 다 학습하지 못했을 수도 있고, 모델의 구조가 지나치게 단순하거나 복잡할 수도 있다. Epistemic uncertainty는 데이터셋 보강, 모델 구조 수정, 학습 방법 변경 등의 방법으로 줄일 수 있다 (reducible).
2. Aleatoric uncertainty (=Data uncertainty)
데이터에 내재된 uncertainty로, 데이터 수집 과정에서 발생하는 불확실성이라고 볼 수 있다. 데이터에 발생하는 noise나 낮은 resolution 등이 여기에 포함된다. Aleatoric uncertainty는 두 가지로 분류될 수 있다.
a. Homoscedastic uncertainty: 모든 데이터에 동일하게 발생하는 uncertainty. 예를 들면 이미지를 촬영하는 기기에 있는 결함으로 인한 noise나, 낮은 resolution 등이 될 수 있다.
b. Heteroscedastic uncertainty: 각 데이터마다 다르게 발생하는 uncertainty. 예를 들면 사진마다 다른 빛의 노출도나 음영도가 여기에 포함된다.
Aleatoric uncertainty는 데이터 자체에서 오는 uncertainty이기 때문에, epistemic uncertainty와 달리 줄이는 것이 불가능하다 (unreducible).
Measuring epistemic uncertainty (Dropout as a Bayesian)
먼저, epistemic uncertainty를 측정하기 위해 Bayesian neural network (BNN)을 이용한다. BNN은 일반적인 neural network처럼 deterministic 한 weight 값을 사용하는 것이 아니라, weight에 대한 Gaussian prior distribution을 상정한다.
$W \sim \mathcal N (0,I)$
이 때 모델의 random output을 $f^W(x)$라고 하면, 모델의 likelihood를 다음과 같은 Gaussian distribution으로 가정할 수 있다.
$p(y|f^W(x))=\mathcal N(f^W(x),\sigma^2)$
이 likelihood를 직접 구하는 것은 힘들기 때문에, dropout variational inference를 이용한다. 이는 "Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning" (ICML 2016) 에서 제시된 개념으로, dropout을 이용해 Bayesian neural network를 구현하는 것이다. Train과 test 과정 모두에서 dropout을 사용하고, 이를 통해 sampling을 구현한다. KL divergence를 이용해 simple distribution $q_\theta^*(W)$을 true model posterior $p(W|X, Y)$에 근사시키는 방법으로 목적함수를 설정하며, 이를 dropout variational inference라고 한다.
$\mathcal L(\theta,p)=-\frac1N \sum_i^N\log p(y_i|f^{\hat W_i}(x_i)) + \frac{1-p}{2N}||\theta||^2$
이 때 $p$는 dropout ratio이며, $\hat W_i$는 sampling된 weight를 가리킨다. ($\hat W_i \sim q_\theta^*(W)$).
Regression에서는, 앞의 항인 negative log likelihood가 다음과 같이 나타내어질 수 있다.
$-\log p(y_i|f^{\hat W_i}(x_i))\propto\frac1{2\sigma^2}||y_i-f^{\hat W_i}(x_i)||^2+\frac12\log\sigma^2$
이 때 $\sigma$는 observation noise parameter을 나타낸다.
Measuring aleatoric uncertainty
본 논문에서는 aleatoric uncertainty 중 각 데이터마다 그 값이 다른 heteroscedastic uncertainty를 측정한다. 이는 위에서 정의한 observation noise $\sigma$를 data-dependent하게 만듦으로써 ($\sigma(x)$) 학습을 통해 측정할 수 있다.
$\mathcal L_{NN}=\frac1N \sum_i^N\frac{1}{2\sigma(x_i)^2}||y_i-f(x_i)||^2+\frac12\log\sigma(x_i)^2$
즉, 모델은 prediction과 uncertainty 두 output을 return하게 된다. Uncertainty가 높은 경우, 전체 loss 값이 작아지게 되고, 해당 sample의 영향이 작아지게 된다. 따라서, uncertainty가 낮은 경우에 loss에 더 가중치를 두어 학습에 더 영향을 주게 한다고 볼 수 있다. 이를 loss attenuation이라고 하며, noisy data에 대한 robustness를 높여주는 역할을 한다.
Combining two types of uncertainty
앞의 두 방법을 결합해, epistemic uncertainty와 aleatoric uncertainty를 한번에 측정할 수 있다.
$\mathcal L_{BNN}=\frac1D \sum_i^D\frac12\hat\sigma_i^{-2}||y_i-\hat y_i||^2+\frac12\log\hat\sigma^2$
실제로는, $\sigma$ 대신에 log variance $s_i:=\log\hat\sigma^2_i$를 predict하도록 훈련되는데, 그 이유는 loss를 0으로 나누는 것을 피하기 위함이다. 따라서 loss는 다음과 같이 바뀌게 된다.
$\mathcal L_{BNN}=\frac1D \sum_i^D\frac12\exp(-s_i)||y_i-\hat y_i||^2+\frac12s_i$
이 때 Predictive unceratinty $Var(y)$는 epistemic uncertainty와 aleatoric uncertainty의 합으로 나타낼 수 있다.
$Var(y)\simeq\frac1T\sum_t^T\hat y_t^2-(\frac1T\sum_t^T\hat y_t)^2+\frac1T\sum_t^T\hat\sigma_t^2$
- $\frac1T\sum_t^T\hat y_t^2-(\frac1T\sum_t^T\hat y_t)^2$ : epistemic uncertainty
- $\frac1T\sum_t^T\hat\sigma_t^2$ : aleatoric uncertainty
Results
CamVid dataset을 이용한 semantic segmentation에 대해 uncertainty를 측정한 결과이다. (d)를 보면, 물체의 가장자리나, 카메라에서 먼 부분에 높은 값이 나타나는 것을 확인할 수 있는데, 이는 aleatoric uncertainty가 관측값 자체에 있는 uncertainty를 포착한 것임을 보여준다. (e)를 보면, 예측값에서 틀린 부분들이 높은 값을 가지는 것을 확인할 수 있는데, 이는 epistemic uncertainty가 모델의 uncertainty를 포착한 것임을 보여준다.
또한 해당 loss를 이용함으로써 uncertainty의 측정 뿐 아니라 segmentation의 성능도 높일 수 있음을 확인할 수 있다. 이는 loss attenuation으로 인해 noisy data에 영향을 적게 받기 때문이다.
위 실험은 train dataset의 크기를 서로 바꿔가며 모델을 훈련시킨 것이다. Train dataset의 크기가 줄어들수록 epistemic uncertainty는 크게 증가하지만, aleatoric uncertainty는 상대적으로 일정하게 유지됨을 확인할 수 있다. 이는 aleatoric uncertainty는 줄일 수 없고(unreducible), epistemic uncertainty는 데이터셋 보강, 모델 구조 수정, 학습 방법 변경 등의 방법으로 줄일 수 있음(reducible)을 나타낸다.
Code (PyTorch implementation)
https://github.com/hmi88/what을 바탕으로 일부 수정함.
GitHub - hmi88/what: Pytorch implementation of "What Uncertainties Do We Need in Bayesian Deep Learning for Computer Vision?"
Pytorch implementation of "What Uncertainties Do We Need in Bayesian Deep Learning for Computer Vision?" - GitHub - hmi88/what: Pytorch implementation of "What Uncertainties Do We Ne...
github.com
본 글에서는 epistemic은 다루지 않고 aleatoric uncertainty를 계산하는 코드만 설명하고자 한다.
Model
prediction(output_mean)과 uncertainty(output_var)를 계산하는 branch를 따로 둠.
class AleatoricModel(nn.Module):
def __init__(self):
self.encoder = ...
self.decoder_mean = ...
self.decoder_var = ...
def forward(self, x):
enc = self.encoder(x)
output_mean = self.decoder_mean(enc)
output_var = self.decoder_var(enc)
return output_mean, output_var
Loss
Model에서 얻은 두 개의 output(mean, var)을 이용해 loss를 계산한다. 이때 논문과 다른 점은 var에 따로 설정할 수 있는 var_weight를 준다.
import torch
import torch.nn as nn
class AleatoricLoss(nn.Module):
def __init__(self, var_weight):
super(AleatoricLoss, self).__init__()
self.var_weight = var_weight
def forward(self, pred, label):
mean, var = pred
var = self.var_weight * var
mse = (mean-label) ** 2
loss = torch.mul(torch.exp(-var), mse) + var
return 0.5 * loss.mean()
참고자료