Decoupled Kullback-Leibler Divergence Loss

Thanks for your suggestions and valuable comments. Here we provide our responses to address your concerns.

Q1: Ablation with TRADES and MART.

Thanks for your suggestion. TRADES is our baseline. We have already listed the results in Table 1 of the main paper. To highlight the significant improvements from our IKL loss, here we present the experimental results with TRADES and MART on CIFAR-100. Equipped with IKL loss, the robustness is enhanced for both methods. We transfer the hyper-parameters from TRADES to MART. Carefully tuning the hyper-parameters can potentially further improve performance. We would include the comparisons in the new revision.

Q2: Evaluation under PGD and CW attacks.

As suggested, we also test the robustness of our models under PGD, and CW attacks with 10 and 20 iterations. The perturbation size and step size are set to 8/255 and 2/255 respectively. As shown in the Table, with increasing iterations from 10 to 20, our models show similar robustness, demonstrating that our models don't suffer from the obfuscated gradients problem.

The worst-case is in bold.

Q3: Applications of IKL loss on semi-supervised learning.

Please refer to our global rebuttal Q1.

Q4: The first terms in Eq5, 6, and 7 are slightly different, but the names are the same; could you please clarify this?

Eq.(5) is equivalent to KL loss regarding gradient optimization. In Eq. (6), the stop gradient operation in $\mathcal{S}(\Delta \mathbf{m})$ is applied because the teacher model is fixed and thus its outputs $o_{m}$ are detached from the gradient backpropagation during distillation. With Eq. (6), we know that the $w$MSE component will have no effects on optimization in knowledge distillation. This issue is caused by the asymmetric property of KL loss. Thus we propose to break the asymmetric gradient property by enabling the gradients of $\Delta \mathbf{n}$ in Eq. (5) and Eq. (6). As a result, we derive the Eq. (7).

Q5: In Table 1, the GI seems more important than BA; could you please explain it and add the results for only the use of GI?

Actually, GI and BA work together for adversarial training. Adversarially-trained models only achieve around 60% natural accuracy on CIFAR-100. The low performance can lead to heavy sample-wise biases for $w$MSE during training. GI can properly address this problem. Additionally, we observe that, without the BA, the model with GI is easy to crush during adversarial training. BA mechanism can make the training stable and efficient.

Q6: Regarding the asymmetric issue, what are the advantages of IKL loss over JSD loss?

That's a good question. With the following JS divergence loss,

$\quad \quad \quad JSD(P||Q) = \frac{1}{2} KL(P|| M) + \frac{1}{2} KL(Q||M), \quad M=\frac{1}{2} P + \frac{1}{2} Q.$

Suppose $P$ and $Q$ are probability distributions of the teacher model and the student model respectively. We calculate its derivatives regarding $\mathbf{o}_{n}$ (the student logits),

$\quad \quad \quad \frac{\partial \mathcal{L}_{JSD}}{\partial \mathbf{o}\_{n}^{i}}=\sum\_{j=1}^{C} \mathbf{w}\_{\mathbf{n}}^{i,j}(\Delta n\_{i,j} - \Delta \mathbf{m}^{'}\_{i,j})$, with the constraint $\quad Softmax(\mathbf{o}\_{m^{'}}) = \frac{1}{2} \mathbf{s}\_{n} + \frac{1}{2} \mathbf{s}\_{m}$,

where $\mathbf{o}\_{m}$ is the logits from the teacher model; $\mathbf{o}\_{m^{'}}$ is a virtual logits satisfying the above constraint; $\mathbf{s}\_{m} = *Softmax*(\mathbf{o}\_{m})$; $\mathbf{s}\_{n}= *Softmax*(\mathbf{o}\_{n})$; $\Delta \mathbf{m}^{'}\_{i,j} = \mathbf{o}\_{m^{'}}^{i} - \mathbf{o}\_{m^{'}}^{j}$; $\Delta \mathbf{n}\_{i,j} = \mathbf{o}\_{n}^{i} - \mathbf{o}\_{n}^{j}$; $\mathbf{w}\_{n}^{i,j}=\mathbf{s}\_{n}^{i} * \mathbf{s}\_{n}^{j}$.

Correspondingly, the derivatives of IKL loss regarding $\mathbf{o}_{n}$ (the student logits),

$\quad \quad \frac{\partial \mathcal{L}\_{IKL}}{\partial \mathbf{o}\_{n}^{i}} = \underbrace{\alpha \sum\_{j=1}^{C} \mathbf{\bar w}\_{m}^{i,j}(\Delta \mathbf{n}\_{i,j} - \Delta \mathbf{m}\_{i,j})}\_{Effects \ \ of \ \ wMSE \ \ loss} + \underbrace{\beta \left( \mathbf{s}\_{m}^{i} * (\mathbf{s}\_{n}^{i} - 1) + \mathbf{s}\_{n}^{i} * (1 - \mathbf{s}\_{m}^{i}) \right)}\_{Effects \ \ of \ \ Cross-Entropy \ \ loss}$

Although JSD can produce gradients on $\mathbf{o}\_{n}$, it has lost the property of KL loss that forces $\mathbf{o}\_{n}$ and $\mathbf{o}\_{m}$ to be similar absolutely in the Euclidean space and relatively in the Softmax space.

First, the soft labels from the teacher models often embed dark knowledge and facilitate the optimization of the student models. However, there are no effects of the cross-entropy loss with the soft labels for JSD loss.

Second, there could be multiple virtual $\mathbf{o}^{'}\_{m}$ satisfying the constraint ($Softmax(\mathbf{o}\_{m^{'}}) = \frac{1}{2} \mathbf{s}\_{n} + \frac{1}{2} \mathbf{s}\_{m}$). Thus, JSD loss is still to make $\mathbf{o}\_{m}$ and $\mathbf{o}\_{n}$ to be similar relatively in the softmax space and suffers from the same problem of KL divergence in scenarios like knowledge distillation.

We also empirically validate that IKL loss achieves better performance on the knowledge distillation task. As shown in Table 33 in the rebuttal_tables.pdf, the model trained with IKL loss surpasses the models trained by KL and JSD losses by 1.44% and 2.03% top-1 accuracy respectively.

Thanks for your suggestions and valuable comments. Here we provide our responses to address your concerns.

Q1: It should be explicitly clarified that Theorem 1 and the subsequent analysis are solely based on the assumption that the probability distribution is a categorical distribution.

Thanks for your suggestions. We will explicitly clarify that our analysis is based on the assumption that the probability distribution is a categorical distribution in the new revision.

Q2: The connection between KL loss and related tasks.

Thanks for the suggestion. We have described the connections between KL loss and the adversarial training and knowledge distillation in Sec. 3.1. For clarification, We will re-organize the section for related work and Section 3.1 and include other related applications of the KL divergence loss in the new revision.

Moreover, our paper focuses on the applications of the KL divergence loss from the gradient optimization perspective rather than the measurement between distributions. We delve into how KL divergence loss optimizes models during training in terms of gradient optimization.

Q3: New applications of our IKL loss on semi-supervised learning, knowledge distillation on imbalanced data, and semantic segmentation knowledge distillation.

Besides the adversarial training task and knowledge distillation on balanced data, We conduct experiments on other tasks including semi-supervised learning, knowledge distillation on imbalanced data, and semantic segmentation knowledge distillation. Experimental results show that our IKL loss can significantly improve model performance when compared with KL loss. On semi-supervised learning, replacing the consistency loss with our IKL loss in the Mean-Teacher algorithm, we achieve 2.67% improvements. For the knowledge distillation on imbalanced data, replacing the KL loss with our IKL loss, our model outperforms the baseline model by 1.44% top-1 accuracy. For the semantic segmentation distillation, replacing the KL loss with our IKL loss, we achieve 0.5 mIoU gains on the ADE20K dataset.

For details, please refer to our global rebuttal Q1.

Q4: Why do we need to break the asymmetric property of KL loss in scenarios like knowledge distillation?

Please refer to our global rebuttal Q3.

Q5: Regarding the asymmetric problem, what is the relation between the proposed DKL and the Jensen-Shannon (JS) divergence?

That's a good question. With the following JS divergence loss,

$\quad \quad \quad JSD(P||Q) = \frac{1}{2} KL(P|| M) + \frac{1}{2} KL(Q||M), \quad M=\frac{1}{2} P + \frac{1}{2} Q.$

Suppose $P$ and $Q$ are probability distributions of the teacher model and the student model respectively. We calculate its derivatives regarding $\mathbf{o}_{n}$ (the student logits),

$\quad \quad \quad \frac{\partial \mathcal{L}_{JSD}}{\partial \mathbf{o}\_{n}^{i}}=\sum\_{j=1}^{C} \mathbf{w}\_{\mathbf{n}}^{i,j}(\Delta n\_{i,j} - \Delta \mathbf{m}^{'}\_{i,j})$, with the constraint $\quad Softmax(\mathbf{o}\_{m^{'}}) = \frac{1}{2} \mathbf{s}\_{n} + \frac{1}{2} \mathbf{s}\_{m}$,

where $\mathbf{o}\_{m}$ is the logits from the teacher model, $\mathbf{o}\_{m^{'}}$ is a virtual logits satisfying the constraint, $\mathbf{s}\_{m} = *Softmax*(\mathbf{o}\_{m})$, $\mathbf{s}\_{n}= *Softmax*(\mathbf{o}\_{n})$, $\Delta \mathbf{m}^{'}\_{i,j} = \mathbf{o}\_{m^{'}}^{i} - \mathbf{o}\_{m^{'}}^{j}$, $\Delta \mathbf{n}\_{i,j} = \mathbf{o}\_{n}^{i} - \mathbf{o}\_{n}^{j}$, $\mathbf{w}\_{n}^{i,j}=\mathbf{s}\_{n}^{i} * \mathbf{s}\_{n}^{j}$.

Correspondingly, the derivatives of IKL loss regarding $\mathbf{o}_{n}$ (the student logits),

$\quad \quad \frac{\partial \mathcal{L}\_{IKL}}{\partial \mathbf{o}\_{n}^{i}} = \underbrace{\alpha \sum\_{j=1}^{C} \mathbf{\bar w}\_{m}^{i,j}(\Delta \mathbf{n}\_{i,j} - \Delta \mathbf{m}\_{i,j})}\_{Effects \ \ of \ \ wMSE \ \ loss} + \underbrace{\beta \left( \mathbf{s}\_{m}^{i} * (\mathbf{s}\_{n}^{i} - 1) + \mathbf{s}\_{n}^{i} * (1 - \mathbf{s}\_{m}^{i}) \right)}\_{Effects \ \ of \ \ Cross-Entropy \ \ loss}$

Although JSD can produce gradients on $\mathbf{o}\_{n}$, it has lost the property of KL loss that forces $\mathbf{o}\_{n}$ and $\mathbf{o}\_{m}$ to be similar absolutely in the Euclidean space and relatively in the Softmax space.

First, the soft labels from the teacher models often embed dark knowledge and facilitate the optimization of the student models. However, there are no effects of the cross-entropy loss with the soft labels for JSD loss.

Second, there could be multiple virtual $\mathbf{o}^{'}\_{m}$ satisfying the above constraint ($Softmax(\mathbf{o}\_{m^{'}}) = \frac{1}{2} \mathbf{s}\_{n} + \frac{1}{2} \mathbf{s}\_{m}$). Thus, JSD loss is still to make $\mathbf{o}\_{m}$ and $\mathbf{o}\_{n}$ to be similar relatively in the softmax space and suffers from the same problem of KL divergence in scenarios like knowledge distillation.

We also empirically validate that IKL loss achieves better performance on the knowledge distillation task. As shown in Table 33 in the rebuttal_tables.pdf, the model trained with IKL loss surpasses the models trained by KL and JSD losses by 1.44% and 2.03% top-1 accuracy respectively.

Thanks for your suggestions and valuable comments. Here we provide our responses to address your concerns.

Q1: Theoretical evaluation is not convincing enough. There are no further justifications for why the specific anti-derivative formulation expressed in Theorem 1 is chosen as the main decoupled KL (DKL) formulation, as there could be multiple anti-derivative formulations for the same function that also lead to equivalent derivative formulations.

We decouple the KL divergence loss into a $w$MSE component and a Cross-entropy component because MSE and Cross-entropy loss are the two most popular objectives for image recognition. There are lots of loss functions that can be used for image classification, like Mean Absolute Error and Hinge loss. However, Cross-entropy loss is the most effective one for gradient optimization. Similarly, MSE loss has also been validated as one of the most effective loss functions in many tasks like object detection and other regression tasks.

With Theorem 1, we know that the KL divergence loss can be decoupled into a $w$MSE and a Cross-entropy loss. Then we can examine how each component affects model performance, which cannot be done by directly operating on KL divergence loss.

Q2: While the authors have intuitively highlighted the issues of the unbalanced and potentially unstable optimization issues with DKL, the authors did not delve deep into the aforementioned issues. Further theoretical proofs or gradient analysis of KL/DKL/IKL optimization would provide a stronger justification, especially where exactly the stop-gradient operations are crucial.

The stop-gradient operation is just used to ensure that our DKL loss (Theorem 1) is equivalent to the KL divergence loss from the perspective of gradient optimization, which implies that KL loss is a special case of our DKL/IKL loss.

By examining each component of DKL loss, we identify the potential issues of DKL loss, i.e., the asymmetric gradient property and possible sample-wise biases from hard examples or outliers. Since the DKL loss is equivalent to the KL loss, we consider that the KL loss also suffers from the two problems. Then we propose the IKL loss to address the potential issues.

Q5: We achieve significant improvements over the baseline and state-of-the-art adversarial robustness on public leaderboard --- RobustBench.

First, with the same experimental settings and only replacing the KL loss with the IKL loss, we significantly outperform our baseline TRADES by 1.63% robustness under auto-attack.

Second, without extra techniques from previous algorithms, we achieve the state-of-the-art adversarial robustness on the public leaderboard --- RobustBench under both settings: with the basic augmentation strategy including random crop and horizontal flip, with advanced data augmentation or synthesized data.

We only replace the KL loss with our IKL loss using the TRADES pipeline and thus our method is much more computationally efficient than LBGAT and ACAT, saving 33.3% training time. Our results are summarized in Table 35 in the rebuttal_tables.pdf.

Q6: Does the assumption of DKL/KL equivalence in Theorem 1 still hold when the parameters $\alpha=1$ and $\beta=1$ are not set to 1?

DKL loss is equivalent to KL loss only when $\alpha=1$ and $\beta=1$. The significance of our Theorem 1 is that: the KL loss can be decoupled into a weighted MSE loss and a Cross-entropy loss. Then the two terms can be adjusted independently for efficient optimization, which can not be achieved by directly tuning hyper-parameters of KL loss.

Q7: Various recent works have highlighted the importance of solving the gradient conflict issue for the accuracy/robustness trade-off in various settings. Some works even proposed that anti-symmetric gradients are essential to solving gradient conflicts. How can the authors justify their approach compared to these opposing approaches?

That's an interesting question. First, it is how the gradients affect training optimization matters rather than the property of symmetric or asymmetric. In our case, the asymmetric gradient property of KL divergence loss causes the $w$MSE loss not to work under scenarios like knowledge distillation, leading to optimization problems during training (please refer to our global rebuttal Q3). Thus, we break the asymmetric property and observe performance improvements on several tasks including KD, AT, and semi-supervised learning (please refer to our global rebuttal Q1).

Second, as mentioned by you, some work like [ref4] proposes that anti-symmetric gradients can solve gradient conflicts between invariance loss and classification loss, and promote model robustness. In the paper, the key is still to remove the harmful gradients. The asymmetric gradient optimization is just one of the ways to solve gradient conflicts in this case. Thus, we should always pay attention to how the gradients affect the optimization process rather than the symmetric or asymmetric properties of gradients. I think it is a promising direction for deep learning optimization not only in robustness. I will include the relevant discussion in the new revision. The mentioned works on stop gradient for unsupervised representation learning will also be discussed.

[ref4] Rethinking Invariance Regularization in Adversarial Training to Improve Robustness-Accuracy Trade-off. arXiv:2402.14648.

[ref5] How Does SimSiam Avoid Collapse Without Negative Samples? A Unified Understanding with Self-supervised Contrastive Learning. ICLR 2022.

[ref6] Towards a Unified Theoretical Understanding of Non-contrastive Learning via Rank Differential Mechanism. ICLR 2023.

[ref7] The Mechanism of Prediction Head in Non-contrastive Self-supervised Learning. NeurIPS 2022.

[ref8] Gradient surgery for multi-task learning. NeurIPS 2020.

Thanks for your suggestions and valuable comments. Here we provide our responses to address your concerns.

Q1: We achieve significant improvements over the baseline and state-of-the-art adversarial robustness on RobustBench

For adversarial robustness

First, with the same experimental settings and only replacing the KL loss with the IKL loss, we significantly outperform our baseline TRADES by 1.63% robustness under auto-attack.

Second, without extra techniques from previous algorithms, we achieve the state-of-the-art adversarial robustness on the public leaderboard --- RobustBench under both settings: with the basic augmentation strategy including random crop and horizontal flip, with advanced data augmentation or synthesized data.

We only replace the KL loss with our IKL loss using the TRADES pipeline and thus our method is much more computationally efficient than LBGAT and ACAT, saving 33.3% training time. Our results are summarized in Table 35 in the rebuttal_tables.pdf

For knowledge distillation

On the knowledge distillation, for the mentioned work "multi-level logit distillation`` [ref3], we observe that it uses a weak augmentation (RandomResizedCrop and random horizontal flip) data batch and a strong augmentation (Autoaug) data batch for each training iteration. For fair comparisons, based on their open-sourced code (https://github.com/Jin-Ying/Multi-Level-Logit-Distillation), we also apply the same pipeline to our IKL loss. The experimental results are summarized in Table 36 in the rebuttal_tables.pdf. Under the same experimental setting, we achieve better performance, outperforming the work [ref3] by 0.35% (ResNet-34--ResNet-18) and 0.33% (ResNet-50--MobileNet) on the ImageNet.

[ref3] Multi-level Logit Distillation. CVPR 2023.

Q2: About the references, terms, and Figure 1.

Thanks for your suggestions. We would add citations properly, revise Figure 1 for a clear explanation, and define the terms systematically in the new version.

Q3: The term "Inserting Class-wise Global Information" used for the weight term replacement in Eq.(5) is misleading and lacks a clear explanation.

In Eq (5), we don't discuss the global class-wise information. We first propose Theorem 1 and identify that the DKL loss can suffer from sample-wise biases because the $w$MSE component depends on the sample-wise prediction scores. Especially in adversarial training, the natural accuracy is only around 60% on CIFAR-100. There can be many hard examples or outliers that cause sample biases. The KL loss is equivalent to DKL loss in terms of gradient optimization. Thus the KL loss also suffers from this problem. Then we address this potential issue in Section 3.3 by introducing the global class-wise information. For more details, please refer to our analysis in Section 3.4.

Q4: The proof approach for decoupling KL divergence is not straightforward. Instead of deriving DKL from KL divergence, the authors present DKL first and then show its equivalence through gradients, which obscures the understanding of the derivation process.

In Section 3.2, we first analyze the gradient information of KL divergence loss regarding $o_{m}$ and $o_{n}$. Then based on the structured gradient information, we construct the DKL formulation and propose Theorem 1. Finally, we give our proof for the equivalence between KL divergence loss and our DKL loss in the Appendix.

Q5: The benefits and reasons for why the term in Eq. (7) should receive gradients are not well explained.

Please refer to our global rebuttal Q3.

Thanks for your suggestions and valuable comments. Here we provide our responses to address your concerns.

Q1: Besides the adversarial training and the knowledge distillation on balanced data, how does the IKL loss perform on other tasks, like semi-supervised learning, knowledge distillation on imbalanced data, and semantic segmentation knowledge distillation?

Thanks for your suggestions. Besides the adversarial training task and knowledge distillation on balanced data, We conduct experiments on other tasks including semi-supervised learning, knowledge distillation on imbalanced data, and semantic segmentation knowledge distillation. Experimental results show that our IKL loss can significantly improve model performance when compared with KL loss. For details, please refer to our global rebuttal Q1.

Q2: Pseudo code of IKL loss and complexity analysis.

Compared with the KL loss, IKL loss is only required to update the global class-wise prediction scores $W \in \mathbb{R}^{C \times C}$ where $C$ is the number of classes during training. This extra computational cost can be nearly ignored when compared with the model forward and backward. Algorithm 1 shows the implementation of our IKL loss in Pytorch style. On dense prediction tasks like semantic segmentation, $\Delta_{a}$ and $\Delta_{b}$ can require large GPU Memory due to the large number of pixels. Here, we also provide the memory-efficient implementations for $w$MSE loss component, which is listed in Algorithm 2.

Algorithm 1: Pseudocode for DKL/IKL loss in Pytorch style

Input: $logits_{a}, logits_{b} \in \mathbb{R}^{B \times C}$; one-hot label $Y$; $W \in \mathbb{R}^{C \times C}$; $\alpha$ and $\beta$.

class_scores = $Y$ @ $W$;

Sample_weights = class_scores.view(-1, C, 1) @ class_scores.view(-1, 1, C);

$\Delta_a$ = $logits_a$.view(-1, C, 1) - $logits_a$.view(-1, 1, C);

$\Delta_b$ = $logits_b$.view(-1, C, 1) - $logits_b$.view(-1, 1, C);

wMSE_loss = (torch.pow($\Delta_{n}$ - $\Delta_{a}$, 2) * Sample_weights).sum(dim=(1,2)).mean() * $\frac{1}{4}$;

score_a = F.softmax($logits_a$, dim=1).detach();

log_score_b = F.log_softmax($logits_b$, dim=-1);

CE_loss = -(score_a * log_score_b).sum(1).mean();

$**return**$ $\beta$ * CE_loss + $\alpha$ * wMSE_loss.

Algorithm 2: Memory efficient implementation for wMSE in Pytorch style.

Input: $logits_{a}, logits_{b} \in \mathbb{R}^{B \times C}$; one-hot label $Y$; $W \in \mathbb{R}^{C \times C}$;

class_scores = $Y$ @ $W$;

loss_a = (class_scores * $logits_a$ * $logits_a$).sum(dim=1) * 2 - torch.pow((class_scores * $logits_a$).sum(dim=1), 2) * 2;

loss_b = (class_scores * $logits_b$ * $logits_b$).sum(dim=1) * 2 - torch.pow((class_scores * $logits_b$).sum(dim=1), 2) * 2;

loss_ex = (class_scores * $logits_a$ * $logits_b$).sum(dim=1) * 4 - (class_scores * $logits_a$).sum(dim=1) * (class_scores * $logits_b$).sum(dim=1) * 4;

wMSE_loss = $\frac{1}{4}$ * (loss_a + loss_b - loss_ex).mean();

$**return**$ wMSE_loss.

Q3: Can the authors provide more insight into the selection of the hyperparameters $\alpha$ and $\beta$ in the IKL loss, and how sensitive are the results to these values?

For tasks using KL divergence loss, like knowledge distillation, adversarial training, and semi-supervised learning, we have priors to set the loss weight $\gamma$ for KL divergence loss. When replacing KL divergence loss with our IKL loss, we first set the $\beta$ =$\gamma$ for the cross-entropy component. Then, we perform the grid search $\frac{\alpha}{4}$ $\in$ {1, 2,3, 4,5}. After determining the value of $\frac{\alpha}{4}$, we again adjust the $\beta$ to achieve the best performance on validation data if needed. As shown in Table 15 of the Appendix, in a reasonable range $\frac{\alpha}{4} \in$ {3,4,5,6} and $\beta \in$ {2,3,4,5}, increasing $\frac{\alpha}{4}$ and $\beta$, the performance firstly increases steadily and then decreases.

Q4: How does the class-wise global information interact with models that are trained in a semi-supervised or unsupervised setting?

In semi-supervised or unsupervised settings, we often can get the pseudo labels for the data. Equipping with a threshold $\gamma$, we filter out the samples that the maximum confidence score is lower than $\gamma$, like in FixMatch [ref2]. Then, the remaining samples can be used to update the global class-wise prediction scores.

[ref2] FixMatch: Simplifying Semi-Supervised Learning with Consistency and Confidence. NeurIPS 2020.

Q5: How does the IKL loss handle adversarial examples that are crafted to be more sophisticated than those generated by the Auto-Attack method used in the paper?

The auto-attack is one of the strongest attack methods. It is an ensemble of several adversarial attacks, including APGD-CE, APGD-DLR, FAB, and Square Attack. Here we also test the robustness of our models under PGD, and CW attacks with 10 and 20 iterations. The perturbation size and step size are set to 8/255 and 2/255 respectively. As shown in the Table, with increasing iterations from 10 to 20, our models show similar robustness, demonstrating that our models don't suffer from the obfuscated gradients problem.

The worst-case is in bold.