May 8, 2022
욕설, 혐오 분류기에 지식 증류(Knowledge Distillation) 적용해보기
얼마전 kcbert-base를 fine-tuning해서 욕설분류기를 만들었습니다. 문제는 bert-base의 400MB에 달하는 모델 크기와 그에 따른 추론속도였습니다. 그래서 지식 증류(Knowledge Distillation)을 이용해서 모델 크기를 확 줄여보고 그냥 작은 모델을 썼을 때보다 성능도 더 뛰어나게 만들어보았습니다.
전체 학습 코드는 아래 저장소에 있습니다.
지식 증류란?
지식 증류란 성능이 뛰어난 대형 모델이나 앙상블 모델을 더 작은 모델이 모방하여 학습하는 방법입니다. 이를 통해 작고 빠른 모델의 성능을 끌어올릴 수 있습니다.
Teacher-Student 모델
지식증류에는 여러 기법들이 있지만 가장 기본적인 Teacher-Student 방식으로 학습해보았습니다. 먼저 큰 모델을 이용해서 성능이 좋은 Teacher를학습하고, 이 Teacher를 이용해서 학습 데이터를 예측한 결과를 얻습니다. 이를 Soft label이라고 합니다. 학습 데이터셋처럼 1, 0으로 나뉜 Hard Label이 아닌 Soft label은, 0.1, 0.9와 같은 예상 확률을 이용한다면, Teacher 모델이 학습한 지식을 Student 모델이 따라 배울 수 있는 것입니다.
예를 들어 승용차는 트럭과 유사한 외형을 갖고 있기 때문에 분류기가 승용차로 분류할 때, 의해 트럭에도 높은 확률을 부여하게 됩니다. Student는 soft label 학습을 통해 이것은 승용차이고 이것은 트럭이다 라는 정보 외에도, 승용차와 트럭이 비슷하다는 지식을 추가로 학습하게 되는 것입니다. 그리고 이러한 정보가 Student의 성능 향상을 이끌어냅니다.
욕설 분류, 혐오 분류와 같은 이진분류에서도 지식증류가 성능을 향상시킬 수 있을까 궁금했는데, 실제로 성능이 향상되었습니다.
실험 데이터셋
욕설 데이터셋은 Curse-Detection-Dataset, 혐오표현 데이터셋은 APEACH를 이용했습니다. 학습은 huggingface에 있는 bert, electra 모델들을 이용했습니다.
실험결과
욕설 분류에서는 사전 학습된 beomi/KcELECTRA-base 모델을 이용해 전이학습한 결과가 0.9167의 정확도로 가장 좋았습니다. 이제 이 모델을 이용해서 지식 증류를 수행했습니다. 이제 이 모델의 soft label을 가지고 student 모델들을 학습한 결과 monologg/koelectra-small-v3-discriminator이 가장 성능이 좋았습니다. 혐오표현 분류 역시 동일한 모델들이 가장 성능이 좋았습니다.
지식 증류를 이용할 때도 사전학습된 모델을 이용하여 전이학습을 하면 성능이 향상됬습니다. 사전학습하지 않은 동일한 모델을 사용하면 정확도가 10%p 가까이 떨어졌습니다.
Student 모델들은 Teacher에 비해 1/8 수준의 크기를 갖고, 추론 속도도 훨씬 빠릅니다. 욕설 분류 모델은 크게 정확도가 떨어지지 않았지만, 혐오표현 분류 모델은 Teacher에 비해 정확도가 벌어진 모습입니다.
데모
더 작아진 모델을 이용해서 욕설/혐오분류 데모를 수정했습니다.
https://share.streamlit.io/heegyukim/cursefilter/main/demo_app.py
앞으로의 개선방향
먼저 욕설분류모델 따로, 혐오분류모델 따로 하지 않고, 한 모델에서 여러 task를 수행하며 학습하는 Multitask Learning을 수행해서 성능을 추가로 개선해보고자합니다. 한 모델로 여러가지 추론을 해서 계산효율을 늘리고, 성능도 향상시켜볼 수 있지 않을까 싶습니다. Multitask Learning이 효과적이라면, 다른 좋은 데이터셋들을 추가로 이용해서 성능을 더 끌어올려볼 수도 있을 것입니다.
추가로 Quantization, Pruning 같은 방법으로도 모델압축을 시도해볼 수 있습니다.
