[논문 리뷰] Hydra: Preserving Ensemble Diversity for Model Distillation

Linh Tran, Bastiaan S. Veeling외 8명|arXiv (Cornell University)|2020. 01. 14.

Adversarial Robustness in Machine Learning참고 문헌 25인용 수 36

한 줄 요약

Hydra는 단일 공유 바디에 다수의 헤드를 두어 앙상블을 증류하고, 멤버별 예측과 불확실성을 보존하여 표준 증류보다 예측 성능과 불확실성 정량화 모두를 향상시킵니다.

ABSTRACT

Ensembles of models have been empirically shown to improve predictive performance and to yield robust measures of uncertainty. However, they are expensive in computation and memory. Therefore, recent research has focused on distilling ensembles into a single compact model, reducing the computational and memory burden of the ensemble while trying to preserve its predictive behavior. Most existing distillation formulations summarize the ensemble by capturing its average predictions. As a result, the diversity of the ensemble predictions, stemming from each member, is lost. Thus, the distilled model cannot provide a measure of uncertainty comparable to that of the original ensemble. To retain more faithfully the diversity of the ensemble, we propose a distillation method based on a single multi-headed neural network, which we refer to as Hydra. The shared body network learns a joint feature representation that enables each head to capture the predictive behavior of each ensemble member. We demonstrate that with a slight increase in parameter count, Hydra improves distillation performance on classification and regression settings while capturing the uncertainty behavior of the original ensemble over both in-domain and out-of-distribution tasks.

연구 동기 및 목표

• 증류 후에도 앙상블 불확실성을 유지해야 할 필요성에서 동기 부여.
• 멤버별 행동을 보존하기 위한 다중 헤드 증류 아키텍처를 제안.
• 분류 및 회귀 작업에서 Hydra를 표준 증류 및 이전 네트워크와 비교 평가.

제안 방법

• Hydra를 도입: 하나의 공유 바디와 M개의 헤드(앙상블 멤버당 하나).
• 각 헤드는 특정 앙상블 멤버를 모방; 바디는 공유 특징 표현을 제공합니다.
• 평균 KL 발산을 각 헤드와 대응하는 앙상블 멤버(분류) 또는 가우시안 출력(회귀) 간에 최소화하여 학습합니다.
• 훈련 중 분포를 가열하기 위해 온도 T를 사용하여 교차지지(Cross-support)를 개선합니다.
• Two-phase training: 먼저 평균 앙상블을 모방(Hinton head), 그런 다음 모든 헤드가 개별 멤버를 맞추도록 학습합니다."
• Knowledge Distillation 및 Prior Networks와 비교하여 데이터셋에 대해 NLL, Brier 점수, 정확도 및 모델 불확실성을 보고합니다.

실험 결과

연구 질문

• RQ1 Hydra가 평균화 기반 증류에 비해 앙상블 다양성을 충실히 보존할 수 있는가?
• RQ2 Hydra가 도메인 내외 데이터에서 예측 성능 및 불확실성 정량화를 향상시키는가?
• RQ3 Hydra가 매개변수 효율성과 앙상블 다양성에 대한 충실도 간 trade-off를 어떻게 달성하는가?
• RQ4 Hydra가 분류 및 회귀 작업에 미치는 영향은 무엇인가?

주요 결과

모델	ACC (MNIST)	NLL (MNIST)	BS (MNIST)	MU (MNIST)	ACC (CIFAR-10)	NLL (CIFAR-10)	BS (CIFAR-10)	MU (CIFAR-10)
Ensemble (M=50)	0.9851	0.0439	-0.9780	9.97e-06	0.9226	0.2392	-0.9033	0.1055
Prior Networks	0.9842	0.0521	-0.9285	0.1158	0.8731	0.4392	-0.8231	0.0280
Knowledge distillation	0.9843	0.0497	-0.9764	N/A	0.8933	0.3598	-0.8373	N/A
Hydra (head=[100,100,10])	0.9857	0.0465	-0.9776	2.28e-05	0.8992	0.3179	-0.8468	0.0074

• Hydra는 MNIST 및 CIFAR-10에서 앙상블 예측 성능을 따라잡거나 능가합니다.
• MNIST에서 Hydra의 NLL 0.0465, Brier −0.9776으로, 앙상블 NLL 0.0439 및 Brier −0.9780에 근접하고 MU 2.28e-5를 기록합니다.
• CIFAR-10에서 Hydra는 ACC 0.8992 및 NLL 0.3179를 달성하여 다른 증류 방법보다 앙상블에 더 근접하고 MU 0.0074를 기록합니다.
• Hydra는 여러 메트릭에서 Knowledge Distillation 및 Prior Networks를 능가하며 특히 불확실성 정량화(MU) 및 NLL에서 우수합니다.
• Hydra는 매모리 증가를 적당히 하면서 앙상블 다양성에 대한 충실도와의 실용적 균형을 제공합니다.