SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Seokil Ham $^{1}$  Jungwuk Park $^{1}$  Dong-Jun Han $^{2*}$  Jaekyun Moon $^{1}$
$^{1}$ KAIST  $^{2}$ Purdue University
{gkatjrdlf, savertm}@kaist.ac.kr, han762@purdue.edu, jmoon@kaist.edu
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 other exits concurrently. This makes multi-exit networks highly vulnerable to simple adversarial attacks. In this paper, we propose NEO-KD, a knowledge-distillation-based adversarial training strategy that tackles this fundamental challenge based on two key contributions. NEO-KD first resorts to neighbor knowledge distillation to guide the output of the adversarial examples to tend to the ensemble outputs of neighbor exits of clean data. NEO-KD also employs exit-wise orthogonal knowledge distillation for reducing adversarial transferability across different submodels. The result is a significantly improved 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 the baselines relying on existing adversarial training or knowledge distillation techniques for multi-exit networks.
Multi-exit neural networks are receiving significant attention [9, 13, 26, 27, 28, 32] for their ability to make dynamic predictions in resource-constrained applications. Instead of making predictions at the final output of the full model, a faster prediction can be made at an earlier exit depending on the current time budget or computing budget. In this sense, a multi-exit network can be viewed as an architecture having multiple submodels, where each submodel consists of parameters from the input of the model to the output of a specific exit. These submodels are highly correlated as they share some model parameters. It is also well-known that the performance of all submodels can be improved by distilling the knowledge of the last exit to other exits, i.e., via self-distillation [15, 20, 24, 27]. There have also been efforts to address the adversarial attack issues in the context of multi-exit networks [3, 12].
Providing robustness against adversarial attacks is especially challenging in multi-exit networks: since different submodels have high correlations by sharing parameters, an adversarial example targeting a specific exit can significantly degrade the performance of other submodels. In other words, an adversarial example can have strong adversarial transferability across different submodels, making the model highly vulnerable to simple adversarial attacks (e.g., an adversarial attack targeting a single exit).
Motivation. Only a few prior works have focused on adversarial defense strategies for multi-exit networks [3, 12]. The authors of [12] focused on generating adversarial examples tailored to multi-exit networks (e.g., generate samples via max-average attack), and trained the model to minimize the sum of clean and adversarial losses of all exits. Given the adversarial example constructed in [12], the authors of [3] proposed a regularization term to reduce the weights of the classifier at each exit during training. However, existing adversarial defense strategies [3, 12] do not directly handle the high correlations among different submodels, resulting in high adversarial transferability and limited robustness in multi-exit networks. To tackle this difficulty, we take a knowledge-distillation-based approach in a fashion orthogonal to prior works [3, 12]. Some previous studies [8, 23, 33, 34] have shown that knowledge distillation can be utilized for improving the robustness of the model in conventional single-exit networks. However, although there are extensive existing works on self-distillation for training multi-exit networks using clean data [15, 20, 24, 27], it is currently unknown how distillation techniques should be utilized for adversarial training of multi-exit networks. Moreover, when the existing distillation-based schemes are applied to multi-exit networks, the dependencies among submodels become higher since the same output (e.g., the knowledge of the last exit) is distilled to all sub-models. Motivated by these limitations, we pose the following questions: How can we take advantage of knowledge-distillation to improve adversarial robustness of multi-exit networks? At the same time, how can we reduce adversarial transferability across different submodels in multi-exit networks?
Main contributions. To handle these questions, we propose NEO-KD, a knowledge-distillation-based adversarial training strategy highly tailored to robust multi-exit neural networks. Our solution is two-pronged: neighbor knowledge distillation and exit-wise orthogonal knowledge distillation.
- Given a specific exit, the first part of our solution, neighbor knowledge distillation (NKD), distills the ensembled prediction of neighbor exits of clean data to the prediction of the adversarial example at the corresponding exit, as shown in Figure 1a. This method guides the output of adversarial examples to follow the outputs of clean data, improving robustness against adversarial attacks. By ensembling the neighbor predictions of clean data before distillation, NKD provides higher quality features to the corresponding exits compared to the scheme distilling with only one exit in the same position.
- The second focus of our solution, exit-wise orthogonal knowledge distillation (EOKD), mainly aims at reducing adversarial transferability across different submodels. This part is another unique contribution of our work compared to existing methods on robust multi-exit networks [3, 12] (that suffer from high adversarial transferability) or self-distillation-based multi-exit networks [15, 20, 24, 27] (that further increase adversarial transferability). In our EOKD, the output of clean data at the  $i$ -th exit is distilled to the output of the adversarial sample at the  $i$ -th exit, in an exit-wise manner. During this exit-wise distillation process, we encourage the non-ground-truth predictions of individual exits to be mutually orthogonal, by providing orthogonal soft labels to each exit as described in Figure 1b. By weakening the dependencies among different exit outputs, EOKD reduces the adversarial transferability across all submodels in the network, which leads to an improved robustness against adversarial attacks.
The NKD and EOKD components of our architectural solution work together to reduce adversarial transferability across different submodels in the network while correctly guiding the predictions of the adversarial examples at each exit. Experimental results on various datasets show that the proposed strategy achieves the best adversarial accuracy with reduced computation budgets, compared to existing adversarial training methods for multi-exit networks. Our solution is a plug-and-play method, which can be used in conjunction with existing training strategies tailored to multi-exit networks.
Knowledge distillation for multi-exit networks. Multi-exit neural networks [9, 13, 26, 27, 28, 32] aim at making efficient inference via early exits in resource-constrained applications. In the multi-exit network literature, it is well-known that distilling the knowledge of the last exit to others significantly improves the overall performance on clean data without an external teacher network, i.e., via self-distillation [15, 20, 24, 27]. However, it is currently unclear how adversarial training can benefit from self-distillation in multi-exit networks. One challenge is that simply applying existing self-distillation techniques increases adversarial transferability across different submodels, since the same knowledge
from the last exit is distilled to all other exits, increasing dependency among different submodels in the network. Compared to the existing ideas, our contribution is to develop a self-distillation strategy that does not increase the dependency of submodels as much; this helps reduce adversarial transferability of the multi-exit network for better robustness.
Improving adversarial robustness. Most existing defense methods [6, 7, 31] have mainly focused on creating new adversarial training losses tailored to single-exit networks. Several other works have utilized the concept of knowledge distillation [8, 23, 33, 34] showing that distilling the knowledge of the teacher network can improve robustness of the student network. Especially in [8], given a teacher network, it is shown that robustness of the teacher network can be distilled to the student network during adversarial training. Compared to these works, our approach can be viewed as a new self-distillation strategy for multi-exit networks where teacher/student models are trained together. More importantly, adversarial transferability across different submodels has not been an issue in previous works as the focus there has been on the single-exit network. In contrast, in our multi-exit setup, all submodels sharing some model parameters require extra robustness against adversarial attacks; this motivates us to propose exit-wise orthogonal knowledge distillation, to reduce adversarial transferability among different submodels.
Some prior works [19, 22, 29, 30] aim at improving adversarial robustness of the ensemble model, by reducing adversarial transferability across individual models. Specifically, the adaptive diversity-promoting regularizer proposed in [22] regularizes the non-maximal predictions of individual models to be mutually orthogonal, and the maximal term is used to compute the loss as usual. While the previous work focuses on reducing the transferability among different models having independent parameters, in a multi-exit network setup, the problem becomes more challenging in that all submodels have some shared parameters, making the models to be highly correlated. To handle this issue, we specifically take advantage of knowledge distillation in an exit-wise manner, which can further reduce the dependency among different submodels in the multi-exit network.
Adversarial training for multi-exit networks. When focused on multi-exit networks, only a few prior works considered the adversarial attack issue in the literature [3, 10, 11, 12]. The authors of [10, 11] focused on generating slowdown attacks in multi-exit networks rather than defense strategies. In [12], the authors proposed an adversarial training strategy by generating adversarial examples targeting a specific exit (single attack) or multiple exits (average attack and max-average attack). However, (i) [12] does not take advantage of knowledge distillation during training and (ii) [12] does not directly handle the high correlations among different submodels, which can result in high adversarial transferability. Our solution overcomes these limitations by reducing adversarial transferability while correctly guiding the predictions of adversarial examples at each exit, via self knowledge distillation.
3 Proposed NEO-KD Algorithm
Consider a multi-exit network with  $L$  exits, which is composed of  $L$  blocks  $\{\phi_i\}_{i=1}^L$  and  $L$  classifiers  $\{w_i\}_{i=1}^L$ . Given the input data  $x$ , the output of the  $i$ -th exit is denoted as  $f_{\theta_i}(x)$ , which is parameterized by the  $i$ -th submodel  $\theta_i = [\phi_1, \dots, \phi_i, w_i]$  that consists of  $i$  blocks and one classifier. Note that all  $L$  submodels produce different predictions  $[f_{\theta_1}(x), \dots, f_{\theta_L}(x)]$ . Here, since each submodel shares several blocks with other submodels, the predictions of any two submodels are highly correlated.
3.1 Problem Setup: Adversarial Training in Multi-Exit Networks
The first step for adversarial training is to generate adversarial examples. Given  $L$  different submodels  $\{\theta_i\}_{i=1}^L$ , clean data  $x$ , and the corresponding label  $y$ , the adversarial example  $x^{adv}$  can be generated based on single attack, max-average attack or average attack, following the process of [12]. More specifically, we have
$$
x _ {\text {s i n g l e}, i} ^ {\text {a d v}} = \underset {x ^ {\prime} \in \{z: | z - x | _ {\infty} \leq \epsilon \}} {\arg \max } | \ell \left(f _ {\theta_ {i}} \left(x ^ {\prime}\right), y\right) |, \tag {1}
$$
$$
x _ {m a x} ^ {a d v} = x _ {i ^ {*}} ^ {a d v}, \text {w h e r e} i ^ {*} = \arg \max  _ {i} \left| \frac {1}{L} \sum_ {j = 1} ^ {L} \ell \left(f _ {\theta_ {j}} \left(x _ {s i n g l e, i} ^ {a d v}\right), y\right) \right|, \tag {2}
$$
$$
x _ {a v g} ^ {a d v} = \underset {x ^ {\prime} \in \{z: | z - x | _ {\infty} \leq \epsilon \}} {\arg \max } \left| \frac {1}{L} \sum_ {j = 1} ^ {L} \ell \left(f _ {\theta_ {j}} \left(x ^ {\prime}\right), y\right) \right|, \tag {3}
$$
(a) Neighbor knowledge distillation
(b) Exit-wise orthogonal knowledge distillation
Figure 1: NEO-KD consists of two parts that together improve the adversarial robustness: NKD and EOKD. (a) NKD guides the output of the adversarial data to mimic the ensemble outputs of neighbor exits of clean data. (b) EOKD reduces adversarial transferability of the network by distilling orthogonal knowledge of the clean data to adversarial data for the non-ground-truth predictions, in an exit-wise manner. Although omitted in this figure, EOKD normalizes the likelihood before distilling the soft labels. The overall process operates in a single model, although we consider two cases depending on the input (clean or adversarial example) for a clear presentation.
which correspond to the adversarial example generated by single attack targeting exit  $i$ , max-average attack, and average attack, respectively.  $\epsilon$  denotes the perturbation degree. In the single attack of (1), the adversarial example  $x_{i}^{adv}$  is generated to maximize the cross-entropy loss  $\ell(\cdot, \cdot)$  of the target exit utilizing an attacker algorithm (e.g., PGD [21]). In the max-average attack of (2), among the adversarial examples generated by the single attack for all exits  $i = 1, 2, \ldots, L$ , the sample that maximizes the average loss of all exits is selected. Finally, the average attack in (3) directly generates an adversarial sample that maximizes the average loss of all exits. Based on the generated  $x^{adv}$ , a typical strategy is to update the model considering both the clean and adversarial losses of all exits as follows:
$$
\mathcal {L} = \frac {1}{N} \sum_ {j = 1} ^ {N} \sum_ {i = 1} ^ {L} [ \ell \left(f _ {\theta_ {i}} \left(x _ {j}\right), y _ {j}\right) + \ell \left(f _ {\theta_ {i}} \left(x _ {j} ^ {a d v}\right), y _ {j}\right) ], \tag {4}
$$
where  $N$  is the number of samples in the training set and  $x_{j}^{adv}$  is the adversarial example corresponding to clean sample  $x_{j}$  generated by one of the attacks described above. However, the loss in (4) does not directly consider the correlation among submodels, which can potentially increase adversarial transferability of the multi-exit network.
3.2 Algorithm Description
To tackle the limitations of prior works [3, 12], we propose neighbor exit-wise orthogonal knowledge distillation (NEO-KD), a self-distillation strategy tailored to robust multi-exit networks. To gain insights, we divide our solution into two distinct components with different roles - neighbor knowledge distillation and exit-wise orthogonal knowledge distillation - and first describe each component separately, and then put together into the overall NEO-KD method. A high-level description of our approach is given in Figure 1.
Neighbor knowledge distillation (NKD). The first component of our solution, NKD, guides the output feature of adversarial data at each exit to mimic the output feature of clean data. Specifically, the proposed NKD loss of the  $j$ -th train sample at the  $i$ -th exit is written as follows:
$$
N K D _ {i, j} = \left\{ \begin{array}{l l} \ell \left(f _ {\theta_ {1}} \left(x _ {j} ^ {a d v}\right), \frac {1}{2} \sum_ {k = 1} ^ {2} f _ {\theta_ {k}} \left(x _ {j}\right)\right) & i = 1 \\ \ell \left(f _ {\theta_ {L}} \left(x _ {j} ^ {a d v}\right), \frac {1}{2} \sum_ {k = L - 1} ^ {L} f _ {\theta_ {k}} \left(x _ {j}\right)\right) & i = L \\ \ell \left(f _ {\theta_ {i}} \left(x _ {j} ^ {a d v}\right), \frac {1}{3} \sum_ {k = i - 1} ^ {i + 1} f _ {\theta_ {k}} \left(x _ {j}\right)\right) & \text {o t h e r w i s e ,} \end{array} \right. \tag {5}
$$
which can be visualized with the colored arrows as in Figure 1a. Different from previous self-knowledge distillation methods, for each exit  $i$ , NKD generates a teacher prediction by ensembling (averaging) the neighbor predictions (i.e., from exit  $i - 1$  and exit  $i + 1$ ) of clean data and distills it to each prediction of adversarial examples.
Compared to other distillation strategies, NKD takes advantage of only these neighbors during distillation, which has the following key advantages for improving adversarial robustness. First,
by ensembling the neighbor predictions of clean data before distillation, NKD provides a higher quality feature of the original data to the corresponding exit; compared to the naive baseline that is distilled with only one exit in the same position (without ensembling), NKD achieves better adversarial accuracy, where the results are provided in Table 7. Secondly, by considering only the neighbors during ensembling, we can distill different teacher predictions to each exit. Different teacher predictions of NKD also play a role of reducing adversarial transferability compared to the strategies that distill the same prediction (e.g., the last exit) to all exit; ensembling other exits (beyond the neighbors) increases the dependencies among submodels, resulting in higher adversarial transferability. The corresponding results are also shown via experiments in Section 4.3.
However, a multi-exit network trained with only NKD loss still has a significant room for mitigating adversarial transferability further. In the following, we describe the second part of our solution that solely focuses on reducing adversarial transferability of multi-exit networks.
Exit-wise orthogonal knowledge distillation (EOKD). EOKD provides orthogonal soft labels to each exit for the non-ground-truth predictions, in an exit-wise manner. As can be seen from the red arrows in Figure 1b, the output of clean data at the  $i$ -th exit is distilled to the output of adversarial example at the  $i$ -th exit. During this exit-wise distillation process, some predictions are discarded to encourage the non-ground-truth predictions of individual exits to be mutually orthogonal. We randomly allocate the classes of non-ground-truth predictions to each exit for every epoch, which prevents the classifier to be biased compared to the fixed allocation strategy. The proposed EOKD loss of the  $j$ -th sample at the  $i$ -th exit is defined as follows:
$$
E O K D _ {i, j} = \ell \left(f _ {\theta_ {i}} \left(x _ {j} ^ {a d v}\right), O \left(f _ {\theta_ {i}} \left(x _ {j}\right)\right)\right). \tag {6}
$$
Here,  $O(\cdot)$  is the orthogonal labeling operation to make the non-ground-truth predictions to be orthogonal across all exits. For each exit,  $O(\cdot)$  randomly selects  $\lfloor (C - 1) / L \rfloor$  labels among a total of  $C$  classes so that the selected labels are non-overlapping across different exits (except for the ground-truth label), as in Figure 1b. Lastly, the probabilities of selected labels are normalized to sum to one. To gain clearer insights, consider a toy example with a 3-exit network (i.e.,  $L = 3$ ) focusing on a 4-way classification task (i.e.,  $C = 4$ ). Let  $[p_1^i, p_2^i, p_3^i, p_4^i]$  be the softmax output of the clean sample at the  $i$ -th exit, for  $i = 1, 2, 3$ . If class 1 is the ground-truth, the orthogonal labeling operation  $O(\cdot)$  jointly produces the following results from each exit:  $[\hat{p}_1^1, \hat{p}_2^1, 0, 0]$  from exit 1,  $[\hat{p}_1^2, 0, \hat{p}_2^2, 0]$  from exit 2,  $[\hat{p}_1^3, 0, 0, \hat{p}_4^3]$  from exit 3, where  $\hat{p}$  indicates the normalized probability of  $p$  so that the values in each vector sum to one.
Based on Eq. (6), the output of the orthogonal label operation  $O(f_{\theta_i}(x_j))$  for the clean data  $x_j$  is distilled to  $f_{\theta_i}(x_j^{adv})$  which is the prediction of the adversarial example of the  $j$ -th sample at the  $i$ -th exit. This encourages the model to self-distill orthogonally distinct knowledge in an exit-wise manner while keeping the essential knowledge of the ground-truth class. By taking this exit-wise orthogonal distillation approach, EOKD reduces the dependency among different submodels, reducing the adversarial transferability of the network.
Overall NEO-KD loss. Finally, considering the proposed loss functions in Eq. (5), (6) and the original adversarial training loss, the overall objective function of our scheme is written as follows:
$$
\mathcal {L} = \frac {1}{N} \sum_ {j = 1} ^ {N} \sum_ {i = 1} ^ {L} [ \ell \left(f _ {\theta_ {i}} \left(x _ {j}\right), y _ {j}\right) + \ell \left(f _ {\theta_ {i}} \left(x _ {j} ^ {a d v}\right), y _ {j}\right) + \gamma_ {i} (\alpha \cdot N K D _ {i, j} + \beta \cdot E O K D _ {i, j}) ], \tag {7}
$$
where  $\alpha, \beta$  control the weights for each component and  $\gamma_{i}$  is the knowledge distillation weight for each exit  $i$ . Since later exits have lower knowledge distillation loss than the earlier exits, we impose a slightly higher  $\gamma_{i}$  for the later exits than  $\gamma_{i}$  of the earlier exits. More details regarding hyperparameters are described in Appendix.
By introducing two unique components - NKD and EOKD - the overall NEO-KD loss function in Eq. (7) reduces adversarial transferability in the multi-exit network while correctly guiding the output of the adversarial examples in each exit, significantly improving the adversarial robustness of multi-exit networks, as we will see in the next section.
In this section, we evaluate our method on five datasets commonly adopted in multi-exit networks: MNIST [18], CIFAR-10, CIFAR-100 [16], Tiny-ImageNet [17], and ImageNet [25]. For MNIST, we
Exit	1	2	3	Average
Adv. w/o Distill [12]	89.74%	95.89%	96.82%	94.15%
SKD [24]	89.77%	96.24%	97.08%	94.36%
ARD [8]	89.65%	95.79%	96.47%	93.97%
LW [3]	87.08%	93.86%	95.42%	92.12%
NEO-KD (ours)	90.56%	96.30%	96.62%	94.49%
(a) Max-average attack
Exit	1	2	3	Average
Adv. w/o Distill [12]	85.17%	95.87%	97.20%	92.75%
SKD [24]	84.35%	96.53%	97.53%	92.82%
ARD [8]	85.07%	96.12%	97.28%	92.82%
LW [3]	85.26%	94.54%	96.05%	91.95%
NEO-KD (ours)	86.17%	96.92%	97.42%	93.50%
Table 1: Anytime prediction setup: Adversarial test accuracy on MNIST.
(b) Average attack
Exit	1	2	3	Average
Adv. w/o Distill [12]	44.77%	46.10%	46.86%	45.91%
SKD [24]	44.75%	44.54%	44.79%	44.69%
ARD [8]	44.50%	45.85%	51.82%	47.39%
LW [3]	37.38%	35.39%	34.85%	35.87%
NEO-KD (ours)	46.53%	47.65%	50.71%	48.30%
(a) Max-average attack
Exit	1	2	3	Average
Adv. w/o Distill [12]	38.00%	41.42%	40.70%	40.04%
SKD [24]	39.36%	41.39%	38.39%	39.71%
ARD [8]	39.37%	41.98%	43.53%	41.63%
LW [3]	31.47%	31.41%	28.98%	30.62%
NEO-KD (ours)	41.67%	45.38%	45.54%	44.20%
(b) Average attack
Table 2: Anytime prediction setup: Adversarial test accuracy on CIFAR-10.
use SmallCNN [12] with 3 exits. We trained the MSDNet [13] with 3 and 7 exits using CIFAR-10 and CIFAR-100, respectively. For Tiny-ImageNet and ImageNet, we trained the MSDNet with 5 exits. More implementation details are provided in Appendix.