[논문 리뷰] On Compressing U-net Using Knowledge Distillation

Karttikeya Mangalam, Mathieu Salzamann|arXiv (Cornell University)|2018. 12. 01.

Energy Efficient Wireless Sensor Networks참고 문헌 11인용 수 7

한 줄 요약

이 논문은 지식 정련 기반 접근법을 제안하여 U-Net 아키텍처를 1000배 이상 압축함으로써 기존 크기의 0.1%로 파라미터를 감소시키면서도 거의 동일한 분할 정확도를 유지한다. 정제 과정 중에 배치 정규화와 클래스 재가중치를 학생 네트워크에 통합함으로써, 학습의 불안정성과 클래스 불균형 문제를 해결하여 표준 정련 기법이 실패하는 상황에서도 신뢰할 수 있는 압축을 가능하게 한다.

ABSTRACT

We study the use of knowledge distillation to compress the U-net architecture. We show that, while standard distillation is not sufficient to reliably train a compressed U-net, introducing other regularization methods, such as batch normalization and class re-weighting, in knowledge distillation significantly improves the training process. This allows us to compress a U-net by over 1000x, i.e., to 0.1% of its original number of parameters, at a negligible decrease in performance.

연구 동기 및 목표

• 모바일 디바이스와 같은 자원이 제한된 환경에서 큰 U-Net 모델을 구현하는 데 도전하는 것.
• 지식 정련이 매우 작은 크기로 U-Net 아키텍처를 효과적으로 압축할 수 있는지 조사하는 것.
• 표준 정련 기법이 매우 작은 U-Net 변형을 압축할 때 실패하는 원인을 규명하고 해결하는 것.
• 아키텍처 및 손실 함수 수정을 통해 소규모 학생 네트워크의 정련 신뢰성을 향상시키는 것.

제안 방법

• 수축 경로의 모든 합성곱 레이어에 배치 정규화 레이어를 도입하여 내부 공변수 이동을 줄이고 학습을 안정화시킨다.
• 클래스 불균형을 해결하기 위해 교차 엔트로피 손실에 클래스 재가중치를 적용하며, 배경 대 대비 배경 픽셀 비율을 바탕으로 전경 가중치를 17.8로 설정한다.
• 혼합 지식 정련을 사용하여 T=5에서 온도 조정된 소프트 레이블 손실을 4-Unet 교사 네트워크로부터 제공하고, 소프트 레이블(T=5)과 하드 레이블을 결합한다.
• 학생 네트워크를 소프트 레이블(정련 손실)과 원래의 교차 엔트로피 손실(하드 레이블)을 모두 사용하여 훈련한 후, 최종 출력을 온도 조정하여 소프트닝한다.
• 하드 레이블로 끝까지 훈련된 4-Unet을 교사 모델로 사용하고, 이를 파라미터가 최소한인 2-Unet으로 정련한다.
• 훈련 중에 소프트 레이블 손실을 T²로 스케일링하여 온도에 관계없이 일관된 기울기 크기를 유지한다.

실험 결과

연구 질문

• RQ1 표준 지식 정련 기법이 U-Net을 원래 파라미터 수의 1% 미만으로 신뢰성 있게 압축할 수 있는가?
• RQ2 왜 표준 정련 기법은 U-Net을 매우 작은 크기로 압축할 때 실패하는가(예: 초기 채널 수가 2인 경우)?
• RQ3 배치 정규화와 클래스 재가중치는 압축된 U-Net 학생 네트워크의 학습을 어떻게 향상시키는가?
• RQ4 정규화 기법을 강화한 정련을 통해 2-Unet이 64-Unet과 유사한 분할 성능을 달성할 수 있는가?

주요 결과

• 122,394개의 파라미터만을 가진 4-Unet은 테스트 손실 0.0974와 IoU 0.807을 기록하여 원래 64-Unet의 0.804 IoU와 유사한 성능을 보이며, 정련 없이도 파라미터 감소가 효과적임을 입증한다.
• 표준 정련(기본형 또는 혼합형)은 2-Unet을 훈련시키지 못하며, 표 2에 따르면 테스트 손실이 0.505~0.507로 나타나 일반화 성능이 열악하고 불안정함을 보여준다.
• 배치 정규화와 클래스 재가중치를 적용한 2-Unet은 혼합 정련을 통해 IoU 0.759와 교차 엔트로피 손실 0.135를 기록하며, 정련 없이 학습된 2-Unet(이상수 0.752, 손실 0.134)보다 우수한 성능을 낸다.
• 최종 2-Unet 모델은 오직 30,902개의 파라미터만을 사용하여 원래 64-Unet의 3100만 개 파라미터의 0.1%에 불과하지만, 거의 동일한 성능을 유지한다.
• 정련 없이도 수정된 방식으로 2-Unet을 훈련한 결과 테스트 손실이 0.307로 나타나 정련이 성능을 확보하는 데 필수적임을 확인한다.
• 배치 정규화와 클래스 재가중치의 조합은 표준 정련 기법이 실패하는 저파라미터 수준에서도 성공적인 정련을 가능하게 한다.