[논문 리뷰] Scaling Ensemble Distribution Distillation to Many Classes with Proxy Targets

Max Ryabinin, Andrey Malinin외 1명|arXiv (Cornell University)|2021. 05. 14.

Domain Adaptation and Few-Shot Learning참고 문헌 31인용 수 8

한 줄 요약

이 논문은 1,000~40,000개의 클래스를 가진 대규모 분류 작업에서 앙상블 분포 복제(EnD)의 안정성을 높이기 위해 프록시-디리클레 복제를 제안한다. 앙상블 출력에서 유도된 프록시 디리클레 타깃으로 반대 KL 발산을 최소화함으로써, 낮은 확률을 가진 꼬리 클래스로 인한 기울기 폭발 문제를 해결하여 수렴 가능하게 하고, ImageNet 및 WMT17 번역 작업에서 기준선 복제 및 단일 모델보다 정확도와 불확실성 추정 모두에서 뛰어난 성능을 달성한다.

ABSTRACT

Ensembles of machine learning models yield improved system performance as well as robust and interpretable uncertainty estimates; however, their inference costs may often be prohibitively high. \emph{Ensemble Distribution Distillation} is an approach that allows a single model to efficiently capture both the predictive performance and uncertainty estimates of an ensemble. For classification, this is achieved by training a Dirichlet distribution over the ensemble members' output distributions via the maximum likelihood criterion. Although theoretically principled, this criterion exhibits poor convergence when applied to large-scale tasks where the number of classes is very high. In our work, we analyze this effect and show that the Dirichlet log-likelihood criterion classes with low probability induce larger gradients than high-probability classes. This forces the model to focus on the distribution of the ensemble tail-class probabilities. We propose a new training objective that minimizes the reverse KL-divergence to a \emph{Proxy-Dirichlet} target derived from the ensemble. This loss resolves the gradient issues of Ensemble Distribution Distillation, as we demonstrate both theoretically and empirically on the ImageNet and WMT17 En-De datasets containing 1000 and 40,000 classes, respectively.

연구 동기 및 목표

• 1,000~40,000개의 클래스를 가진 고차원 분류 작업에서 앙상블 분포 복제(EnD)의 수렴 불량 문제를 해결하기 위해.
• 낮은 확률을 가진 꼬리 클래스가 기울기의 주요 기여를 차지하는 디리클레 로그우도 기준에서 기울기 불균형의 근본 원인을 규명하기 위해.
• 예측 성능과 앙상블 모델의 불확실성 추정을 모두 유지하는 안정적인 학습 목표를 개발하기 위해.
• 실제의 컴퓨팅 자원 제약이 있는 위험 중심 응용 분야에서 앙상블의 확장 가능한 복제를 가능하게 하기 위해.

제안 방법

• 앙상블 구성원 출력의 경험적 분포에서 프록시-디리클레 분포를 유도하여 핵심 통계적 성질을 유지하기 위해.
• 학생 모델의 디리클레 출력과 프록시-디리클레 타깃 간의 반대 KL 발산을 주 학습 목표로 사용하기 위해.
• 표준 디리클레 로그우도 기준을 반대 KL 목표로 대체하여 최적화의 안정성을 향상시키기 위해.
• 복제 과정에서 GPU 메모리에 큰 앙상블(10개 모델)을 수용하기 위해 기울기 누적 및 혼합 정밀도 학습을 적용하기 위해.
• 디코딩을 위해 범위 탐색을 사용하고, 범위 크기 5로 설정하며, 불확실성 평가를 위해 EoE 및 RMI 지표를 사용하기 위해.
• 배치 크기를 32K 토큰으로 증가시키고 32단계에 걸쳐 기울기 누적을 수행하며, 총 20,000 스텝 동안 복제 모델을 학습하기 위해.

실험 결과

연구 질문

• RQ1 왜 표준 앙상블 분포 복제가 1,000~40,000개의 클래스를 가진 대규모 분류 작업에서 수렴하지 못하는가?
• RQ2 앙상블 출력에서의 클래스 확률 분포가 복제 과정 중 기울기 크기에 어떤 영향을 미치는가?
• RQ3 앙상블 출력에서 유도된 프록시 분포로 학습된 복제 모델의 안정성이 향상되는가?
• RQ4 프록시-디리클레 타깃으로 반대 KL 발산을 최소화하는 것이 정방향 KL 또는 표준 로그우도보다 복제 성능에서 뛰어나지 않는가?
• RQ5 제안된 방법이 고카디널리티 작업에서 원래 앙상블과 비교해 정확도와 불확실성 추정을 모두 유지할 수 있는가?

주요 결과

• 디리클레 로그우도 기준은 낮은 확률을 가진 꼬리 클래스에서 상당히 큰 기울기를 유도하여 최적화의 불안정성을 초래한다.
• ImageNet(1,000개 클래스)에서 프록시-디리클레 방법은 상위 1위 정확도 77.1%와 이상치 탐지에 대한 RMI 11.8%를 달성하여 표준 EnD 및 단일 모델을 능가했다.
• WMT17 En-De 번역(40,000개 클래스)에서 이 방법은 BLEU 29.5점과 RMI 35.8점을 기록하여 앙상블의 성능을 맞추거나 초월했으며, 불확실성 탐지 측면에서도 유사하거나 뛰어난 성능을 보였다.
• 프록시-디리클레 타깃으로 반대 KL 발산을 최소화하는 것이 정방향 KL 또는 로그우도와 달리 학습을 안정화시키고 기울기 폭발을 방지했다.
• 프록시-디리클레 방법은 고출력 카디널리티를 가진 이미지 및 시퀀스 모델링 작업에서 수렴성과 신뢰할 수 있는 불확실성 추정을 모두 보장했다.
• 이 방법은 컴퓨팅 자원 제약이 있는 안전 중심 응용 분야에서 앙상블 수준의 불확실성 추정을 유지하면서 추론 비용을 줄일 수 있도록 앙상블 복제를 가능하게 했다.