[論文レビュー] Decoupled Kullback-Leibler Divergence Loss

Jiequan Cui, Zhuotao Tian+4 more|arXiv (Cornell University)|May 23, 2023

Adversarial Robustness in Machine Learning被引用数 18

ひとこと要約

本論文は KL loss が重み付き MSE と soft labels を用いたクロスエントロピーからなるデカップルド形 (DKL) に同値であることを示し、非対称性を崩しグローバル情報を付加することで IKL を提案し、敵対的学習と知識蒸留において最先端の結果を達成する。

ABSTRACT

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 decomposed formulation of DKL loss, we have identified two areas for improvement. Firstly, we address the limitation of KL/DKL in scenarios like knowledge distillation by breaking its asymmetric optimization property. This modification ensures that the $\mathbf{w}$MSE component is always effective during training, providing extra constructive cues. Secondly, we introduce class-wise global information into KL/DKL to mitigate bias from individual samples. With these two enhancements, we derive the Improved Kullback-Leibler (IKL) Divergence loss and evaluate its effectiveness by conducting experiments on CIFAR-10/100 and ImageNet datasets, focusing on adversarial training, and knowledge distillation tasks. The proposed approach achieves new state-of-the-art adversarial robustness on the public leaderboard -- RobustBench and competitive performance on knowledge distillation, demonstrating the substantial practical merits. Our code is available at https://github.com/jiequancui/DKL.

研究の動機と目的

• KL発散の勾配最適化機構を説明し、それが Decoupled Kullback-Leibler (DKL) loss への等価性を示す。
• 非対称性とグローバルな文脈の欠如におけるDKLの限界を特定し、特に知識蒸馏において。
• 非対称性を崩し、グローバル情報を取り入れて Improved Kullback-Leibler (IKL) loss を形成することで改善を提案する。
• CIFAR-10/100 および ImageNet での敵対的学習と知識蒸留実験を通じて IKL の有効性を示す。

提案手法

• KL発散を、weighted MSE 項と soft labels を用いたクロスエントロピー項からなる Decoupled Kullback-Leibler (DKL) loss として再定式化する。
• DKL は入力間で非対称な勾配フローを持つ可能性があり、訓練中に wMSE 成分を減少させる可能性があることを示す。
• IKL を導入する: (i) 非対称性を崩して両方の損失項が最適化に寄与するようにし、(ii) グローバル（クラスごとの）情報を正則化として取り入れる。
• グローバル情報を weighted MSE に注入するために、各真のクラスの平均予測を用いてクラスごとの重みを定義する。
• CIFAR-10/100 および ImageNet にわたる敵対的学習と知識蒸留タスクで IKL を経験的に評価する。

実験結果

リサーチクエスチョン

• RQ1 勾配ベースの最適化の下で KL 発散はどのように挙動し、一般性を損なうことなく相補的な成分（wMSE と cross-entropy）に分解できるか？
• RQ2 勾配の非対称性に起因する知識蒸留などのトレーニングシナリオにおける DKL の制限は何か？
• RQ3 非対称性を崩し、グローバル情報を追加することで敵対的訓練と知識蒸留の性能を改善できるか？
• RQ4 提案された IKL loss は、対象タスクで CIFAR-10/100 および ImageNet において最先端の結果をもたらすか？

主な発見

• KL loss は、weighted MSE loss と soft labels を用いたクロスエントロピー損失からなるデカップルド形 (DKL) に数学的に等価である。
• DKL の最適化は入力に対して非対称であり、知識蒸留のようなシナリオで wMSE 成分を無視する原因となる。
• IKL は非対称性を崩し、グローバル情報を取り入れることで、クラス内の一貫性を向上させる。
• 実験により、IKL が CIFAR-10/100 および ImageNet で敵対的学習と知識蒸留の新しい最先端結果を達成することが示される。