Summary
In this paper, the authors analyzed the optimization gradient of the commonly used KL Divergence loss metric in adversarial training and distillation. KL loss can be reformulated as a Decoupled KL (DKL) Divergence loss term through antiderivative operations. The gradients of both terms are equivalent in the case where the parameters of the DKL are set to $\alpha$ = 1 and $\beta$ = 1. The DKL loss term consists of 1) a weighted MSE (wMSE) loss term, which encourages similarity between output logits, and 2) a CE loss term, which encourages accurate predictions. The authors highlighted two potential drawbacks of the KL/DKL formulations : 1) asymmetric optimization of both loss terms, 2) sample-wise conditioning term $w_m$ which induces strong variance into the optimization process. The authors proposed two modifications to the DKL loss term, namely, 1) Enabling the gradients of $\Delta n$ instead of $\Delta m$ in the formulation, thereby overcoming asymmetric optimization, 2) using global class-wise conditioning term $\bar{w_y}$, conditioned on global ground truth, instead of $w_m$. The novel Improved KL Divergence Loss (IKL) is shown to be quantitatively and qualitatively superior to the classic and decoupled KL divergence terms. The authors also performed thorough adversarial robustness and distillation experiments, achieving improved performance compared to accuracy and robustness performance in various natural image datasets ( CIFAR-10, CIFAR-100, ImageNet ), augmentation strategies , and network backbones, outperforming multiple benchmarks.
Strengths
1. Experiments : Thorough experiments were performed, incorporating various datasets, network backbones, and augmentation strategies.
2. Clarity : The intuition behind the main concept is clear and straightforward. Readers can easily follow the thought process and justifications of the authors.
3. Originality : The paper shows the promising potential of IKL in two different application domains, thus highlighting the IKL metric's relative flexibility. Furthermore, the hyperparameter ablations suggest that one can tune the predictive accuracy and adversarial robustness tradeoff. The competitive Robustbench results further substantiate the authors' claims.
Weaknesses
1. 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. The mathematical proof provided in the appendix has merely shown that performing chain rules on KL/DKL loss terms justifies the equivalence of the derivatives.
2. 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 following papers have performed gradient analysis and a thorough study into stop-gradient mechanisms.
- Zhang, C., Zhang, K., Zhang, C., Pham, T. X., Yoo, C. D.,
and Kweon, I. S. How does simsiam avoid collapse without negative samples? a unified understanding with selfsupervised contrastive learning. International conference on learning representations, 2022. https://arxiv.org/abs/2203.16262
- Zhuo, Z., Wang, Y., Ma, J., and Wang, Y. Towards a unified theoretical understanding of non-contrastive learning via rank differential mechanism. International conference on learning representations, 2023.
https://arxiv.org/abs/2303.02387
3. The accuracy and robustness results are generally only marginally better than the benchmarks. t-SNE visualization of the purported benefits of the IKL formulation does not additionally convey a strong qualitative advantage. The authors highlighted the lower computational costs of the approach in some of their experiments. I suggest that the authors pivot more on this to demonstrate that the IKL is a competitive, relatively computationally lighter method.
Minor : While the structure of the paper is clear, there are multiple grammatical and expression errors. I would recommend stronger proofreading and reformulation for improvement.
Questions
1. Does the assumption of DKL/KL equivalence in Theorem 1 still hold when the parameters $\alpha$ and $\beta$ are not set to 1? The authors have experimented with the values of these parameters while assuming that the DKL/IKL formulation still results in equivalent gradient formulations.
2. While various recent works have highlighted the importance of solving the gradient conflict issue for accuracy/robustness tradeoff in various settings. Some works even proposed that anti-symmetric gradients are essential to solving the gradient conflicts. How can the authors justify their approach compared to these opposing approaches/perspectives?
- Zhang, C., Zhang, K., Zhang, C., Pham, T. X., Yoo, C. D.,
and Kweon, I. S. How does simsiam avoid collapse without negative samples? a unified understanding with selfsupervised contrastive learning. International conference on learning representations, 2022. https://arxiv.org/abs/2203.16262
- Zhuo, Z., Wang, Y., Ma, J., and Wang, Y. Towards a unified theoretical understanding of non-contrastive learning via rank differential mechanism. International conference on learning representations, 2023.
https://arxiv.org/abs/2303.02387
- Yu, T., Kumar, S., Gupta, A., Levine, S., Hausman, K., and Finn, C. Gradient surgery for multi-task learning. Advances in Neural Information Processing Systems, 33: 5824–5836, 2020.
https://arxiv.org/abs/2001.06782
- Wen, Z. and Li, Y. The mechanism of prediction head in noncontrastive self-supervised learning. Advances in Neural
Information Processing Systems, 35:24794–24809, 2022.
https://arxiv.org/abs/2205.06226
- Waseda, F. et al. Rethinking Invariance Regularization in Adversarial Training to Improve Robustness-Accuracy Trade-off. https://arxiv.org/abs/2402.14648, 2024.
Limitations
The authors have not addressed the limitations of their work. Please provide some potential limitations that can be improved on for future works.