Knowledge distillation based adversarial training for robust multi-exit neural network

Abstract

While multi-exit neural networks are regarded as a promising solution for making efficient inference via early exits, combating adversarial attacks remains a challenging problem. In multi-exit networks, due to the high dependency among different submodels, an adversarial example targeting a specific exit not only degrades the performance of the target exit but also reduces the performance of all the other exits concurrently. This fundamental challenge makes multi-exit networks highly vulnerable to simple adversarial attacks. In this paper, we propose a knowledge distillation based adversarial training strategy tailored to multi-exit neural networks. The first component of our architectural solution, adversarial knowledge distillation, guides the output of the adversarial examples to mimic the output of the clean data. The second component, exit-wise orthogonal knowledge distillation, reduces the adversarial transferability across different submodels, significantly improving the robustness against adversarial attacks. Experimental results on various datasets/models show that our method achieves the best adversarial accuracy with reduced computation budgets, compared to other baselines relying on existing adversarial training or knowledge distillation techniques for multi-exit networks.