Generalized Kullback-Leibler Divergence Loss

在本文中，我们深入探讨了 Kullback-Leibler (KL) 散度损失，并从数学上证明了它等价于由两部分组成的 Decoupled Kullback-Leibler (DKL) 散度损失：(1) 加权均方误差 (wMSE) 损失，以及 (2) 融入软标签的交叉熵损失。得益于 DKL 损失的解耦结构，我们发现了两个可以改进的方向。首先，针对知识蒸馏等场景中 KL 损失的局限性，我们通过打破其非对称优化特性并引入更平滑的权重函数来解决这一问题。这种修改有效缓解了优化过程中特别是对于软标签中高预测分数类别收敛困难的问题。其次，我们将类级别的全局信息引入 KL/DKL 损失中，以减少单个样本带来的偏差。基于这两项改进，我们推导出了广义 Kullback-Leibler (GKL) 散度损失，并通过在 CIFAR-10/100、ImageNet 和视觉-语言数据集上的实验评估其有效性，重点考察了其在对抗训练和知识蒸馏任务中的表现。具体来说，我们在公开排行榜 RobustBench 上取得了最新的对抗鲁棒性最佳结果，并在 CIFAR/ImageNet 模型和 CLIP 模型的知识蒸馏任务中展现了具有竞争力的性能，充分证明了其实际应用价值。我们的代码已开源，地址为 https://github.com/jiequancui/DKL。

该论文旨在解决Kullback-Leibler (KL) 散度损失在知识蒸馏等场景中的优化难题，特别是对于高预测分数的软标签类别所面临的收敛挑战。此外，它还试图减少由于个别样本引起的偏差问题。

论文的关键思路是通过将传统的KL散度损失分解为两部分：加权均方误差（wMSE）和包含软标签的交叉熵损失，形成Decoupled Kullback-Leibler (DKL) 散度损失。进一步地，作者提出了Generalized Kullback-Leibler (GKL) 散度损失，引入了类别的全局信息，并改进了权重函数，以克服原有KL损失的不对称优化特性。这一方法不仅解决了收敛性问题，还减少了样本带来的偏见。

实验设计涵盖了CIFAR-10/100、ImageNet以及视觉-语言数据集，在对抗训练和知识蒸馏任务中进行了评估。研究结果表明，GKL损失在RobustBench上实现了新的最先进对抗鲁棒性，并在CIFAR/ImageNet模型和CLIP模型的知识蒸馏性能方面表现出色。代码已开源至GitHub (https://github.com/jiequancui/DKL)，便于后续研究者复现和扩展。

近期相关研究包括《Distilling the Knowledge in a Neural Network》探讨了神经网络的知识蒸馏，《Adversarial Examples Are Not Bugs, They Are Features》分析了对抗样本的本质。其他相关工作如《Improved Regularization of Convolutional Neural Networks with Cutout》也涉及到了模型训练中的正则化技术。