[論文レビュー] Scaling Ensemble Distribution Distillation to Many Classes with Proxy Targets

Max Ryabinin, Andrey Malinin+1 more|arXiv (Cornell University)|May 14, 2021

Domain Adaptation and Few-Shot Learning参考文献 31被引用数 8

ひとこと要約

本稿では、1,000〜40,000クラスのスケールで大規模分類を対象としたアンサンブル分布蒸留（EnD）の安定化を図るため、Proxy-Dirichlet蒸留を提案する。アンサンブル出力から導出された代理ディリクレ分布をターゲットとして逆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）の収束不良を解消すること。
• 低確率の末尾クラスが勾配に与える影響を引き起こすディリクレ対数尤度基準における勾配不均衡の根本的原因を特定すること。
• アンサンブルモデルが持つ予測性能と不確実性推定を両立する安定な学習目的を構築すること。
• 実世界の計算リソース制限やリスクセンシティブな応用に適したスケーラブルなアンサンブル蒸留を可能にすること。

提案手法

• アンサンブルメンバー出力の経験的分布から、主要な統計的性質を保持したProxy-Dirichレ分布を導出する。
• 学生モデルのディリクレ出力とProxy-Dirichレターゲットとの間の逆KLダイバージェンスを主な学習目的として用いる。
• 標準的なディリクレ対数尤度基準に代えて、逆KL目的を導入することで最適化の安定性を向上させる。
• 蒸留中にGPUメモリに収まるように、勾配蓄積とミックスドプレシジョン学習を適用し、大規模アンサンブル（10モデル）を処理する。
• デコードにはビームサーチ（ビームサイズ5）を用い、不確実性評価にはEoEおよびRMI指標を適用する。
• バッチサイズを32Kトークンに増加させ、32ステップにわたる勾配蓄積を伴い、全20,000ステップで蒸留モデルを学習する。

実験結果

リサーチクエスチョン

• RQ1 なぜ標準的なアンサンブル分布蒸留は、1,000〜40,000クラスの大規模分類タスクで収束しないのか？
• RQ2 アンサンブル出力におけるクラス確率分布は、蒸留過程における勾配の大きさにどのように影響を与えるか？
• RQ3 アンサンブル出力から導出された代理分布を用いることで、蒸留モデルの学習が安定化できるか？
• RQ4 Proxy-Dirichレターゲットに対する逆KLダイバージェンスの最小化は、前方KLや標準的対数尤度基準を上回る性能を発揮するか？
• RQ5 提案手法は、高基数タスクにおいて、元のアンサンブルと同等の精度と不確実性推定を保持できるか？

主な発見

• ディリクレ対数尤度基準は、低確率の末尾クラスに対して顕著に大きな勾配を誘発し、最適化の不安定性を引き起こす。
• ImageNet（1,000クラス）では、Proxy-Dirichレ手法がトップ-1正解率77.1%、OoD検出のRMI指標11.8%を達成し、標準EnDおよび単一モデルを上回った。
• WMT17 En-De翻訳（40,000クラス）では、29.5 BLEUおよび35.8 RMIを達成し、不確実性検出においてアンサンブルと同等またはそれを上回る性能を示した。
• Proxy-Dirichレターゲットに対する逆KLダイバージェンス最小化により、勾配爆発を防ぎ、前方KLや対数尤度基準と比較して学習が安定化した。
• Proxy-Dirichレ手法により、出力基数が大きい画像および系列モデルタスクにおいて、収束と信頼性のある不確実性推定が可能になった。
• 本手法により、推論コストを削減しながらアンサンブルレベルの不確実性を保持できるため、計算リソース制限や安全性が求められる実用的応用に適したアンサンブル蒸留が可能になった。