Distill-Whisper: Robust Knowledge Distillation via Large-Scale Pseudo Labelling

References

https://arxiv.org/pdf/2311.00430

 

 

Introduction

Whisper를 아는가?

Speech-to-Text(STT) 분야에서 오늘날 가장 많이 사용되고 있는 모델이 OpenAI에서 2022년 개발한 Whisper라는 모델이다.

Whisper는 seq2seq기반 구조를 택한 모델로 강력한 text decoder를 가지고 있고, 680,000시간 분량의 pretrain을 거쳐 아직까지도 이 분야에서 최고 성능을 내는 모델로 남아 있다.

 

그러나 이로 인해 모델이 너무 무겁다는 단점을 동시에 갖고 있다. 온디바이스 적용과 같은 실용적 접근을 위해선 파라미터 수 절감 및 연산 속도 개선 등 경량화 작업이 필수적인데, 동시에 경량화로 인해 기존 모델의 성능이 너무 떨어지지 않게 방법론을 잘 설계하는 것이 중요하다.

 

2023년, Hugging Face는 이 whisper의 파라미터 수를 줄이고 연산 속도를 대폭 올리면서 Word Error Rate는 거의 떨어뜨리지 않는 Distill-Whisper라는 방법론을 제시하였다. 이 방법론은 지식 증류 기반 경량화를 수행하나, 증류 방법론 자체에 큰 변동은 주지 않고 데이터 관점에서의 개선을 주로 타겟팅한다.

유념해야 할 점은, 이 논문은 지식 증류 자체보다는 Whisper의 경량화에 더 포커싱이 맞추어져 있다는 점이다.

 

주요 방법론을 요약하자면

• Teacher-pseudo labeled data 기반 지식 증류
• Pseudo-labeling 퀄리티 보존을 위한 WER 기반 휴리스틱 데이터 필터링

두 가지로 볼 수 있겠다.

 

아래 섹션에서 이 두 가지 방법론을 간략히 설명하고자 한다.

 

Teacher-pseudo labeled data 기반 지식 증류

사전에 먼저 짚고 넘어가자면,

본 논문에서 사용한 지식 증류 방법론은 초기 Geoffery Hinton이 제시한 바와 같이, Logit의 KL-Divergence를 바탕으로 수행하는 전통적 방법이다.

 

다만 지식 증류 과정의 효율성 증대를 위하여 앞서 이야기한 두 가지 방법론을 사용하는 것인데, 먼저 설명할 내용은 pseudo-labeling을 통한 지식 증류 방법론이다.

 

원래는 지식 증류로 학습할 경우 task loss(ex. Cross-entropy), KD loss(ex. KL-Divergence) 두 가지 loss를 사용한다.

그리고 ASR에서의 KD loss는 각 단어 단위로 예측 확률 분포가 teacher와 유사해지도록 유도하는 것을 목표로 한다.

 

그러나 본 방법론은, 단순 단어 단위 뿐 아니라 문장 단위에서의 예측 분포 또한 유사해지도록 유도하고자 하였고, 이를 위해 pseudo label을 사용한 것이다.

 

이 개념을 지식 증류에 활용된 방법은 단순하다. Teacher가 샘플에 대한 transcript(문장)을 예측하고(=pseudo-label 생성), student가 이 문장을 따라가도록 별도의 task loss를 주는 것이다.

Teacher도 모든 문장을 항상 정확하게 인식하는 것은 아니므로 틀릴 때의 패턴 & 맞을 때의 패턴 등이 정해져 있을 텐데, 이 loss를 student에게 줌으로써 teacher의 불완전한(?) 예측 패턴을 따라가면서 원래 teacher 분포에 더 가까워지는 것이다.

 

논문에선 이러한 방법론을 sequence-level KD, 기존 방법론을 word-level KD라고 명명한다.

그리고 이 방법론을 사용하는 대신 Ground-Truth 기반 task loss는 아예 적용하지 않는다.

 

즉 학습 시의 Loss는,

Total Loss = KD Loss(a*PL loss + b*KL loss)
* PL loss: sequence-level KD task에서의 loss

로 정리할 수 있다.

참고로, 논문에서는 PL loss의 가중치인 a에 1을, KL loss의 가중치인 b에 0.8을 부여하였다.

 

Pseudo-labeling 퀄리티 보존을 위한 WER 기반 휴리스틱 데이터 필터링

주의해야 할 점은, teacher 분포를 잘 따라가게 해보겠다고 pseudo labeling을 활용하였지만,

앞서 언급한 것처럼 항상 teacher가 모든 답을 바람직하게 맞추진 않는다는 것이다.

때에 따라 할루시네이션이 발생할 수도 있고, 너무 정답과 멀리 떨어진 예측을 내놓는데 그걸 student가 따라가게 유도하는 것은 비효율적이다.

 

그렇기 때문에 본 논문에서는 너무 퀄리티가 낮은 pseudo label은 버리는 방법을 취하였다.

이 방법도 상당히 단순하다.

• 임의의 샘플이 들어오면, Whisper 학습에 사용된 Text-Normalizer를 활용하여 Text Normalization을 수행한다.
  (Text Normalization은 '백 점' != '100점'과 같이, 실제로는 동일한 의미이나 표기가 달라 예측이 오답 처리되는 것을 막기 위해 일관적으로 non-speech 표기를 speech 표기로 바꾸는 작업이다. 즉 100 → '백' 이 되는 것)
• 샘플에 대한 Normalize된 GT와 teacher 예측 간의 WER을 측정한다. 이때 WER이 사전 정의 휴리스틱 값인 lambda보다 크면, 그 샘플은 PL loss 계산에 사용하지 않는다.

논문에선 lambda=10일 때가 전체 학습에서 가장 좋은 성능을 내었다고 보고한다.

이와 같은 방법을 통하여 PL loss가 모델을 성능이 악화되는 방향으로 유도하지 않도록 필터링을 수행하였다.

 

Student와 Teacher는 어떻게 구성하였나?

논문에서는 distill-large-v2, distill-medium라는 두 가지의 모델을 제시하였다.

두 모델은 각각 Whisper-large, Whisper-medium을 teacher로 두는데,

 

teacher와의 유일한 차이는 decoder 개수가 원래 개수에서 단 2개로만 바뀌었다는 점이다.

심지어 encoder는 freeze시켜 두고 아예 학습조차 하지 않는다.

 

원래 Whisper 사이즈별 스펙과 제안된 모델의 스펙을 참고 바란다.

Original Whisper

vs

Distill-Whisper

 

 

넘어가기 전에 드는 의문,

왜 인코더는 그대로 두고 디코더 사이즈만 줄였을까?

 

Appendix에서 이와 관련된 내용을 서술하였는데, 요약하자면

1) 인코더를 줄이는 실험을 해 보았을 때 성능 손실이 너무 커지는 문제가 있었고

2) Whisper 추론에서 가장 속도 병목인 부분이 autoregressive 연산을 수행하는 디코더 파트이기 때문에(Transformer 구조상),

 

인코더보다는 디코더를 경량화하는 데에 더 중점을 두었다고 한다.

 

지금까지의 파이프라인을 요약하면 아래 그림처럼 이해해볼 수 있겠다.

 

 

참고로 학습에 활용한 방법론 중 Chunked Long-Form Transcription, Speculate Decoding 과 같은 추가 내용도 존재하는데,

이들은 distill 관점에서는 벗어나는 부수적인 학습 파이프라인 요소들이라 제외하였다. 필요할 경우 별도로 읽어보길 권한다.

 

Experiments & Results

우선 학습을 위해 10가지 도메인의 데이터셋을 활용하였다.

 

그리고 평가는 short-form / long-form 에서 각각 수행하고, 이때 데이터 도메인이 겹치지 않게 아래 각 4개의 데이터셋을 섞어 평가 데이터셋을 구축하였다.

 

이때의 결과는 아래와 같다.

• distill-medium.en이 medium.en(teacher)보다 2배 가량 작은 사이즈인데,
  - Long-form에서 성능은 증류 이전과 거의 동일하고,
  - Short-form에서의 성능도 small.en수준 으로 선방한 것을 확인할 수 있다.
  주목할 점은 이 모델의 속도는 기존의 tiny보다 더 빨라졌다는 것이다.
• distill-large-v2는 medium.en 정도 사이즈에 large 수준의 성능(Long-form 기준)을 보이고, short-form에서도 small 정도의 성능을 낸다. 속도는 기존 tiny보다 약간 느린 수준이다.

 

 

Conclusion

본 논문에서는 distillation 방법 자체를 크게 건드리지 않으면서도, pseudo-labeling을 활용하여 증류 효율을 크게 향상시켰다.

증류를 '어떻게'가 아닌 '무엇을' 이라는 관점에서 접근한다는 점이 신선했다.

 

다만 두 가지의 아쉬운 점이 존재했다.

1) KD의 Encoder 영향 여부 불투명. Whisper는 Encoder → Decoder의 흐름으로 학습하고, KD loss는 decoder 끝단에서 매겨지기 때문에 원래는 backpropagation 과정에서 encoder도 KD loss의 영향을 받게 된다. 그러나 저자들은 encoder를 freeze한 상태에서 학습한 결과만을 보여주었고, 이로 인해 이 방법론이 인코더에 어떠한 영향을 미치는지는 불투명하다. Appendix에서 따로 보여줬어도 좋았을 건데, 왜 이 파트는 뺐는지 의문이다.

2) 동일한 Width 내 시나리오 한계. 일반적으로 소형 모델을 만들기 위한 지식 증류를 수행할 때에는 teacher-student 간 width가 달라지는 경우가 많은데, 이 논문에서는 실험의 범위가 다소 제한적이었다고 생각한다. 어쩌면 width 감소로 인해 발생하는 (속도)이득이 이 방법론과 독립적으로 작용하기 때문이라고 판단하여 보여주지 않은 것일 수도 있겠다 싶으나, 개인적으로는 더 작은 모델에서의 성능 추이도 확인하고 싶었다. 뭐 꼭 확인하고자 한다면 직접 실험하는 수밖에..

 

몇 안되는 도메인 특화 지식증류 논문이라 반가웠다.

기회가 된다면 이같은 STT 특화 경량화 논문을 더 찾아서 정리해보면 좋겠다.