앞 게시글에 이어서 Kim, H. et al. (2021) Multi-domain Knowledge Distillation via Uncertainty-Matching for End-to-End ASR Models 논문 리뷰를 이어가겠습니다.
본 논문이 제안하는 방법론인 KD와 uncertainty 중 지식 증류에 대해 먼저 살펴보겠습니다.
E2E ASR 에 지식 증류를 적용하기 위해 기존 KD 처럼 두가지 loss function 을 사용합니다.
먼저 첫번째 loss 를 살펴보겠습니다.
먼저 E2E ASR은 시퀀스 라벨링 작업으로서, 입력 토큰 s 를 가지고 negative log likelihood 를 최소화하여 출력 t 를 예측하는데, 여기서 yi 는 대상 텍스트의 i번째 ground truth 이며, p 에서는 이전 토큰의 예측에 대한 의존성을 억제시켜줍니다. C는 출력어휘를 나타내고 이 c는 클래스 인덱스입니다.
델타는 kronecker 델타로 yi 와 c 가 같으면 함숫값이 1이고, 아니면 0을 나타냅니다. 이 loss 가 앞 게시글의 지식 증류 그림에서 아래에 해당하는 일반 student loss 에 해당하는 부분입니다.
이 Loss function 의 목표는 타겟과 출력 분포 사이의 교차 엔트로피를 최소화하는 것입니다.
두번째 loss function 은 E2E ASR 모델 압축을 위한 토큰 레벨의 KD 입니다. 여기서 p hat i 는 위의 NLL 손실의 델타 대신에, teacher 모델의 pretrained 파라미터인 세타T 에 의한 teacher 모델의 I 번째 토큰 출력 분포이고, 이 loss function 은 이 그림에서 distillation loss 라고 할 수 있습니다. 이 loss function 을 최소화하는 것이 teacher 과 student model 확률 분포 사이의 KL 분산을 최소화하는 것으로 볼 수 있습니다.
이제 uncertainty 개념을 지식 증류에 적용하는 과정입니다.
본 논문에서 많이 언급되는 What uncertainties do we need in Bayesian deep learning for computer vision? 논문에서는 MC-dropout 및 input-dependent variance modeling 을 사용하여 모델과 데이터에서 uncertainty, 즉 episdemic 과 aleatoric uncertainty 구하는 것을 제안하고, 이를 컴퓨터비전 semantic segmentation, depth estimation 에 적용하였습니다.
이 논문을 많이 참고하여 본 논문에서는 sequence-to-sequence 모델에 uncertainty 를 활용한 KD 를 적용했습니다.
야린 갈의 이 논문에서 MC-dropout 을 사용해서 uncertainty 를 구할 수 있었던 배경에는, 야린 갈이 그 전에 냈던 dropout as a Bayesian approximation 논문을 참고하면 좋습니다.
베이지안 뉴럴 네트워크는 이전의 CNN 처럼 모델의 weights 를 확정적인 숫자로 구하는 것이 아니라 확률 분포로 학습하게 되는데, 이 과정을 objective function 을 구성하여 수행할 수도 있지만, 더 쉬운 방법으로 dropout 을 사용하는 것만으로도 이 효과를 얻을 수 있다고 이 전 논문에서 말하고 있습니다.
Dropout 이 확률적으로 weight 를 사용해서 test 때마다 다른 결과를 얻을 수 있기 때문에 똑같은 효과를 보일 수 있는 것입니다.
딥러닝 모델은 굉장히 고차원의 파라미터를 가지고 있어서 posterior 분포를 구하기가 동전던지기 사례처럼 쉽지 않기 때문에 ‘근사’ 하는 방식을 사용해야 하고, 주로 사용되는 방법이 variational inference 입니다. Posterior 분포를 근사하기 위해 variational 분포를 가정하고, 이 둘 사이의 거리를 좁히는 것을 variational 분포의 파라미터를 찾는 최적화 문제로 변환하여 stochastic 방식으로 구할 수 있게 됩니다. 이게 사실 복잡한 것인데 최근 연구들에서 이 방식이 CNN 에 MC Dropout 을 사용하는 것과 근사한다는 것이 증명되어서, MC 드롭아웃을 통해 베이지안 네트워크를 근사적으로 구현할 수 있게 되었습니다.
What uncertainties do we need in Bayesian deep learning for computer vision? 논문에 대해 조금 더 알아봅시다.
이 논문은 본 논문에서도 많이 언급된 것처럼 uncertainty 에서 아주 중요한 논문입니다.
이 논문의 main contribution 은,
Epistemic uncertainty, aleatoric uncertainty 의 두 가지 uncertainty 값을 직접 구할 수 있다 라는 것과 또 실제 성능이 좋다 라는 것입니다.
output 의 값 자체를 그냥 찾는 것이 아니라 output 의 분포를 모델링할 수 있으면, 그걸 통해서 aleatoric uncertainty 를 모델링할 수 있고, 또 variance를 decompose 하여 aleatoric uncertainty 와 epistemic uncertainty 로 나눌 수 있다 라는게 중요한 포인트입니다.
aleatoric uncertainty와 epistemic uncertainty 은 그럼 어떻게 모델링할까요?
aleatoric uncertainty 는 output 의 분포를 모델링함으로써 모델링할 수 있고, epistemic uncertainty는 우리가 모델에 대해 prior 을 놓고, 데이터가 주어졌을 때 prior 이 얼마나 달라지는지 보는 것. 마치 베이지안 뉴럴 넷에서 보는 variance term 을 모델링하는 것과 같습니다.
다시 reference 논문 보충 설명으로 돌아간다면, 먼저 이 논문은 semantic segmentation 과 depth estimation 의 두 가지 실험을 했는데, 그 중 semantic segmentation 의 실험 결과를 보겠습니다.
(d)와 (e) 처럼 실제로 uncertainty 를 직접 계산해서 나타냈을 때, (c) segmentation model 이 정확히 구분하지 못했던 부분에서 uncertainty 의 값이 크다는 것을 확인할 수 있습니다.
(d) 의 aleatoric uncertainty 를 보면 aleatoric 불확실성은 measurement noise 를 의미하는데, (c)의 semantic segmentation 에서 가장자리 부분에 aleatoric uncertainty 가 많은 것을 볼 수 있습니다. 그 이유는, 실제로 input image가 주어졌을 때 이걸 사람이 다 가장자리를 따서 labeling 을 진행하는 것이라서 semantic labeling 과정에서 생길 수 있는 inconsistency로 인해 aleatoric 불확실성이 나타날 수 있고, 따라서 d 를 확인했을 때 각 object 들의 경계면 부분에서 aleatoric 불확실성이 높게 나타나는 걸 볼 수 있었습니다. 또한 e 의 경우에는 training example 에서 볼 수 없었던 object 가 새로운 image 에서 출현한 경우 epistemic 불확실성이 높게 나타났습니다.
이렇게, 이 논문에서는 d 와 e 처럼 실제로 이 두 가지 불확실성을 구했다는 점과, 또 classification 의 성능이 좋다는 것이 point 입니다.
몬테카를로 드롭아웃을 이용해서 얻어지는 네트워크는 모델 자체의 classification accuracy 는 떨어질 수 있다고 봤다면 이 논문의 classification 성능은 좋습니다.
그렇다면 어떻게 성능이 좋을 수 있었을지에 대해 알아보겠습니다.
해당 그림의 density net 학습과정에서 loss function 을 minimize 하는데 weight 를 그냥이 아니라 드롭아웃을 걸어서 뽑는다는 점이 중요합니다.
일반적으로 input 이 있고 y 를 찾는 regression 일 때, y 에 대한 분포 자체가 가우시안이면 loss function 도 가우시안 likelihood 를 maximize 하면 됩니다.
그런데 regression 이 아닌 분류 문제에서는 어떤 입력에 대해서 output 들. Logit 들에 대한 확률 분포가 나왔을 때 logit 에 대한 확률 분포가 가우시안을 따르는 어떤 값이 아닌, 이것을 softmax 취해서 분류를 푸는게 이 문제의 목적입니다.
따라서 y 를 softmax 취해서 나오는 그 loss 의 cross entropy 를 키워야합니다. 이렇게 logit 에 대한 확률 분포가 나왔을 때 이 것을 한번 더 가공해서 새롭게 loss 를 취하게 되니까 바로 negative log likelihood 를 쓸 수 없습니다.
Regression 처럼 입력이 x 고 output 이 y 일 때 y를 가우시안으로 모델링하면 x 에서 y 로 가는 뮤와 시그마 함수를 바로 찾을 수 있습니다. 그런데 나온 output y 가 분류를 푸는게 목적이면 이런식으로 이 softmax 취한 값의 loss 의 cross entropy 를 키워야합니다.
뉴럴넷 아웃풋 y 의 mean 과 variance 를 가지고, 그걸 sampling 해서 얻어진 값을 softmax 에 집어넣어서 그 때의 cross entropy loss 를 backpropagation 시키는 방식이고, 그렇게 한 이후 output의 분포인 uncertainty 를 구할 수 있습니다.
결과적으로 이 식 자체가 뉴럴넷 아웃풋 y 에 대한 variance 인데,
이 뉴럴넷 아웃풋 y 에 대한 total variance 식에서, 아웃풋에서 얻어지는 시그마가 measurement noise 에 해당하고 그게 aleatoric uncertainty 입니다. 그리고 그 나머지는 epistemic uncertainty 가 되는 것입니다. y hat 은 주어진 입력에 대한 mean 들이고, 시그마 hat 은 주어진 입력에 대한 variance 들로 다 샘플된 값들이다. 네트워크를 돌릴 때마다 weight 에 드롭아웃을 걸기 때문에 매번 다른 y hat, 시그마 hat, T 가 나오고, 그렇게 해서 이런 식이 나올 수 있는 것이다.
위의 식처럼 이런식으로 total variance 를 모델링할 수 있고, epistemic 과 aleatoric uncertainty 를 구할 수 있다라는 것이 이 논문의 포인트였습니다.
또 이 분류기가 성능이 좋을 수 있었던 이유는, 높은 불확실성을 가진다고 생각되는 부분에 영향력을 줄여버릴 수 있었기 때문인데, heteroscedastic uncertainty 는 각 위치마다 불확실성이 달라지는 것을 말합니다. 지금 mean 과 variance 를 모델링한 상황에서 입력이 달라지면 위치가 달라지는 variance 는, 입력에 따라서 noise level 혹은 variance 를 다르게 학습할 수 있습니다. 어떤 라벨 아웃풋이 있을 때 그게 문제가 있으면 그곳에 대한 variance 를 그냥 키움으로써 그 영향력을 줄일 수 있다는 것입니다. 애초에 학습 데이터 자체에 문제가 있으면 그곳에 대한 영향력을 줄여버림으로써 전반적인 나머지를 잘 맞출 수 있습니다. 실제로 semantic segmentation task 에 대해서는 이 방법론이 state of the art 를 보였습니다.
댓글