[论文解读] How and When Adversarial Robustness Transfers in Knowledge Distillation?

Rulin Shao, Jinfeng Yi+2 more|arXiv (Cornell University)|Oct 22, 2021

Adversarial Robustness in Machine Learning参考文献 30被引用 5

一句话总结

本文提出 KDIGA（基于输入梯度对齐的知识蒸馏）方法，以将对抗鲁棒性从教师模型迁移至学生模型，而标准知识蒸馏无法保留这种鲁棒性。通过在蒸馏过程中对齐学生与教师网络的输入梯度，KDIGA 确保学生模型的认证鲁棒性不低于教师模型，在 CIFAR-10 上实证性能提升高达 17.81%，在 ImageNet 上提升高达 37.5%，同时保持或提升干净样本准确率。

ABSTRACT

Knowledge distillation (KD) has been widely used in teacher-student training, with applications to model compression in resource-constrained deep learning. Current works mainly focus on preserving the accuracy of the teacher model. However, other important model properties, such as adversarial robustness, can be lost during distillation. This paper studies how and when the adversarial robustness can be transferred from a teacher model to a student model in KD. We show that standard KD training fails to preserve adversarial robustness, and we propose KD with input gradient alignment (KDIGA) for remedy. Under certain assumptions, we prove that the student model using our proposed KDIGA can achieve at least the same certified robustness as the teacher model. Our experiments of KD contain a diverse set of teacher and student models with varying network architectures and sizes evaluated on ImageNet and CIFAR-10 datasets, including residual neural networks (ResNets) and vision transformers (ViTs). Our comprehensive analysis shows several novel insights that (1) With KDIGA, students can preserve or even exceed the adversarial robustness of the teacher model, even when their models have fundamentally different architectures; (2) KDIGA enables robustness to transfer to pre-trained students, such as KD from an adversarially trained ResNet to a pre-trained ViT, without loss of clean accuracy; and (3) Our derived local linearity bounds for characterizing adversarial robustness in KD are consistent with the empirical results.

研究动机与目标

• 探究为何标准知识蒸馏无法在学生模型中保留对抗鲁棒性。
• 开发一种方法，使对抗鲁棒性能从教师模型成功迁移至学生模型，即使在不同架构之间亦可实现。
• 在温和假设下，理论上证明使用 KDIGA 训练的学生模型的认证鲁棒性不低于教师模型。
• 分析局部线性度在知识蒸馏过程中对抗鲁棒性迁移的作用。
• 在多种模型（包括 ResNets 和视觉Transformer）上验证该方法在 ImageNet 和 CIFAR-10 上的性能。

提出的方法

• 提出 KDIGA，一种知识蒸馏框架，通过最小化学生与教师模型之间输入梯度差异的 L2 范数来实现训练。
• 修改标准 KD 损失，引入梯度对齐正则化项，促使学生模型模仿教师模型的输入梯度方向。
• 基于局部线性度量（LLM）、交叉熵损失和梯度对齐范数，推导认证鲁棒性的理论边界。
• 将该方法应用于从对抗训练的 ResNets 蒸馏至小型模型（如 MobileNetV2）以及预训练的 ViTs。
• 将 KDIGA 与对抗训练（ARD）结合，进一步提升鲁棒性，实现更低的理论边界和更优的实证性能。
• 使用不同扰动半径的 PGD 攻击，在 CIFAR-10 和 ImageNet 上评估鲁棒准确率。

实验结果

研究问题

• RQ1 为何标准知识蒸馏无法在教师到学生模型之间保留对抗鲁棒性？
• RQ2 在何种条件下，可成功实现知识蒸馏中从教师到学生模型的对抗鲁棒性迁移？
• RQ3 能否在根本不同的架构之间（如 ResNets 与视觉Transformer）实现对抗鲁棒性迁移？
• RQ4 KDIGA 中的输入梯度对齐如何提升学生模型的局部线性度与鲁棒性？
• RQ5 KDIGA 是否可应用于预训练的学生模型，而不会降低干净样本准确率，同时提升鲁棒性？

主要发现

• 当从对抗训练的 WideResNet 蒸馏至 MobileNetV2 学生模型时，KDIGA 将其在 CIFAR-10 上的鲁棒准确率从 5.97% 提升至 17.81%，且干净准确率高于标准 KD。
• 在 ImageNet 上，KDIGA 将学生模型的鲁棒准确率从标准 KD 的 1.5% 提升至 37.5%，显著优于基线方法。
• KDIGA 实现了从对抗训练的 ResNet 到预训练视觉Transformer的鲁棒性迁移，使 ViT 的鲁棒准确率提升 11.1 倍，且未牺牲干净准确率。
• 经 KDIGA 蒸馏的学生模型，其鲁棒性达到或超过教师模型，即使学生模型架构不同（如从 ResNet 到 MobileNetV2）。
• 基于局部线性度量的理论边界与实证鲁棒性趋势高度吻合，KDIGA 与 ARD 的理论边界均低于标准 KD。
• 将 KDIGA 与对抗蒸馏（ARD）结合后，理论鲁棒性边界进一步降低，在 CIFAR-10 上 ε=4/255 时达到最低边界 1.081。