Source-linked AI summary
Decoupled Kullback-Leibler Divergence Loss
Jiequan Cui, Zhuotao Tian, Zhisheng Zhong, Xiaojuan Qi, Bei Yu, Hanwang Zhang
TL;DR
KL is widely applied, but its gradient-optimization mechanism has not been thoroughly investigated. The paper proves KL’s equivalence to a decomposed loss, then uses that analysis to develop IKL. IKL achieves state-of-the-art adversarial robustness and competitive knowledge-distillation performance across the evaluated datasets.
Problem
KL is widely applied, but its underlying gradient-optimization mechanism may not be thoroughly investigated.
Source: §1 · p. 1Method
The paper proves KL and DKL gradient-equivalent, decomposes KL into wMSE and soft-label Cross-Entropy, then modifies asymmetry and adds class-wise global information to derive IKL.
Source: §1 · pp. 1–2Results
IKL achieves state-of-the-art adversarial robustness and competitive knowledge-distillation performance on CIFAR-10/100 and ImageNet.
Source: §1 · p. 2, §5 · p. 10Takeaways & Limitations
The reported results support IKL's effectiveness in adversarial robustness and knowledge distillation on the evaluated datasets.
Source: §5 · p. 10In knowledge distillation, the teacher is well-trained and fixed, and its probability vector is detached from gradient backpropagation.
Source: §3.1 · p. 4
Abstract
from arXiv · showhide
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.
1 Introduction
Although KL is widely used, its gradient-optimization mechanism may not be thoroughly examined. The paper analyzes this mechanism, decomposes KL into two loss terms, and uses the decomposition to derive IKL.
1 Introduction
- KL and DKL are mathematically equivalent for gradient optimization, with DKL comprising wMSE and soft-label Cross-Entropy.The decomposition makes it more convenient to analyze how KL works in training optimization.Source: §1 · pp. 1–2
- KL/DKL optimization is asymmetric, so wMSE can be overlooked in knowledge distillation when teacher logits are detached; IKL enables wMSE gradients on the other input.The introduction describes this change as alleviating the problem.Source: §1 · p. 2
- IKL achieves state-of-the-art adversarial robustness on RobustBench and competitive knowledge-distillation performance on CIFAR-10/100 and ImageNet.The experiments evaluate adversarial training and knowledge distillation.Source: Figure 2 · p. 2, §1 · p. 2
- Because sample-wise predictions can make wMSE optimization challenging on hard examples, IKL adds class-wise global information to regularize training.The paper combines this change with breaking optimization asymmetry to derive IKL.Source: §1 · p. 2
2 Related Work
Related work covers adversarial training, knowledge distillation, and other applications of KL divergence. The paper positions its analysis and IKL loss in relation to these uses.
Adversarial Robustness
Knowledge Distillation
- Knowledge distillation uses KL to align student and teacher output probabilities, with logits-based and feature-based approaches among later methods [42] [44] [45] [46] [19] [47] [48] [49] [50] [51] [31] [53] [54] [55] [56].The paper decouples KL and addresses a limitation in settings such as distillation.Source: §2 · p. 3
3 Method
The method analyzes KL through gradient optimization, derives an equivalent decoupled form, then uses that decomposition to address asymmetric optimization and sample-wise weighting in IKL. It evaluates the designs in adversarial training and knowledge distillation.
3.1 Preliminary
- In adversarial training, KL aligns adversarial-example outputs with clean-image outputs; in distillation, it aligns student predictions with a fixed teacher’s predictions.The teacher output is detached from backpropagation during distillation.Source: §3.1 · pp. 3–4
3.2 Decoupled Kullback-Leibler Divergence Loss
- For softmax models, KL is equivalent in gradient optimization to DKL, a composition of weighted MSE and soft-label Cross-Entropy.The paper says the decomposition makes KL’s optimization behavior easier to analyze.Source: §3.2 · p. 4
- KL/DKL’s asymmetric gradients can leave wMSE ineffective when teacher logits are detached, as in knowledge distillation.The paper identifies this as a potential source of performance degradation.Source: §3.2 · p. 4
3.3 Improved Kullback-Leibler Divergence Loss
- IKL breaks gradient asymmetry to keep wMSE active and replaces sample-wise weights with class-wise global weights to mitigate bias and regularize intra-class consistency.Class-wise weights are described as especially useful late in training, when they can provide correct predictions for optimizing wMSE.Source: §3.3 · p. 5
- Sample-dependent wMSE weights can misdirect optimization toward a hard example’s predicted class rather than its ground-truth class, destabilizing training.The authors also report that the choice of these weights significantly affects adversarial robustness.Source: §3.2 · pp. 4–5, §3.3 · p. 5
3.4 A Case Study and Analysis
- Ablations confirm DKL’s comparable performance to KL and the effectiveness of global information and broken asymmetry; class-wise weights mitigate sample-noise bias, with IKL-AT showing mostly larger CIFAR-100 margins and more compact, separable features than TRADES.The authors link the larger margins to stronger robustness.Source: §3.4 · pp. 6–7, Figure 3 · p. 6
4 Experiments
The experiments evaluate IKL in adversarial training and knowledge distillation on CIFAR-10/100 and ImageNet, including imbalanced data and robustness ablations.
4.1 Adversarial Robustness
- With 50M generated images, IKL-trained WideResNet-28-10 reaches 73.85% natural accuracy and 39.18% Auto-Attack robustness on CIFAR-100.Without extra data, IKL with AutoAug reaches 32.53% robustness, a new state of the art in that setting.Source: §4.1 · p. 8
- On CIFAR-10, IKL reaches 57.09% robustness with basic augmentation, 0.92% above AWP, and 67.75% with generated data, improving the state of the art by 0.44%.With basic augmentation, natural-image top-1 accuracy is 84.80%.Source: §4.1 · p. 8
4.2 Knowledge Distillation
- IKL-KD matches or exceeds prior methods across CIFAR-100 settings, achieving the best performance in 4 of 6 same-structure settings and 3 of 5 different-structure configurations.The comparisons cover teacher-student architectures with the same and different unit structures.Source: §4.2 · p. 9
- With a ResNet34 teacher, IKL-KD gives ResNet18 71.91% top-1 accuracy; with a ResNet50 teacher, MobileNet reaches 72.84%, while computation costs fall 38% and 52% versus ReviewKD.The cost reductions correspond respectively to ResNet34–ResNet18 and ResNet50–MobileNet distillation.Source: §4.2 · p. 9
4.3 Knowledge Distillation on Imbalanced Data
- On ImageNet-LT, IKL-KD consistently outperforms KL-KD across the reported Many-shot, Medium-shot, and All classes.Evaluation reports top-1 accuracy for the three class-frequency groups.Source: §4.3 · p. 9, Table 5 · p. 9
4.4 Ablation Studies
- Increasing τ from 1 to 4 improves clean accuracy by 5.75% while keeping robustness comparable; IKL also exceeds TRADES across tested perturbation sizes by 1.34% average robustness.The study selects τ = 4 for strongest robustness; robustness declines as perturbation size increases for both methods.Source: §4.4 · p. 10
- Under PGD and CW attacks, IKL models show similar robustness at 10 and 20 iterations, which the authors interpret as evidence against obfuscated gradients.The evaluations use perturbation size 8/255 and step size 2/255.Source: §4.4 · p. 10
5 Conclusion and Limitation
The paper concludes that IKL combines the KL decomposition with two targeted modifications, reporting strong adversarial robustness and competitive distillation performance.
5 Conclusion and Limitation
- IKL decouples KL into weighted MSE and soft-label cross-entropy, then breaks asymmetric optimization and adds class-wise global information.Experiments on CIFAR-10/100 and ImageNet report new state-of-the-art adversarial robustness and competitive knowledge distillation performance.Source: §5 · p. 10
A.1 Proof to Theorem 1
The proof establishes that KL and DKL have identical gradients for the same inputs, making them equivalent for training optimization.
A.1 Proof to Theorem 1
- KL and DKL produce the same derivatives for identical inputs, establishing their equivalence in gradient-based training optimization.The proof compares the derivatives obtained by expanding DKL and applying the chain rule to KL.Source: §A.1 · pp. 14–15
A.2 New state-of-the-art robustness on CIFAR-100/10
The CIFAR-10/100 results report new state-of-the-art robustness with and without generated data.
A.2 New state-of-the-art robustness on CIFAR-100/10
- 32.53% robustness on CIFAR-100 without generated data exceeds the previous best by 0.68% while saving 33.3% computation; with generated data, natural accuracy reaches 73.85%, 1.27% above the previous best, with strongest robustness.The paper reports new state-of-the-art robustness on CIFAR-10 and CIFAR-100 under both data settings.Source: §A.2 · p. 15
A.3 Comparisons on CIFAR-100 for Knowledge Distillation
On CIFAR-100, IKL-KD performs best in most tested teacher–student architecture settings and slightly improves results under a strong training scheme. The reported wins span both matching and differing architectures.
A.3 Comparisons on CIFAR-100 for Knowledge Distillation
- IKL-KD achieves the best results in 4 of 6 same-architecture settings and 3 of 5 different-architecture settings on CIFAR-100.The comparisons use Tables 13 and 14, respectively.Source: §A.3 · p. 16
- IKL-KD shows slightly better results with BEiT-Large as teacher and ResNet-50 as student under a strong training scheme.This comparison follows the concurrent work [71].Source: §A.3 · p. 16
A.4 Other Applications with IKL
The paper tests IKL in semi-supervised learning and semantic-segmentation distillation by replacing the consistency or KL loss, respectively. Reported results favor IKL in both applications.
A.4 Other Applications with IKL
- With two labeled and 500 unlabeled images per CIFAR-100 class, Mean-Teacher with IKL surpasses FixMatch.The comparison uses ViT-small and is reported in Table 16.Source: §A.4 · p. 16, Table 16 · p. 16
- Replacing APD’s KL loss with IKL yields better performance in semantic-segmentation distillation on ADE20K.The comparison is reported in Table 17.Source: §A.4 · p. 16, Table 17 · p. 16
A.5 Complexity of IKL
IKL updates global class-wise prediction scores during training, with nearly negligible extra computation relative to model forward and backward passes. The paper also provides a memory-efficient wMSE implementation for dense prediction tasks.
A.5 Complexity of IKL
- Updating the global class-wise prediction scores W ∈ R^(C×C) adds computational cost that can be nearly ignored relative to model forward and backward passes.C is the number of classes.Source: §A.5 · p. 16
- For dense prediction, the ∆a and ∆b tensors can require large GPU memory, so the paper supplies a memory-efficient wMSE implementation.The implementation is given in Algorithm 2.Source: §A.5 · p. 17
- Algorithm 1 computes weighted MSE and soft-label cross-entropy from two logits inputs, then returns their weighted sum.Its inputs include one-hot labels, W, and coefficients α and β.Source: §A.5 · p. 17
A.6 Connection between IKL and the Jensen-Shannon (JS) Divergence
On ImageNet-LT knowledge distillation, the paper reports that IKL outperforms Jensen–Shannon divergence, whose teacher soft labels do not contribute through a cross-entropy term. The checklist also reports proof, reproducibility, and compute details, while noting that most tables lack error bars because of computation cost.
A.6 Connection between IKL and the Jensen-Shannon (JS) Divergence
- In the authors’ account, JSD’s teacher soft labels have no cross-entropy effect, which may underlie its poorer results than KD and IKL-KD.The passage describes this as a possible underlying reason.Source: §A.6 · p. 17
- IKL performs better than JSD on knowledge distillation, with the comparison reported for ImageNet-LT.The cited passage identifies Table 18 as the empirical comparison.Source: §A.6 · p. 17, Table 18 · p. 18
NeurIPS Paper Checklist
- The checklist states that Theorem 1 is proved in Appendix A.1 and that experimental settings and hyperparameter ablations are provided.The settings are in Section 4 and ablations in Appendix 4.4.Source: NeurIPS Paper Checklist · pp. 19–20
- The checklist says the paper provides a code link as its open-access avenue for reproduction.The justification states that the code link is provided in the paper.Source: NeurIPS Paper Checklist · p. 20
- Error bars appear in Appendix Table 14, while other tables omit them because of heavy computation.The paper also reports adversarial-training runs taking around two days with four GPUs and ImageNet distillation runs taking about one day on eight Nvidia GeForce 3090 GPUs.Source: NeurIPS Paper Checklist · p. 21