Improving Defensive Distillation using Teacher Assistant

KD was originally introduced as a training procedure that aims to transfer the knowledge learned by a larger network (Teacher) to another relatively smaller one (Student), while maintaining the same level of performance [13]. The motivation behind this technique is to reduce the computational complexity of some operations, or compression of large networks, such that devices with limited capacity (e.g. smartphones) could benefit from the achievements of deep learning. In KD, the student network is trained not with the original hard binary labels of the dataset, but instead with the soft labels taken from the output probability of the teacher network.

The technique of defensive distillation introduced by [11] heavily relies on the intuition behind knowledge distillation (KD) and the common methods of attacks. The intuition behind this procedure is that the knowledge acquired by the teacher network is not only embedded in the weight parameters, but also in the probability vector output by the network. Consider the MNIST digits dataset, where the images of digits 7 and 1 share similar spatial structures. The output of the teacher network thus contains relatively similar probabilities for these two digits. This extra entropy in the probability vector contains structural information about the data compared with the hard binary labels, and thus may be helpful when the student network is faced with adversarial samples that could be “unnatural”, or outside of the data distribution.

One important factor in the distillation procedure is the choice of distillation temperature ($T$). The effect of the value of $T$ on the resulting defensive distillation is perhaps clearer when we consider the method used for adversarial attacks. Most attacks assume that the adversary has access to the gradient of the network, and can thus estimate the model’s sensitivity to the input data (see section II). When the gradients are large, the model is more sensitive because simple perturbations of the input can lead to large changes of the model output, and thus is more prone to attacks. Therefore, a major goal of the defense against adversarial attacks could be to smooth the gradients and to stabilize the model around the input data in the sample space. In this sense, acquiring “softer” labels from the teacher network (corresponding to larger $T$) should produce more robust student networks, as confirmed by the results from [11].

The authors in [11] found that using defensive distillation reduces the success rate of adversarial sample crafting from 95.89% to 0.45% on the MNIST dataset [14], and from 87.89% to 5.11% on the CIFAR10 dataset [15].

Knowing that defensive distillation could produce student networks that are more robust to adversarial attacks, a natural question one might ask is, what if we apply distillation multiple times? Another motivation is the observation of the teacher-student gap, where the performance of the student is often far than ideal when there is a large difference between the capacity of the student and the teacher.

[12] explored this multi-step distillation strategy in normal KD settings with the task of image recognition. Instead of training one student directly from the teacher, this method employs a teacher assistant, a network that has an architecture of intermediate size between the teacher and the student, and serves as an intermediate step in the distillation process. In other words, the TA is first trained with the soft labels from the teacher, and then produces another set of soft labels to train the student. The authors observed that student models trained with TAs outperform those trained directly with the teacher. They also discuss the choice of TA size and the distillation path for multi-step distillation beyond one TA. They observed that the more TAs employed in between distillations, the better the student.

In this work, we combine TAKD with defensive distillation, and explore whether adding intermediate steps in the distillation process can lead to even larger improvement in the robustness of the student model.

Instead of optimizing for $\delta X$, the objective is optimized over a transformed variable $w$, s.t.

As, $-1\leq\tanh(w_{i})\leq 1$ implies $0\leq X_{i}+\delta X_{i}\leq 1$. If $X$ is the input sample, and $t$ is the chosen ‘wrong’ target class ($t\neq F(X)$), then the attack finds $w$ solving

Here $f$ is the best suited objective function proposed in the paper, and $\kappa$ controls the confidence of misclassification. It encourages the solver to find an adversarial sample $X^{*}$ which will be classified as class $t$ with high confidence. For the purpose of experiments in this project, $\kappa$ has been set to zero. To avoid getting stuck at a local minimum, gradient descent is run with multiple random starting points close to $X$ for a fixed number of iterations. The random points are sampled at random from a norm ball of radius $r$ centered at $X$, where $r$ is the closest adversarial sample found so far.

The $L_{0}$ metric is non-differentiable and therefore standard gradient descent cannot be applied. The $L_{0}$ attack algorithm uses the $L_{2}$ adversary to eliminate pixels that are unimportant. Starting with the original image $X$, let $\delta X$ be the solution found by $L_{2}$ adversary, such that $X^{=}X+\delta X$ is an adversarial sample. The gradient of the objective function is computed ($g=\nabla f(X*)$). The pixel with the minimum gradient ($i=\operatorname{arg\min}_{i}g_{i}\cdot\delta X_{i}$) is removed from the allowed set, as it will have the least impact on the output. On each iteration, pixels are eliminated and the adversary is restricted to modify only the allowed set. This process is repeated until the $L_{2}$ adversary fails. By process of elimination, a final subset of important pixels can be identified. The constant $c$ is initially set to a very low value, and $L_{2}$ attack is run using this value. Upon failure, the value is doubled and the attack is run again till $c$ exceeds a threshold, which is declared as a failure. At each iteration, the solution found in the previous one is used as the starting point for gradient descent, and as such, the attack algorithm is much more efficient than previous $L_{0}$ attacks in the literature.

The $L_{\infty}$ distance is not fully differentiable and the authors observed that the performance was poor when using the basic objective function. They also noticed that as the $\left\|\delta X\right\|_{\infty}$ penalizes the largest element, gradient descent gets stuck oscillating between two suboptimal options. To circumvent this, the $L_{\infty}$ term is replaced in the objective function with apenalty for any element of $\delta X$ that exceeds threshold $\tau$. This prevents the oscillation in gradient descent, as all the large values are penalized. The resulting objective function is

The threshold $\tau$ is initially set to $1$, and decreased in each iteration. If after an iteration, all of the components of $\delta X$ are less than $\tau$, it is reduced by a factor of $0.9$ and repeated. The constant $c$ is chosen similarly to the $L_{0}$ case. Initially set to a low value, and if the $L_{\infty}$ adversary fails, the value of $c$ is doubled until it exceeds a threshold, when it is declared as a failure. Similar to $L_{0}$ adversary, gradient descent at each iteration is performed with a “warm-start” for efficiency.

For training the DNN classifier $F$, the training algorithm minimizes the empirical risk:

So, to decrease the log-likelihood $\ell(F,X,Y(X))=-Y(X)\cdot\log F(X)$ of $F$ on $(X,Y(X))$, the optimizer adjusts $\theta_{F}$ so that $F(X)$ can get close to $Y(X)$. As, $Y(X)$ is a one-hot vector, the optimization problem can be written as

So, the optimizer effectively pushes the model to make overconfident predictions on $t(X)$, while pushing other probabilities to zero. In contrast, when the student network is trained using soft labels $F(X)$, the optimization problem is given by:

Using probabilities $F_{j}(X)$ ensures that the training algorithm constrains the output neurons $F_{j}^{d}(X)$ proportionally to their likelihood when updating $\theta_{F}$. This ensures that the model learns the relative likelihood among classes and is not forced into hard predictions, this increasing its generalizability. Introducing $F^{a}$ in between them further increases this trend. Ideally, with enough samples and training, $F^{a}$ would eventually converge to $F$, and consequently $F^{d}$ would converge to $F^{a}$, but empirically the model robustness is seen to increase with distillation.

The training procedure gives an intuition about the higher generalizability of the distilled models, but a sensitivity metric shows how crafting adversarial samples gets harder with increasing the distillation temperature. As discussed before, adversarial attacks identify inputs that have a higher chance of misclassification with lower perturbations. In other words, models that are more sensitive to perturbations will cause a higher change in output with small input variations. The authors in [11] have shown that the sensitivity of a model with respect to its input variations is dictated by the magnitude of its Jacobian or gradients. The expression for the $(i,j)$ component of the Jacobian for model $F$ at distillation temperature $T$ is:

Here $z_{0}(X),...,Z_{N-1}(X)$ are the logits. For fixed values of logits, increasing distillation temperature $T$ will lead to decreasing the magnitude of the model gradients as can be inferred from the above expression. This reduces the model sensitivity to adversarial perturbations, and larger input variations are needed to craft the samples. In other words, increasing distillation temperature reduces the magnitude of the adversarial gradients. It is to be noted that the distillation temperature is introduced only during training, and it is not used in the evaluation, i.e. for testing the models, $T$ is set to 1.

To measure the efficacy of defensive distillation, the robustness metric was introduced in [18]. The robustness of a model refers to its resistance to adversarial perturbations. A robust DNN should perform well (in terms of prediction accuracy) both on the training and the outside data, and should be consistent in its class predictions for inputs in the neighborhood of a given sample. Robustness is achieved when this consistency can be ensured in a closed neighborhood. The larger this neighborhood, the higher the robustness of the DNN. The neighborhood cannot be extended indefinitely, as that is the case of a constant function. The robustness of a trained DNN classifier $F$ is given as:

$X$ is the input data drawn from the true distribution $\mu$, and $\Delta_{adv}(X,F)$ is the minimum perturbation that causes a misclassification. Therefore,

The distance metric can be chosen accordingly. For our project, we compute the average value of $\Delta_{adv}(X,F)$ computed over all test samples. Higher the value, more robust is the model.

For our project, we used two legacy image classification datasets - MNIST [14] and CIFAR10 [19]. The MNIST dataset is used for classifying handwritten digits (0-9). It consists of 60,000 training samples and 10,000 testing samples, and the pixels are encoded to $[0,1]$. The CIFAR10 dataset consists of 60,000 color images (three color components) divided into 50,000 training and 10,000 testing samples. Similar to MNIST, CIFAR10 images are classified into 10 mutually exclusive classes. We chose these simple datasets because we could construct relatively shallow models to achieve satisfactory performance. Also, we needed to train and evaluate a large number of models for different temperatures and multiple steps, so we wanted to work with small datasets that can be trained efficiently. Moreover, the previous papers on defensive distillation also dealt with these two simple multi-class classification problems, and we wanted to demonstrate improvements over them.

To remain consistent with previous works on defensive distillation, we used the same models as [11] and [6]. Both the DNN architectures have 9 layers consisting of 4 convolution layers, 2 max pooling layers, and 3 fully connected dense layers in the head. Momentum and parameter decay were used to ensure convergence, and dropout was used to prevent overfitting. For the softmax layer, we experimented on distillation temperatures $T=\{1,2,5,10,20,30,40\}$. For training, we used a batch size of 128, and a learning rate of $\eta=0.01$, decay rate $10^{-6}$, momentum 0.9, and SGD optimizer for 50 epochs.

We experimented with the three attacks proposed in [6] as described in section II. For each of the attacks, we evaluated the robustness (average perturbation) of the teacher, assistant, and student models for the entire temperature range. For $L_{2}$ attack, we used a learning rate of $10^{-2}$, maximum number of iterations as 10,000, initial value of constant $c$ as $10^{-3}$, and zero confidence value. For $L_{0}$ attack, we reduced the maximum number of iterations to 1,000, and the upper limit of the constant to $2\times 10^{-6}$. For $L_{\infty}$ attack, the learning rate was fixed at $5\times 10^{-3}$, the initial value of the constant was set to $10^{-}5$, and the upper threshold was set to $20$. The number of maximum iterations was kept the same as $L_{0}$ attack.

In the first set of experiments, we wanted to observe the variation of the accuracy of the distilled models as compared to the original teacher model when trained at different distillation temperatures. For each of the datasets, MNIST and CIFAR10, we evaluate the performance of the models on the test dataset (10,000 images). As mentioned before, for testing we keep the distillation temperature to be $1$. The baseline accuracy for both the datasets are measured by evaluating the teacher model trained with temperature $T=1$. The accuracy for baseline MNIST model $F_{MNIST}$ was $99.38\%$, and for $F_{CIFAR10}$ was $77.72\%$. From figure 3 we can notice that the accuracy for either of the datasets varies very little with distillation. For MNIST, the maximum variation that can be noticed w.r.t the teacher is about $0.15\%$, whereas for CIFAR10, the maximum variation noticed is $1.2\%$. It should also be noted that for each of the models, there is a trend of accuracy decreasing with temperature. With the increase in temperature, most of the class outputs have higher probability values and that might make the class distinctions a bit difficult (As $T\to\infty$, the output of softmax converges to $1/N$). It is also interesting to note that in most cases, for a fixed temperature, the distilled models leads to better performance. This may be due to the increase in generalization ability of the network that leads to better prediction on unobserved samples.

The second set of experiments pertains to the effect of distillation on the sensitivity of the models. As described in Section III, the sensitivity of a model is measured using the magnitude of its gradients. For each of the two datasets, we computed the gradient magnitudes for the 10,000 test samples as inputs, and took their average. According to the hypothesis presented in the analysis, increasing the distillation temperature should result in decreasing the magnitude of the gradients. Small gradients mean the function is smoother around the sample points in the distribution. A smoother function with lower gradient magnitudes would need higher perturbation for misclassification, which we will show later. The effect of temperature is illustrated in figure 4. For the teacher models, each trained at different temperatures, we compute the magnitude of the gradients averaged over all input dimensions. Repeating this for 10,000 test samples, we have 10,000 instances of mean gradient amplitudes. We then calculate the proportion of samples for which the gradient magnitude is close to zero. For simplifying the illustration, we consider any gradient with amplitude $<10^{-10}$ to be close to zero. The proportion of such gradients is plotted for both models. As expected, higher temperatures cause the trained models to produce smaller gradients and thus smoother outputs.

Robustness is measured by computing the minimum perturbation needed for misclassification (Sec III). In the third set of experiments, we evaluate the efficacy of robustness against the attacks stated in section II. Our goal is to observe how temperature affects the robustness of a model, and if the introduction of the assistant in distillation improves the robustness of the student model. To evaluate this, we craft adversarial samples using each attack on all 10,000 inputs from the test set. For each sample input, the minimum perturbation for misclassifying it is measured (distance between the input and the crafted sample). Finally, we calculate the average of the perturbations over the entire test set which is the robustness value of the model. We repeat this experiment for both the MNIST and CIFAR10 datasets.

Figure 5 is a collection of plots for different models. In each subplot, the orange line indicates the variation of adversarial perturbation for the assistant model, the blue line indicates that of the student model, and the red dotted line is that of the baseline teacher model. For each dataset, the plots on the top row are of the mean deviations (or perturbations), which is the robustness metric, and the bottom row is the maximum deviation required to craft an adversarial sample. Irrespective of the attack, it can be observed that, in general, increasing the distillation temperature improves the robustness of the model. Also, on average, the student model has higher robustness than the teacher and assistant models. This shows that introducing an assistant network did improve the robustness against adversarial attacks in most cases. Among the different attacks, it can be noticed that $L_{2}$ is the most efficient as it generates adversarial samples successfully with the least perturbations among the three, whereas $L_{0}$ is the least efficient. From these sets of experiments, it can be inferred that an assistant model in distillation does help in defending better against adversarial attacks with effectively very minimal hit to performance.

The final experiments that we conducted were computing the confidence values of the previously stated models. The confidence of a model calculated over a dataset $\mathcal{X}$ is the average of the following quantity over all $X\in\mathcal{X}$:

Here $t(X)$ is the correct class for input $X$. In [11], the authors have stated the confidence values increase with distillation temperature for the models trained on CIFAR10. We repeated the same experiment for our models but did not observe any concrete trend over temperature. We also did not observe any substantial difference among the teacher, assistant, and student models as shown in Fig. 6. The variation in confidence values is very little in either case.