提出改进版KL散度损失，提升模型鲁棒性与知识蒸馏效果。
Generalized Kullback-Leibler Divergence Loss
- 将KL散度拆解为wMSE与交叉熵组合，实现对称优化
- 在CIFAR/ImageNet/CLIP上实现新SOTA对抗鲁棒性
- 适合对抗训练与知识蒸馏场景的模型优化
本文深入研究KL散度损失，数学证明其等价于由加权均方误差（wMSE）和含软标签的交叉熵组成的解耦KL（DKL）损失。基于DKL的解耦结构，提出两项改进：一是通过平滑权重函数打破原KL损失的非对称优化特性，缓解高预测得分类别的收敛难题；二是引入类别级全局信息，降低单样本偏差。由此构建广义KL（GKL）损失，并在CIFAR-10/100、ImageNet及视觉语言数据集上验证其有效性，尤其在对抗训练任务中达到RobustBench公开排行榜新SOTA性能，同时在各类模型的知识蒸馏任务中表现优异，充分展现其实用价值。代码已开源。
原文摘要 · Abstract (English)
In this paper, we delve deeper into the Kullback-Leibler (KL) Divergence loss and mathematically prove that it is equivalent to the Decoupled Kullback-Leibler (DKL) Divergence loss that consists of (1) a weighted Mean Square Error (wMSE) loss and (2) a Cross-Entropy loss incorporating soft labels. Thanks to the decoupled structure of DKL loss, we have identified two areas for improvement. Firstly, we address the limitation of KL loss in scenarios like knowledge distillation by breaking its asymmetric optimization property along with a smoother weight function. This modification effectively alleviates convergence challenges in optimization, particularly for classes with high predicted scores in soft labels. Secondly, we introduce class-wise global information into KL/DKL to reduce bias arising from individual samples. With these two enhancements, we derive the Generalized Kullback-Leibler (GKL) Divergence loss and evaluate its effectiveness by conducting experiments on CIFAR-10/100, ImageNet, and vision-language datasets, focusing on adversarial training, and knowledge distillation tasks. Specifically, we achieve new state-of-the-art adversarial robustness on the public leaderboard -- RobustBench and competitive knowledge distillation performance across CIFAR/ImageNet models and CLIP models, demonstrating the substantial practical merits. Our code is available at https://github.com/jiequancui/DKL.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。