Kim, H. et al. (2021) Multi-domain Knowledge Distillation via Uncertainty-Matching for End-to-End ASR Models 논문 리뷰입니다.
Multi-domain Knowledge Distillation via Uncertainty-Matching for End-to-End ASR Models
( End-to-End ASR 모델을 위한 불확실성 매칭을 통한 다중 도메인 지식 증류)
논문에 대해 설명드리기 전에, 먼저 본 논문에서 사용한 방법론에서 핵심인 두 가지 개념에 대해 짚고 넘어가려고 합니다.
본 논문에서는 uncertainty와 knowledge distillation 을 다루고 있는데, 먼저 uncertainty에 대해 간단히 소개해드리겠습니다.
딥러닝에서 uncertainty가 중요해진 두 가지 예시가 있는데요.
딥러닝 기술들이 실생활로 넘어오면서 uncertainty를 구하는 것과, 그 uncertainty를 고려할 수 있는 방법론이 점점 중요해지고 있습니다.
첫번째 사진은 테슬라의 차 사고 사진인데요. 자동주행 중 오른쪽의 동그라미 부분을 하늘로 인식하여 돌진하다가 사고가 난 경우입니다. 두번째 사진은 구글 포토에서 흑인을 고릴라로 판단하여 문제가 생긴 경우입니다.
이렇게 딥러닝이 일상생활에서 점점 많이 쓰이게 되면서 uncertainty 를 구하고, 잘 모르는 판단에는 모른다라고 말할 수 있는 uncertainty 가 점점 중요해지고 있습니다.
딥러닝에서 uncertainty 는 3가지로 분류할 수 있습니다.
1. 한 번도 본 적이 없는 공간에 데이터가 들어왔을 때, 즉 테스트 데이터가 내가 한 번도 본적이 없는 데이터인 경우는, out of distribution test data 라고 합니다.
2. 두번째로는 트레인 데이터 자체에 노이즈가 껴있을 때 혹은 그 데이터가 보여지는 과정 자체에 노이즈가 있을 때를 aleatoric uncertainty 라고 합니다. 예를 들어 아이의 사진을 찍을 때, 아이가 사진을 찍을 때마다 찡그린 표정을 짓는다던가 이런식으로 데이터 자체를 관측하는 과정에서 노이즈가 있을 때 aleatoric uncertainty 라고 할 수 있습니다.
3. 세번째로는 내 모델이 과연 training data를 정말 잘 피팅했을까에 대한 불확실성을 epistemic uncertainty 라고 합니다.
Uncertainty 에 대해서 본 논문에서도 많이 언급한 중요한 논문이 있는데, 바로 이 what uncertainties do we need in Bayesian deep learning for computer vision? 이라는 논문입니다. 이 논문에서는 aleatoric 불확실성을 관측에 있는 노이즈를 측정할 수 있는 것이라고 정의했고, epistemic 불확실성을 충분한 데이터가 주어지면 설명될 수 있는 불확실성, 모델 자체가 얼마나 실릴 수 있을지에 대한 uncertainty 라고 정의했습니다.
그리고 본 논문에서 다루는 두번째 중요한 개념은 Knowledge Distillation 이라는 ‘지식 증류’ 기법입니다. 딥러닝에서의 지식 증류는 학습된 모델로부터 지식을 추출하는 과정입니다. 큰 모델인 teacher network 로부터 증류한 지식을 작은 모델인 student network 로 transfer 하는 일련의 과정입니다.
지식 증류는 모델을 배포하는 과정에서 등장한 개념인데, 모델을 실제 서비스로 배포하는 경우에 복잡한 모델보다는 단순한 모델이 더 적합할 때가 있습니다. 따라서 복잡한 모델이 학습한 일반화 능력을 단순한 모델로 전달해주는 것을 지식증류라고 할 수 있습니다.
Distilling the Knowledge in a Neural Network 의 논문에서 제프리 힌튼이 KD 의 개념에 대해 처음으로 제시했는데, 정답이라고 생각하는 값만 중요한 게 아니라, 어느 클래스와 어느 클래스가 더 비슷하다고 생각했는지와 같은 것도 모두 정보의 일종이라는 것입니다.
실제 학습 과정에서는, 먼저 큰 모델을 학습 시키고, 그 후 작은 모델을 학습시키면서 두 개의 loss function을 구하고, 그 합으로 최종 loss 를 만듭니다. 첫번째 loss인 distillation loss 는 student network 의 soft prediction 과 teacher network 의 soft label 간의 loss 이고, 두번째 loss 는 student network 의 hard prediction 과 실제 ground truth 간의 loss 로 classification loss 입니다.
T 가 1 일 경우, 우리가 알고 있는 softmax function 으로 가장 큰 값을 제외하고 다 버리고 그 값과 ground truth 를 비교해서 loss 를 계산하고, T가 1보다 큰 어떤 상수가 될 경우, p 가 soften 되어 ‘덜 확신하는 ‘ 값인 soft predictions 이 되고, 이 것과 teacher 모델로부터 나온 soft labels 간의 loss 를 구하게 되는 것입니다.
즉, teacher 모델 추론과 student 모델 추론의 차이인 distillation loss 와 실제 레이블 데이터와의 loss 인 student loss 를 더해서 최종 loss 를 산출하게 됩니다.
Distillation loss 는 teacher network 가 알고 있는 지식을 흉내냄으로써 학습하고, student loss 는 right answer 인 정답을 학습합니다.
이 때 distillation loss 를 구할 때 사용하는 “soft label” 이라는 개념은 원핫벡터인 hard label 과는 반대로 모델의 knowledge 를 나타내주기 위해 softmax 를 통과시킨 확률값을 사용합니다. 추가적으로 T 라는 temperature 값을 사용해 T 로 나누어 softmax 를 취하게 해서 output value 를 더 극단적이지 않게, smooth 하게 만들어 줄 수 있습니다.
여기서 큰 모델의 지식을 이용한다는 점에서 transfer learning 과 비슷하다고 생각되어 다른 점에 대해 알아보았습니다.
Transfer learning 과 knowledge distillation 모두 지식을 전달한다는 점에서는 둘 모두 동일하지만, 지식 증류는 처음에 모델 배포, 압축 단계에서 나왔던 개념인만큼 사이즈를 줄이는데에 더 목적이 있습니다. 또 transfer learning 은 오토바이의 바퀴를 찾고 싶을 때 자동차 바퀴를 구별해내는 모델을 이용하는 것처럼 도메인이 다르다면, KD 는 대형 모델이 학습한 정보를 더 작은 별개의 소형 모델로 전달해주는 방식으로 도메인이 같은 task 에서 적용하는 기법입니다.
Transfer learning 은 학습된 모델에서 fully connected layer 과 softmax 분류 출력 부분만 제거하고, 새로운 fully connected layers 와 내가 원하는 새로운 출력층을 넣어, 마지막 파트만 바뀐 모델을 다시 학습시키는 것입니다. 이 때 convolution layers 의 weight는 기존에 학습된 모든 정보들을 가지고 있기 때문에 freeze 시켜 사용함으로써 pre-trained task 에서 내가 원하는 new task 로 knowledge를 transfer 하는 과정이라고 할 수 있습니다.
Knowledge Distillation 은 teacher model 이 학습한 정보를 더 작은 별개의 student model 이 따라하는 방식으로 학습되면서 teacher model 의 지식이 student model 로 전달되는 과정이라고 볼 수 있습니다.
이제 본 논문에 대해 설명드리겠습니다.
먼저 기존 convolutional neural network 모델의 단점은 오버피팅되기 쉽다는 것과, 빈도주의적 관점에서 학습되었기 때문에 모르는 부분이 어떤 것인지 알아내기 어렵고, 또 모르는 것에 있어서 모른다고 대답하지 못한다는 것입니다. 즉 uncertainty 를 측정할 수 없는 문제가 있습니다.
하지만 베이지안 관점으로 딥러닝 모델의 파라미터를 학습하면, 기존 CNN의 deterministic한 weight 가 아닌, 각 weight 의 분포를 학습하기 때문에 모델의 판단이 얼마나 uncertain 한지 측정할 수 있습니다.
이 전에는 단순히 예측 결과만 보여줌으로써 신뢰도에 대한 정보를 ‘암시적’으로 보여줬다면, 본 논문에서는 uncertainty 를 실제로 계산함으로써 uncertainty 에 대한 개념을 model 의 신뢰도로 판단할 수 있고, 또 단순히 예측만 하는 모델이 아닌, uncertainty 를 함께 제공함으로써 모델이 어느 정도의 신뢰도를 가지는지 명시적으로 제시가 가능하다는 것입니다.
특히 본 논문은 speaker variations 이나 헷갈리는 발음과 같은 데이터 관찰에 내재된 불확실성을 직접 모델링하고 전송하는 자동음성인식을 위한 새로운 지식 증류 방법을 제안합니다. 또 추가적으로 다양한 영역에서 학습한 여러 teacher 를 사용하여 지식을 보다 효과적으로 전달하는 효과를 보여줍니다.
또 WSJ 데이터셋으로 테스트했을 때, 본 논문에서 제안하는 지식 증류 방법은 student baseline performance 에 비해 상당한 개선을 보여줍니다.
먼저 앞에서 설명드린 것처럼, 지식 증류는 더 많은 데이터로 훈련된 교사 네트워크의 지식을 학생 네트워크에 전달함으로써 제한된 용량이나 데이터를 갖는 대상 네트워크의 성능을 향상시키는 것을 목표로 합니다. 또 제한된 자원 조건 하에서도 성능 저하를 최소화하기 위해서 제안된 것으로, 최근 몇 년동안은 소프트맥스 계층에 대한 정보를 넘어, 중간 계층 또는 어텐션 맵의 지식도 matching 하는데 사용되고 있습니다. 또한, 지식 증류는 컴퓨터 비전에서도 많이 쓰이고, ASR 자동 음성인식 등 다양한 분야에 적용할 수 있는 범용 압축 기술이라고 할 수 있습니다.
KD의 핵심 원리는 교사 네트워크에서 생성된 소프트맥스 확률이 출력 단일 지점보다 훨씬 더 유용한 정보를 포함하기 때문에, 예측 분포를 매칭시킴으로써 교사 네트워크의 지식을 암시적으로 활용할 수 있습니다.
이제 본 논문에서는 이 지식 증류에 “불확실성” 개념을 하나 더 추가해 교사 네트워크의 불확실성에 대한 정보도 활용하여 증류하는 방법론을 제안합니다.
이 불확실성 정보를 활용하기 위해 먼저 uncertainty 를 측정할 수 있어야 하는데, 불확실성을 측정하는 작업은 매우 어렵고 많은 연구가 수행되었었습니다. 그 중 이 분야의 중요한 연구인 “Dropout as a Bayesian approximation : representing model uncertainty in deep learning” 논문은 드롭아웃 정규화된 딥 뉴럴 네트워크는 가우시안의 posterior 에 대한 변형 근사치이고, 드롭아웃 정규화된 네트워크는 uncertainty 를 측정할 수 있다는 것을 보여주었습니다.
본 논문에서는 프라이버시와 같은 문제로 인해 학생 네트워크가 거의 데이터를 사용하지 않고 학습하도록 제한되는 지식 증류 방법을 연구하고, 특히 heteroscedastic, aleatoric 유형인 토큰별 uncertainty 를 모델링하여 정보를 전송합니다.
여기에 추가적으로 동적 가중치를 적용하여 여러 도메인에서 여러 교사 네트워크의 지식을 증류하는 실험을 하고, 결과적으로 교사 모델이 LibriSpeech 및 Tedium2 로 훈련될 때, 본 논문이 제안하는 방법론이 WSJ 의 student baseline 에 비해 크게 향상된다는 것을 보여줍니다.
댓글