[论文解读] 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

一句话总结

该论文提出了一种代理狄利克雷蒸馏（Proxy-Dirichlet distillation）方法，以稳定大规模分类任务（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.

研究动机与目标

• 解决在具有大量类别的高维分类任务中，集成分布蒸馏（EnD）收敛性差的问题。
• 识别狄利克雷对数似然准则中梯度不平衡的根本原因，即低概率尾部类别主导了梯度。
• 开发一种稳定的训练目标，以保留集成模型的预测性能和不确定性估计。
• 实现在计算资源受限、风险敏感的现实应用中，对大规模集成模型进行可扩展的蒸馏。

提出的方法

• 从集成成员输出的经验分布中推导出代理狄利克雷分布，以保留其关键统计特性。
• 使用学生模型的狄利克雷输出与代理狄利克雷目标之间的反向KL散度作为主要训练目标。
• 用反向KL目标替代标准狄利克雷对数似然准则，以稳定优化过程。
• 在蒸馏过程中，通过梯度累积和混合精度训练，将大型集成模型（10个模型）适配到GPU内存中。
• 使用束搜索（束大小为5）进行解码，并通过EoE和RMI指标评估不确定性。
• 以32K token的批量大小训练蒸馏模型共20,000步，并在32个步骤内进行梯度累积。

实验结果

研究问题

• RQ1 为何标准的集成分布蒸馏在1,000–40,000类的大规模分类任务上无法收敛？
• RQ2 集成输出中的类别概率分布如何影响蒸馏过程中的梯度大小？
• RQ3 能否通过从集成输出中推导出的代理分布来稳定蒸馏模型的训练？
• RQ4 与前向KL或标准对数似然相比，最小化到代理狄利克雷目标的反向KL散度是否在蒸馏中表现更优？
• RQ5 所提出的方法是否能在高基数任务中同时保持与原始集成模型相当的准确率和不确定性估计？

主要发现

• 狄利克雷对数似然准则导致低概率尾部类别的梯度显著增大，引发优化不稳定性。
• 在ImageNet（1,000类）上，代理狄利克雷方法实现了77.1%的top-1准确率和11.8%的RMI（分布外检测），优于标准EnD和单模型。
• 在WMT17 En-De翻译任务（40,000类）上，该方法实现了29.5 BLEU和35.8 RMI，其不确定性检测性能与集成模型相当或更优。
• 与前向KL或对数似然相比，对代理狄利克雷目标的反向KL散度稳定了训练过程，避免了梯度爆炸。
• 代理狄利克雷方法在高输出基数的图像和序列建模任务中均实现了收敛和可靠的不确定性估计。
• 该方法通过在降低推理成本的同时保留集成级别的不确定性，使得在计算受限、安全性关键的应用中实现集成蒸馏成为可能。