MixACM: Mixup-Based Robustness Transfer via Distillation of Activated Channel Maps

1 Introduction . Deep learning models have achieved impressive performance on a wide variety of challenging tasks such as image recognition , natural language generation , game playing , etc . However , these models are susceptible to even small changes in the input space . It is possible to craft tiny changes in the input such that a model classifies unaltered input correctly but classifies the same input incorrectly after small and visually imperceptible perturbations [ 71 , 27 ] . These altered inputs are known as adversarial examples , and this vulnerability of the state-of-the-art models has raised serious concerns [ 41 , 58 , 10 , 47 ] . Many defense mechanisms have been proposed to train deep models to be robust against such adversarial perturbations . These techniques include defensive distillation [ 57 ] , gradient regularization [ 29 , 58 , 62 ] , model compression [ 21 , 45 ] , activation pruning [ 23 , 59 ] , adversarial training [ 49 ] , etc . Among them , Adversarial Training ( AT ) is one general strategy that is the most effective [ 4 ] . ∗Equal Contribution . †This work was carried out at Huawei Noah ’ s Ark Lab . The webpage for the project is available at : awais- rauf.github.io/MixACM ‡Corresponding Author . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . In general , adversarial training is a kind of data augmentation technique that trains a model on examples augmented with adversarial changes [ 49 ] . The adversarial training consists of an inner , iterative maximization loop to augment natural examples with adversarial perturbations , and an outer minimization loop similar to normal training . Many different methods have been introduced to improve robustness [ 76 , 15 , 52 , 91 , 79 , 66 , 94 , 78 , 7 , 81 , 39 , 44 , 56 ] , but all of them are fundamentally based on the principle of training on adversarially augmented examples . Adversarial training is more challenging compared with normal training . Better robust generalization requires larger capacity models [ 49 , 53 , 80 ] . Even over-parameterized models that can easily fit data for normal training [ 92 ] may have insufficient capacity to fit adversarially augmented data [ 97 ] . The sample complexity of adversarial training can be significantly higher than normal training , and it requires more labeled [ 64 ] or unlabeled data [ 76 , 15 , 52 , 91 ] . The inner maximization for the generation of adversarial perturbations for adversarial training is significantly more expensive computationally as it requires iterative gradient steps with respect to the inputs ( e.g. , adversarial training takes ∼7x more time compared with normal training ) . Adversarial training also degrades the performance of a model on natural examples significantly [ 53 , 94 ] . Can we attain robustness more efficiently , i.e. , with less data , small capacity models , without extra back-propagation steps , and at no significant sacrifice of clean accuracy ? For normal training , knowledge transfer is one possible paradigm for the efficient training of deep neural networks . In normal knowledge transfer , a pre-trained teacher model is leveraged to efficiently train a student model on the same or similar dataset [ 70 ] . However , knowledge transfer methods are designed to transfer features related to normal generalization which may be at odds with robustness [ 74 , 38 ] . A recent line of work has explored transferring or distilling robustness from pre-trained models [ 32 , 67 , 25 , 16 ] . Although these techniques are effective , they also require extra gradient computation and may not work in some cases . In this work , we argue that a better approach for robustness transfer is to distill intermediate features of a robust teacher , generated on mixup examples . Our proposed approach does not require any additional gradient computation and works well with smaller models and fewer data samples . It can also effectively transfer robustness across models and datasets . We begin with the theoretical analysis and show that adversarial loss of a student model can be bounded with two terms : natural loss and distance between the student and teacher model on mixup examples . To minimize the distance between robust teacher and student , we proposed a new distillation method . Our proposed distillation method is based on channel-wise activation analysis of robustness by Bai et al . [ 7 ] and other studies showing that adversarially trained models learn fundamentally different features [ 38 , 74 , 99 ] . Channels of a convolutional neural network learn different disentangled representations which , when combined , describe specific semantic concepts [ 8 ] . Samples of different classes activate different channels for an intermediate layer . However , adversarial examples make these channels activate more uniformly thereby destroying class-related information . Adversarial training solves this issue by forcing a similar channel-activation pattern for normal and adversarial examples [ 7 ] . Our proposed distillation method generates activated channel maps from a robust teacher which has already learned robust channel-activation patterns . The student , then , is forced to match these activated channel maps . We have conducted extensive experiments to show the effectiveness of our method using various datasets and under different learning settings . We start by exhibiting the ability of our method to transfer robustness without requiring adversarial examples . We then show that our method can distill robustness from large pre-trained models to smaller models and it can transfer robustness across datasets . We also show that our method is capable to robustify a model even with a small number of examples . Concisely , our contributions are as follows : 1 . We show that adversarial loss of a student model can be bounded with the distance between robust teacher and student model on mixup examples . 2 . We propose a new method to transfer robustness from the intermediate features of robust pre-trained models . 3 . We have demonstrated effectiveness of our method with experiments on CIFAR-10 , CIFAR100 , and ImageNet datasets . We showed that MixACM can achieve adversarial robustness comparable to the state-of-the-art methods without generating adversarial examples . Our method also improves clean accuracy significantly . 2 Related Work . Adversarial Training and Robust Features . Adversarial attacks are considered to be a security concern [ 9 , 71 , 27 , 14 ] , and many methods have been proposed to defend against such attacks [ 29 , 27 , 57 , 30 , 58 , 73 , 49 , 12 , 43 , 62 , 21 , 45 , 23 , 59 ] . Adversarial training [ 49 ] is the most effective defense [ 4 ] among these methods . The fundamental principle of adversarial training is to generate adversarial perturbation during training by doing iterative back-propagation w.r.t . input . The model is then trained on these perturbed examples . Building upon this principle , different aspects of adversarial training have extensively been studied , e.g. , effect of large unlabeled data [ 76 , 15 , 91 ] , theoretical aspects to improve the trade-off between accuracy and robustness [ 94 ] , improving parts of adversarial training for better robustness [ 78 , 81 , 56 , 7 , 28 ] , making adversarial training fast [ 52 , 66 , 79 , 1 ] , pruning with adversarial training [ 87 , 65 ] , adversarial training with unlabeled data [ 39 ] , etc . Our work is different from these methods as our main concern is robustness transfer without generating adversarial perturbations . Another line of work studied the effects of adversarial training and showed that adversarial training learns ‘ fundamentally different ’ [ 53 ] features compared with normal training [ 38 , 86 , 100 ] . Our method is motivated by these studies and is designed to leverage these robust features . Knowledge and Robustness Transfer . Knowledge transfer is training a student model more efficiently on a dataset with the help of a teacher model trained on a similar or related dataset [ 70 ] . Many different settings have been explored to achieve this objective by using soft labels produced by teacher [ 11 , 35 ] , knowledge learned by intermediate layers of teacher also called feature distillation [ 61 , 89 , 37 , 88 , 75 , 40 , 68 , 34 , 33 , 72 ] , or matching input gradients [ 89 , 70 ] . However , the main objective of these knowledge transfer methods is improving the generalization on natural , unperturbed images . Our work is fundamentally different from these works as we want to transfer robustness . Our work is motivated by [ 32 , 67 , 25 , 16 , 6 ] showing transferability of robustness . Hendrycks et al . [ 32 ] showed that robust features learned on large datasets like ImageNet can improve adversarial training . Shafahi et al . [ 66 ] showed that robust pre-trained models can act as feature extractors in transfer learning . These two methods are limited in application as pre-trained models are used as backbone models . In [ 25 ] , the authors showed that the robustness can be distilled more efficiently from a large pre-trained teacher model to a smaller student model by using the teacher-produced class scores on adversarial examples . Compared with [ 25 ] , our method does not require any adversarial training . In [ 16 ] , the distillation is performed by matching the gradient of the teacher and student , requiring fine-tuning teacher on target dataset and back-propagation to get gradients of the teacher and student w.r.t . inputs and training of a discriminator . Compared with [ 16 ] , our proposed method does not require any back-propagation steps in addition to those carried out in normal training . Activation and Robustness . Many works have explored the role of activations in the robustness from different angles , such as modifying output of intermediate layers to improve adversarial training [ 82 , 23 , 48 , 51 , 84 ] , using different activation functions [ 59 , 77 , 85 , 28 , 55 , 7 , 5 ] or interpretation of what is learned by adversarial training with activations [ 86 , 7 ] . However , we leveraged activations to transfer robustness . Our distillation method is inspired by the channel-wise analysis of robustness proposed by Bai et al . [ 7 ] . Mixup Augmentation and Robustness . Mixup augmentation was introduced to improve the generalization however , it also improved performance against one-step adversarial attacks [ 95 ] . This motivated many mixup-based adversarial defense methods , e.g. , Pang et al . [ 54 ] , Lee et al . [ 42 ] , Archambault et al . [ 3 ] , Bunk et al . [ 13 ] , Chen et al . [ 17 ] . These defense methods primarily improve adversarial training by mixup . Our work , however , is different as we explore the role of the mixup in robustness transfer from both theoretical and algorithmic perspectives . Our theoretical work is based on [ 98 ] that showed a connection between mixup and robustness theoretically . Goldblum et al . [ 25 ] also briefly empirically explored the effect of some augmentation methods on robustness transfer . 3 Setup . We consider task of mapping input x ∈ X ⊆ Rd to label y ∈ Y = { 1 , 2 , . . . , K } . Given the training data D = { ( x1 , y1 ) , ... , ( xn , yn ) } , the goal is to learn a classifier f : X → Rk from a hypothesis space F . Let L ( f ( x ) , y ) be the loss function that measures that how poorly the classifier f ( x ) predicts the label y . Suppose f can be decomposed into g ◦ h , where h ∈ H is a feature extractor and g ∈ G stands for the classifier on top of feature extractor . Adversarial Augmentation : Adversarial training augments natural examples with adversarial perturbations which are attained by maximizing the loss , i.e. , x̃ = x+δ where δ = arg max‖δ′‖p≤ L ( f ( x+ δ′ ) , y ) . Here ‖ · ‖p stands for Lp norm . Throughout this paper , we take p =∞ . We denote the adversarial loss [ 49 ] as : L̃ ( f , D ) = 1 n n∑ i=1 max ‖δ‖∞≤ L ( f ( xi + δ ) , yi ) , where is the perturbation budget that governs the adversarial robustness of the model . Mixup Augmentation : Mixup constructs virtual examples by linearly combining two examples [ 54 ] : x̃ij ( λ ) = λxi + ( 1− λ ) xj , and ỹij ( λ ) = λyi + ( 1− λ ) yj , where λ ∈ [ 0 , 1 ] follows the distribution Pλ . The mixup loss is Lmix ( f , D ) = 1 n2 n∑ i=1 n∑ j=1 Eλ∼Pλ [ L ( f ( x̃ij ( λ ) ) , ỹij ( λ ) ) ] .

Summary: This paper proposes a robustness transfer method that combines mixup data augmentation, distillation from both prediction and intermediate features.   It is motivated by the theoretical insight that robust test error can be decomposed into adversarial loss plus distillation loss, and both of them can be upper bounded by mixup loss. Motivated by some other previous work, the KL divergence with the teacher model’s prediction and L2 distance with the teacher model’s intermediate activated feature maps are also added to the final loss of the student model. So the loss is composed of three terms: mixup loss, KL divergence between predictions, and L2 distance between the intermediate activation maps.   Experiments on popular datasets such as CIFAR-10, CIFAR-100 and ImageNet demonstrate it’s superiority compared with a variety of baselines.