Robust Mixture-of-Expert Training for Convolutional Neural Networks
Sparsely-gated Mixture of Expert (MoE), an emerging deep model architecture,
has demonstrated a great promise to enable high-accuracy and ultra-efficient
model inference. Despite the growing popularity of MoE, little work
investigated its potential to advance convolutional neural networks (CNNs),
especially in the plane of adversarial robustness. Since the lack of robustness
has become one of the main hurdles for CNNs, in this paper we ask: How to
adversarially robustify a CNN-based MoE model? Can we robustly train it like an
ordinary CNN model? Our pilot study shows that the conventional adversarial
training (AT) mechanism (developed for vanilla CNNs) no longer remains
effective to robustify an MoE-CNN. To better understand this phenomenon, we
dissect the robustness of an MoE-CNN into two dimensions: Robustness of routers
(i.e., gating functions to select data-specific experts) and robustness of
experts (i.e., the router-guided pathways defined by the subnetworks of the
backbone CNN). Our analyses show that routers and experts are hard to adapt to
each other in the vanilla AT. Thus, we propose a new router-expert alternating
Adversarial training framework for MoE, termed AdvMoE. The effectiveness of our
proposal is justified across 4 commonly-used CNN model architectures over 4
benchmark datasets. We find that AdvMoE achieves 1% ~ 4% adversarial robustness
improvement over the original dense CNN, and enjoys the efficiency merit of
sparsity-gated MoE, leading to more than 50% inference cost reduction. Codes
are available at https://github.com/OPTML-Group/Robust-MoE-CNN.
Yihua Zhang, Ruisi Cai, Tianlong Chen, Guanhua Zhang, Huan Zhang, Pin-Yu Chen, Shiyu Chang, Zhangyang Wang, Sijia Liu
# Robust Mixture-of-Expert Training for Convolutional Neural Networks

###### Abstract

Sparsely-gated Mixture of Expert (MoE), an emerging deep model architecture, has demonstrated a great promise to enable high-accuracy and ultra-efficient model inference. Despite the growing popularity of MoE, little work investigated its potential to advance convolutional neural networks (CNNs), especially in the plane of adversarial robustness. Since the lack of robustness has become one of the main hurdles for CNNs, in this paper we ask: How to adversarially robustify a CNN-based MoE model? Can we robustly train it like an ordinary CNN model? Our pilot study shows that the conventional adversarial training (AT) mechanism (developed for vanilla CNNs) no longer remains effective to robustify an MoE-CNN. To better understand this phenomenon, we dissect the robustness of an MoE-CNN into two dimensions: Robustness of routers (i.e., gating functions to select data-specific experts) and robustness of experts (i.e., the router-guided pathways defined by the subnetworks of the backbone CNN). Our analyses show that routers and experts are hard to adapt to each other in the vanilla AT. Thus, we propose a new router-expert alternating Adversarial training framework for MoE, termed AdvMoE. The effectiveness of our proposal is justified across 4 commonly-used CNN model architectures over 4 benchmark datasets. We find that AdvMoE achieves \(1\%\sim 4\%\) adversarial robustness improvement over the original dense CNN, and enjoys the efficiency merit of sparsity-gated MoE, leading to more than \(50\%\) inference cost reduction. Codes are available at [https://github.com/OPTML-Group/Robust-MoE-CNN](https://github.com/OPTML-Group/Robust-MoE-CNN).

+
Footnote †: Correspondence to: Yihua Zhang\(<\)han1908@msu.edu\(>\)

## 1 Introduction

Despite the state-of-the-art performance achieved by the outrageously large networks [1, 2, 3, 4, 5] in various deep learning (DL) tasks, it still remains challenging to train and deploy such models cheaply. A major bottleneck is the lack of parameter efficiency [6]: A single data prediction only requires activating a small portion of the parameters of the full model. Towards efficient DL, sparse Mixture of Experts (MoE) [7, 8, 9, 10, 11, 12, 13, 14, 15] aims to divide and conquer the model parameters based on their optimal responses to specific inputs so that inference costs can be reduced. A typical MoE structure is comprised of a set of 'experts' (_i.e._, sub-models extracted from the original backbone network) and 'routers' (_i.e._, additional small-scale gating networks to determine expert selection schemes across layers). During inference, sparse MoE only activates the most relevant experts and forms the expert-guided pathway for a given input data. By doing so, sparse MoE can boost the inference efficiency (see 'GFLOPS' measurement in **Fig. 1**). Architecture-wise, sparse MoE has been used for both CNNs [8, 16] and vision transformers (VITs) [7, 9, 10, 11, 12, 13, 14, 15, 17]. Yet, we will focus on the former since sparse MoE for CNNs is under-explored compared to non-sparse MoE for CNNs [18, 19, 20], and adversarial robustness (another key performance metric of our work) was extensively studied in the context of CNNs.

It is known that a main weakness of DL is the lack of adversarial robustness [21, 22, 23]. For example, CNNs can be easily fooled by adversarial attacks [21, 22, 23], in terms of tiny input perturbations generated to direct to erroneous predictions. Thus, adversarial training (**AT**) of CNNs has become a main research thrust [24, 25, 26, 27, 28, 29]. However, when CNN meets sparse MoE, it remains elusive if the improved inference efficiency brought by the sparse MoE comes at the cost of more complex adversarial training recipes. Thus, we ask:

[leftmargin=*]
**(Q)** _What will be the new insights into adversarial robustness of sparse MoE-integrated CNNs? And what will be the suited AT mechanism?_

To our best knowledge, problem (**Q**) remains open in the literature. The most relevant work to ours is [30], which investigated the adversarial robustness of MoE and lever
aged the ordinary AT recipe [24] to defend against adversarial attacks. However, it only focused on the ViT architecture, making a vacancy for the research on robustification for the sparse MoE-based CNN (termed **MoE-CNN** in this work). Most importantly, we find that the vanilla AT [24, 25] (widely used to robustify CNNs) is _no longer_ effective for MoE-CNN. Thus, new solutions are in demand.

To address **(Q)**, we need to (1) make careful sanity checks for AT in MoE-CNN, (2) make an in-depth analysis of its failure cases, and (3) advance new AT principles that can effectively improve robustness without losing the generalization and efficiency from sparse MoE. Specifically, our **contributions** are unfolded below.

* We dissect the MoE robustness into two new dimensions (different from CNNs): routers' robustness and experts' robustness. Such a robustness dissection brings novel insights into the (in)effectiveness of AT.
* Taking inspiration from the above robustness dissection, we propose a new Adversarial training framework for MoE, termed AdvMoE, which enforces routers and experts to make a concerted effort to improve the overall robustness of MoE-CNN.
* We conduct extensive experiments to demonstrate the effectiveness of AdvMoE across \(4\) CNN architectures and 4 datasets. For example, AdvMoE outperforms AT on the original dense CNN model (termed Dense) by a substantial margin: \(1\%\sim 4\%\) adversarial robustness improvement and over \(50\%\) reduction of inference overhead; see **Fig. 1** for illustrations on different CNN types and highlighted performance achieved.

## 2 Related Work

**Sparsely-activated Mixture of Experts (Sparse MoE).** As a special instance of compositional neural architectures [31, 32, 33], MoE [4, 7, 8, 9, 10, 11, 16, 18, 19, 20, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141, 142, 143] aims at solving ML tasks in a divide-and-conquer fashion, which creates a series of sub-models (known as the _experts_) and conducts input-dependent predictions by combing the output of sub-models. As an important branch of MoE, sparsely gated MoE [4, 7, 8, 9, 10, 11, 16, 18, 19, 20, 39, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111,
robustness and standard generalization ability. Other work [26, 28, 47, 48, 49, 50, 51, 52, 53, 54] aims at trimming down the computational costs of robust training while maintaining robustness. The work [30] studies the robustness of MoE-based architectures for the first time. Yet, its focus stays on MoE for ViTs and the relationship between model capacity and robustness.

## 3 Problem Statement

In this section, we start by presenting the setup of MoE-CNN in this work and then introduce the robust learning paradigm. The lack of adversarial robustness of deep models inspires us to investigate whether the adversarial training (AT) approach designed for vanilla CNNs keeps effective for MoE-CNN. Through a motivating example, we show that the conventional AT recipe is _incapable_ of equipping MoE-CNN with desired robustness. The resulting performance is even worse than that of AT-produced S-Dense, which has a much smaller model capacity than MoE-CNN. Thus, the question of how to robustify MoE-CNN arises.

Model setup.We consider a CNN-backboned MoE that consists of multiple MoE layers. Each MoE layer involves a router and a vanilla convolutional layer from the backbone CNN model. Within one MoE layer, we define \(N\) experts, each of which picks a subset of the channels from the convolutional layer. Specifically, suppose the \(l\)-th layer contains \(C_{l}\) channels, one expert will contain \(r\times C_{l}\) channels, where we call the ratio \(r\in[0,1]\)_model scale_ and keep it the same across different layers (see **Fig. 1a**). It is worth noting that as \(r\) increases, the per-expert model capacity increases (_i.e._, with more parameters) at the cost of the efficiency reduction. In a forward path, the router first makes an _input-specific_ expert selection. These selected layer-wise experts then form an end-to-end pathway to process this input. We use "_pathway_" to describe one experts-guided forward path (see **Fig. 1a**). We summarize the model setup in **Fig. A1**.

Further, we introduce different model types considered in this work and shown in **Fig. 1a**. First, we term the original dense CNN model '**Dense**', which serves as the _model basis_ for other model types that derive from. Second, we directly shrink the channel number of each layer in Dense (based on the model scale parameter \(r\)) to obtain the'small dense' model (termed '**S-Dense**'). Notably, S-Dense has the size _equivalent to a single pathway_ in MoE-CNN. Third, we use the structured pruning method [50] to create a sparse subnetwork from Dense, with the weight remaining ratio same as the model scale parameter \(r\) in MoE-CNN, which we call '**Sparse-CNN**'. In summary, S-Dense has the smallest model capacity (comparable to a single pathway of MoE-CNN), and should provide the _performance lower-bound_ for MoE-CNN. By contrast, Sparse-CNN has a larger model capacity but is smaller than MoE-CNN as it encodes a data-agnostic pathway of Dense, while MoE-CNN yields data-specific pathways at the same scale. Dense has the largest model capacity but the least inference efficiency.

Adversarial robustness: From CNN to MoE-CNN.It has been known that current machine learning models (_e.g._, CNNs) are vulnerable to adversarial attacks [21, 22, 23]. Towards the robust design, a variety of AT (adversarial training) methods have been developed. The predominant ones include the min-max optimization-based vanilla AT [24] and its TRADES variant [25] that strikes a balance between generalization and adversarial robustness. Throughout the paper, we adopt TRADES as the default conventional AT recipe, which solves the following problem:

\[\min_{\mathbf{\theta}}\mathbb{E}_{(\mathbf{x},y)\in\mathcal{D}}\left[\ell(\mathbf{ \theta};\mathbf{x},y)+\frac{1}{\lambda}\max_{\|\mathbf{\delta}\|_{\infty}\leq\epsilon }\ell_{\mathrm{KL}}(\mathbf{f_{\theta}}(\mathbf{x}),\mathbf{f_{\theta}}(\mathbf{x}+ \mathbf{\delta}))\right]\] (AT)

where \(\mathbf{\theta}\) denotes model parameters to be robustified, \((\mathbf{x},y)\in\mathcal{D}\) is a training sample, drawn from the training set \(\mathcal{D}\), with input feature \(\mathbf{x}\) and label \(y\), \(\ell(\mathbf{\theta},\mathbf{x};y)\) denotes the cross-entropy loss using model \(\mathbf{\theta}\) at data point \((\mathbf{x},y)\), \(\mathbf{\delta}\) signifies the input perturbation variable subject to the \(\ell_{\infty}\)-norm ball of radius \(\epsilon\), \(\mathbf{f_{\theta}}(\cdot)\) denotes the model's predictions, \(\ell_{\mathrm{KL}}\) is the KL divergence loss that characterizes the worst-case prediction stability at the presence of \(\mathbf{\delta}\), and \(\lambda>0\) is a regularization parameter to strike the tradeoff between empirical risk minimization and the robustness of model predictions.

Although AT has been well studied for adversarial robustness of CNNs, there exists few attempts to robustify MoE-CNN. This raises the problem of our interest:

**(Problem statement)** Can MoE-CNN be robustified as effectively as an ordinary CNN using AT? If not, how to robustly train MoE-CNN to achieve robustness not worse than AT-oriented S-Dense, Sparse-CNN, and Dense while preserving MoE's efficiency?

Warm-up study: AT for MoE-CNN is _not_ trivial.Our goal to robustify MoE-CNN includes (1) achieving high robustness, (2) maintaining high prediction accuracy, and (3) making full use of MoE routing to keep the model's high efficiency and expressiveness. Nonetheless, the routing system in MoE brings extra robustification challenges, which never exist in ordinary CNNs. Specifically, the input-specific expert selection in MoE could make the attacker easier to succeed, since input perturbations can _either_ mislead routers to select incorrect experts _or_ fool the pathophysiological predictor. Such a '_two-way attack mode_' makes AT for MoE-CNN highly non-trivial.

**Fig. 2** empirically justifies that the direct application of (AT) to MoE-CNN is problematic. In Fig. 2, we consider ResNet-18 as the model backbone (Dense) and CIFAR-10 for image classification. We apply (AT) to train MoE-CNN
and S-Dense, and report the robust accuracy (RA), _i.e._, testing accuracy over adversarial examples generated by 50-step PGD attacks [24], against different attack strengths \(\epsilon\).

As we can see, although MoE-CNN has a much larger model capacity than S-Dense, it leads to a significant RA drop when the conventional AT approach is applied. This implies that the design of AT for MoE-CNN is far from trivial. A new robust learning protocol is thus needed to improve the robustness of MoE-CNN without losing its merits in efficiency and generalization.

## 4 Methods

In this section, we start by peering into the failure case of (AT) in MoE-CNN by understanding the roles of the routers and pathways in (AT). We empirically show that these individual components are hard to adapt to each other and cannot make a concerted effort in AT. Based on that, we develop a new AT framework for MoE-CNN, AdvMoE, which also takes inspiration from bi-level optimization.

**Dissecting robustness of MoE-CNN: Routers' robustness vs. pathways' robustness.** The main puzzle in robustifying MoE-CNN comes from the coupling between the robustness of routers (which are responsible for expert selection across layers) and the robustness of the input-specific MoE pathways (which are in charge of the final prediction of an input). Given the failure case of AT for MoE-CNN in Fig. 2, we need to understand the roles of routers and pathways in AT, _i.e._, how the adversarial robustness of MoE-CNN is gained in the presence of the 'two-way attack mode'. To this end, we begin by assessing the influence of the routers' robustness on the overall robustness. This is also inspired by the recent pruning literature [50] showing that model robustness can be gained solely from network's sparse topology (regardless of model weights). We thus ask:

**(Q1)** Is improving routers' robustness sufficient to achieve a robust MoE-CNN?

To tackle **(Q1)**, we first split the parameters of MoE-CNN (_i.e._, \(\mathbf{\theta}\)) into two parts, the parameters of routers \(\mathbf{\phi}\) and the parameters of the backbone network \(\mathbf{\psi}\). This yields \(\mathbf{\theta}=[\mathbf{\phi}^{\top},\mathbf{\psi}^{\top}]^{\top}\), where \(\top\) is the transpose operation. We then call (AT) to robustly train routers (\(\mathbf{\phi}\)) but _fix_ the backbone network (\(\mathbf{\psi}\)) at its standard pre-trained weights. We denote this partially-robustified model by \(\mathbf{\bar{\theta}}=[\mathbf{\phi}^{\top},\mathbf{\psi}^{\top}]^{\top}\), where \(\bar{}\) indicates the updated parameters. To answer **(Q1)**, we assess the robustness gain of \(\mathbf{\bar{\theta}}\) vs. 3 baselines (M1-M3): (**M1**) the standard MoE-CNN \(\mathbf{\theta}\), (**M2**) AT-robustified S-Dense, and (**M3**) Sparse-CNN achieved by the robust sparse mask learning method [50] over the original Dense model.

Based on Insight 1, we further peer into the resilience of expert selection decisions to adversarial examples. If expert selections in _all_ MoE layers keep intact in the presence of an adversarial perturbation, we say that the routing system of MoE-CNN is robust against this adversarial example. We then divide adversarial examples into **four categories** according to whether they successfully attacked routers and the router-oriented pathways: _unsuccessful_ attack on _both_ routers and MoE pathways, _successful_ attack on routers but _not MoE_ pathways, _successful_ attack on MoE pathways but _not routers_, and _successful_ attack on _both_ routers and MoE pathways. Here \(\mathbf{\Theta}\) + \(\mathbf{\Theta}\) characterizes the robustness of routers, while \(\mathbf{\Theta}\) + \(\mathbf{\Theta}\) represents that of MoE. Thus, if \(\mathbf{\Theta}\) or \(\mathbf{\Theta}\) takes a large portion of generated adversarial examples, it implies that the routers' robustness does _not_ directly impact the MoE pathway-based predictor's robustness. **Fig. 4** shows the above categories \(\mathbf{\Theta}\)-\(\mathbf{\Theta}\) when attacking the router-robustified MoE-CNN (_i.e._, \(\mathbf{\bar{\theta}}\)). As we can see, routers' robustness indeed improves prediction robustness (as shown by \(31.74\%\) unsuccessful attacks against the MoE predictor in \(\mathbf{\Theta}\)). However, in the total number of unsuccessful attacks against routers (_i.e._, \(\mathbf{\Theta}\)+\(\mathbf{\Theta}\)\(=76.27\%\)), more than half of them successfully fool

Figure 3: Robustness comparison of router-robustified MoE-CNN (_i.e._\(\mathbf{\hat{\theta}}\)) and baseline models (M1 – M3) for different model scales under CIFAR-10 given the backbone network ResNet-18.

Figure 2: Performance of MoE-CNN and S-Dense robustly trained using (AT) on CIFAR-10 with ResNet-18 as the backbone.
the MoE predictor (_i.e._, ). The above results provide us an additional insight:

**Insight 2:** Improving routers' robustness is _not_ sufficient for the MoE predictor to gain satisfactory robustness although the former makes a positive impact.

Both **Insight 1** and **Insight 2** point out that only improving routers' robustness is _not_ adequate to obtain the desired robustness for the overall MoE-CNN. Thus, we next ask:

**(Q2)** Given the router-robustified model \(\bar{\mathbf{\theta}}\), can we equip \(\bar{\mathbf{\theta}}\) with additional robustness by robustly training expert weights (\(\mathbf{\psi}\))? And how does it further impact routers?

To answer **(Q2)**, we call (AT) to further robustly train the backbone network \(\mathbf{\psi}\) on top of the router-robustified model \(\bar{\mathbf{\theta}}\). We denote the resulting model by \(\bar{\mathbf{\theta}}=[\bar{\mathbf{\phi}}^{\top},\bar{\mathbf{\psi}}^{\top}]\).

**Fig. 5** shows the dissection of the robustness of \(\bar{\mathbf{\theta}}\) in the same setup of Fig. 4. Obviously, the overall prediction robustness (**0+\(\mathbf{\Theta}\)**) is further enhanced after updating \(\bar{\mathbf{\theta}}\) to \(\bar{\mathbf{\theta}}\). Thus, the gains in the robustness of experts' weights indeed further help improve the overall robustness. However, this leads to a surprising drop in the router's robustness (**0+\(\mathbf{\Theta}\)**) when comparing \(\bar{\mathbf{\theta}}\) with \(\bar{\mathbf{\theta}}\). This shows that routers' robustness is _not_ automatically preserved if experts are updated. We obtain the following insight into **(Q2)**:

**Input 3:** Robustness of MoE weights (a), (b), (c).

**Output 4:** Adversarial attack success analysis on dissected MoE-CNN models \(\bar{\mathbf{\theta}}=[\bar{\mathbf{\phi}}^{\top},\mathbf{\psi}^{\top}]\) (model scale \(r=0.5\)), where only \(\bar{\mathbf{\phi}}\) is (AT)-robustified. The adversarial evaluation is based on 50-step PGD attack [24] to fool \(\bar{\mathbf{\theta}}\), and other experiment setups align with Fig. 3. The evaluation is carried out on the test set with a total number of \(10000\) samples.

**Input 4:** Adversarial attack success analysis on dissected MoE-CNN models \(\bar{\mathbf{\theta}}=[\bar{\mathbf{\phi}}^{\top},\mathbf{\psi}^{\top}]\) (model scale \(r=0.5\)), where only \(\bar{\mathbf{\phi}}\) is (AT)-robustified. The adversarial evaluation is based on 50-step PGD attack [24] to fool \(\bar{\mathbf{\theta}}\), and other experiment setups align with Fig. 3. The evaluation is carried out on the test set with a total number of \(10000\) samples.

**Input 5:** Adversarial attack success analysis on forest (a), (b), (c), (d), (e), (f), (g), (h),
are fixed. We term the resulting algorithmic framework as Adversarially robust learning for MoE-CNN (AdvMoE); see Algorithm 1 for a summary.

```
1:Initialize: backbone network \(\mathbf{\psi}\), routers \(\mathbf{\phi}\), batch size \(b\), attack generation step \(K\).
2:for Iteration \(t=0,1,\dots,\)do
3: Pick different random data batches \(\mathcal{B}_{\psi}\) and \(\mathcal{B}_{\phi}\) for backbone and router training
4: Lower-level \(\mathbf{\phi}\)-update (with fixed \(\mathbf{\psi}\)): Given \(\mathbf{\psi}\), update \(\mathbf{\phi}\) by minimizing \(\ell_{\text{TRADES}}\) using \(K\)-step PGD attack [24] generator and SGD (with \(\mathcal{B}_{\mathbf{\phi}}\))
5: Upper-level \(\mathbf{\psi}\)-update (with fixed \(\mathbf{\phi}\)): Given \(\mathbf{\phi}\), update \(\mathbf{\psi}\) by minimizing \(\ell_{\text{TRADES}}\) using \(K\)-step PGD attack gen-erator and SGD (with \(\mathcal{B}_{\mathbf{\psi}}\))
6:endfor
```

**Algorithm 1** The AdvMoE algorithm

We highlight that AdvMoE will train robust routers and robust MoE pathways to 'accommodate' each other. In contrast to the conventional AT framework, AdvMoE delivers the coupled \(\mathbf{\phi}^{\star}(\mathbf{\psi})\) and \(\mathbf{\psi}\), where both parts make a concerted effort to improve the overall robustness. We also remark that AdvMoE does not introduce additional hyper-parameters, since in practice we found routers and experts can share the same learning rate and schedules. More implementation details are provided in Appendix B. In the meantime, we remark that since our proposal is a BLO with non-convex lower and upper-level objectives (1). It is difficult to prove the convergence of AdvMoE. Existing theoretical analysis of BLO typically relies on strongly convex assumptions of lower-level problems [58, 59]. Although without a proper theoretical analysis framework, our method converges well in practice (see Appendix C).

## 5 Experiments

In this section, we will demonstrate the effectiveness of our proposed AdvMoE approach on diverse datasets and models. We will also make an in-depth analysis of the router utility and the expert selection distribution for AdvMoE-trained MoE-CNN.

### Experiment Setup

**Model and dataset setups.** To implement MoE-CNN and other baselines, we conduct experiments on ResNet-18 [60], Wide-ResNet-28-10 [61], VGG-16 [62], and DenseNet [63]. Towards fair assessment, our performance comparison between different model types is restricted to using the same model scale parameter \(r\) (see Fig. 1 for an example). By doing so, an input example will leverage the same amount of model parameters for decision-making. For MoE-CNN, we consider \(N=2\) experts with \(r=0.5\) by default, see Appendix B for more details. Dataset-wise, we focus on the commonly used ones to evaluate the adversarial robustness of image classification [24, 25, 64], including CIFAR-10 [65], CIFAR-100 [65], TinyImageNet [66], and ImageNet [66].

**Baselines.** To make our performance comparison informative and comprehensive, we consider three kinds of baselines that are fairly comparable to (AdvMoE). 1 AT (S-Dense): we apply AT to S-Dense; 2 AT (Sparse): we apply the robustness-aware (structured) sparse mask learning method [50] to obtain Sparse-CNN; 3 AT (MoE): we directly apply AT to MoE-CNN, which co-trains the routers and backbone network. Note this method is also adopted in the latest robust training algorithm [30] for ViT-based MoE architectures. It is worth noting that the above baselines use the same number of model parameters as the pathway of MoE-CNN during model prediction. In addition, we cover 4 AT (Dense) (applying AT to Dense) to acquire a robustness performance reference. Yet, we remark that it is _not_ quite fair to directly compare Dense with the aforementioned smaller counterparts, since the former uses a larger model scale (\(r=1.0\)) at test-time inference.

**Training and evaluation.** We use TRADES [25] as the default robust training objective for all baselines. We also follow the literature [24, 25, 27, 64] to set the attack strength by \(\epsilon=8/255\) for CIFAR-10 and CIFAR-100, and \(\epsilon=2/255\) for TinyImageNet and ImageNet. To implement AdvMoE (Algorithm 1), we mimic the TRADES training pipeline but conduct the proposed BLO routine to robustify routers and backbone parameters in an interactive mode. We adopt \(2\)-step PGD attack [24] at training time for _all_ the methods, supported by the recent work [67] showing its compelling performance in AT. We refer readers to Appendix B for more training details. During evaluation, we report standard accuracy (**SA**) on the clean test dataset and robust accuracy (**RA**) against test-time 50-step PGD attacks [24] with the attack strength same as the training values. We also report **GFLOPs** (FLOPS \(\times 10^{9}\)) as an indicator of the test-time inference efficiency.

### Experiment Results

**Overall performance.** Tab. 1** presents the overall performance of our proposed AdvMoE algorithm vs. baselines. We make several key observations below.

**First.** AdvMoE yields a significant robustness enhancement over all the baselines in every data-model setup. Specifically, AdvMoE consistently yields an improvement of around \(1\%\sim 5\%\) on the robustness measured by RA against PGD attacks. Notably, AdvMoE can also outperform \(\bullet\) AT (Dense) in most cases, around \(1\%\sim 4\%\) robustness improvement (see highlighted results in [green]). This is remarkable since Dense (\(r=1.0\)) is twice larger
than an MoE pathway (\(r=0.5\)). **Second**, we observe that AdvMoE has a preference on wider models. For instance, when WRN-28-10 (the widest model architecture in experiments) is used, AdvMoE yields better robustness over the Dense counterpart across all the dataset setups. **Third**, we also observe that the direct AT application to MoE-CNN, _i.e._, AT (MoE), is worse than AT (S-Dense) and AdvMoE in all setups. This is consistent with our findings in Sec. 4. We remark that although the usefulness of AT (MoE) was exploited in [30] for the MoE-type ViT, it is _not_ effective for training MoE-type CNNs anymore. **Fourth**, AdvMoE can retain the high inference efficiency for MoE-CNN, as evidenced by the GFLOPS measurements in Tab. 1. Compared to S-Dense, MoE-CNN introduces minor computational overhead due to the routing system. However, it saves more than \(50\%\) of the inference cost vs. Dense. This implies that our proposal AdvMoE can preserve the efficiency merit of the MoE structure while effectively improving its adversarial robustness.

**Robust evaluation on AutoAttack [68].** In Tab. 2, we provide additional experiments evaluated by AutoAttack [68] (termed **RA-AA**), a popular robustness evaluation benchmark [69]. The experiment setting in Tab. 2 follows **Tab. 1**. We report RA-AA on CIFAR-10 and CIFAR-100 with ResNet-18 and WRN-28-10. As we can see, although AutoAttack leads to a lower RA-AA compared to RA evaluated using PGD attacks (termed

\begin{table}
\begin{tabular}{l|c|c c c|c|c c c} \hline \hline
**Method** & **Backbone** & **RA (\%)** & **SA (\%)** & **GFLOPS(\#)** & **Method** & **Backbone** & **RA (\%)** & **SA (\%)** & **GFLOPS (\#)** \\ \hline \multicolumn{10}{c}{**CIFAR-10**} \\ \hline \(\bullet\) AT (Dense) & & 50.13\(\pm\)0.13 & 82.99\(\pm\)0.11 & 0.54 & \(\bullet\) AT (Dense) & & 51.75\(\pm\)0.12 & 83.54\(\pm\)0.15 & 5.25 \\ \(\circ\) AT (S-Dense) & & 48.12\(\pm\)0.09 & 80.18\(\pm\)0.11 & 0.14 (74\(\%\)\(\downarrow\)) & \(\circ\) AT (S-Dense) & & 50.66\(\pm\)0.13 & 82.24\(\pm\)0.10 & 1.31 (75\(\%\)\(\downarrow\)) \\ \(\circ\) AT (Sparse) & ResNet-18 & 47.93\(\pm\)0.17 & **80.45\(\pm\)**0.13 & 0.14 (74\(\%\)\(\downarrow\)) & \(\circ\) AT (Sparse) & WRN-28-10 & 48.95\(\pm\)0.14 & 82.44\(\pm\)0.17 & 1.31 (75\(\%\)\(\downarrow\)) \\ \(\circ\) AT (MoE) & & 45.57\(\pm\)0.51 & 78.84\(\pm\)0.75 & 0.15 (72\(\%\)\(\downarrow\)) & \(\circ\) AT (MoE) & & 46.73\(\pm\)0.46 & 77.42\(\pm\)0.73 & 1.75 (67\(\%\)\(\downarrow\)) \\ \(\cline{2-10} \(\Delta\)**DivMoE** & & **51.83**\(\pm\)0.12 & 80.15\(\pm\)0.11 & 0.15 (72\(\%\)\(\downarrow\)) & \(\Delta\)**DivMoE** & & **55.73**\(\pm\)0.13 & **84.32**\(\pm\)0.18 & 1.75 (67\(\%\)\(\downarrow\)) \\ \hline \(\bullet\) AT (Dense) & & 46.19\(\pm\)0.21 & 82.18\(\pm\)0.23 & 0.31 & \(\bullet\) AT (Dense) & & 44.52\(\pm\)0.14 & 74.97\(\pm\)0.19 & 0.07 \\ \(\circ\) AT (S-Dense) & & 45.72\(\pm\)0.18 & **80.10**\(\pm\)0.16 & 0.07 (77\(\%\)\(\downarrow\)) & \(\circ\) AT (Sparse) & & 38.07\(\pm\)0.13 & 69.63\(\pm\)0.11 & 0.02 (71\(\%\)\(\downarrow\)) \\ \(\circ\) AT (Sparse) & VGG-16 & 46.13\(\pm\)0.15 & **79.32**\(\pm\)0.18 & 0.07 (77\(\%\)\(\downarrow\)) & \(\circ\) AT (Sparse) & DenseNet & 37.73\(\pm\)0.13 & 67.35\(\pm\)0.12 & 0.02 (71\(\%\)\(\downarrow\)) \\ \(\circ\) AT (MoE) & & 43.37\(\pm\)0.46 & 76.49\(\pm\)0.65 & 0.12 (61\(\%\)\(\downarrow\)) & \(\Delta\)**DivMoE** & & 35.21\(\pm\)0.74 & 64.41\(\pm\)0.81 & 0.03 (57\(\%\)\(\downarrow\)) \\ \(\cline{2-10} \(\Delta\)**DivMoE** & & **49.82**\(\pm\)0.11 & 80.03\(\pm\)0.10 & 0.12 (61\(\%\)\(\downarrow\)) & \(\Delta\)**DivMoE** & & **39.97**\(\pm\)0.11 & **70.13**\(\pm\)0.15 & 0.03 (57\(\%\)\(\downarrow\)) \\ \hline \multicolumn{10}{c}{**CIFAR-100**} \\ \hline \(\bullet\) AT (Dense) & & 27.23\(\pm\)0.08 & 58.21\(\pm\)0.12 & 0.54 & \(\bullet\) AT (Dense) & 27.90\(\pm\)0.13 & 57.60\(\pm\)0.09 & 5.25 \\ \(\circ\) AT (Sparse) & & 26.41\(\pm\)0.16 & 57.02\(\pm\)0.14 (74\(\%\)\(\downarrow\)) & \(\circ\) AT (S-Dense) & & 26.30\(\pm\)0.10 & 56.80\(\pm\)0.08 & 1.31 (75\(\%\)\(\downarrow\)) \\ \(\circ\) AT (Sparse) & ResNet-18 & 26.13\(\pm\)0.14 & 57.24\(\pm\)0.12 & 0.14 (74\(\%\)\(\downarrow\)) & \(\circ\) AT (Sparse) & WRN-28-10 & 25.83\(\pm\)0.16 & 57.39\(\pm\)0.14 & 1.31 (75\(\%\)\(\downarrow\)) \\ \(\cline{2-10} \(\Delta\)**DivMoE** & & 22.72\(\pm\)0.42 & 53.34\(\pm\)0.61 & 0.15 (72\(\%\)\(\downarrow\)) & \(\Delta\)**DivMoE** & & 22.92\(\pm\)0.55 & 53.39\(\pm\)0.49 & 1.75 (67\(\%\)\(\downarrow\)) \\ \(\cline{2-10} \(\Delta\)**DivMoE** & & **28.05**\(\pm\)0.13 & **57.73**\(\pm\)0.11 & 0.15 (72\(\%\)\(\downarrow\)) & \(\Delta\)**DivMoE** & & **28.82**\(\pm\)0.14 & **57.56**\(\pm\)0.17 & 1.75 (67\(\%\)\(\downarrow\)) \\ \hline \(\bullet\) AT (Dense) & & 22.37\(\pm\)0.15 & 52.36\(\pm\)0.17 & 0.31 & \(\bullet\) AT (Dense) & & 21.72\(\pm\)0.13 & 48.64\(\pm\)0.14 & 0.07 \\ \(\circ\) AT (S-Dense) & & 20.58\(\pm\)0.13 & **48.89**\(\pm\)0.14 & 0.07 (77\(\%\)\(\downarrow\)) & \(\circ\) AT (S-Dense) & & 16.86\(\pm\)0.21 & 39.97\(\pm\)0.11 & 0.02 (71\(\%\)\(\downarrow\)) \\ \(\circ\) AT (Sparse) & VGG-16 & 21.12\(\pm\)0.22 & 48.03\(\pm\)0.17 & 0.07 (77\(\%\)\(\downarrow\)) & \(\circ\) AT (Sparse) & DenseNet & 17.72\(\pm\)0.14 & 41.03\(\pm\)0.16 & 0.02 (71\(\%\)\(\downarrow\)) \\ \(\circ\) AT (MoE) & & 19.34\(\pm\)0.43 & 45.51\(\pm\)0.75 & 0.12 (61\(\%\)\(\downarrow\)) & \(\circ\) AT (MoE) & & 14.45\(\pm\)0.45 & 36.72\(\pm\)0.71 & 0.03 (57\(\%\)\(\downarrow\)) \\ \(\cline{2-10} \(\Delta\)**DivMoE** & & **21.21**\(\pm\)0.21 & **48.33**\(\pm\)0.17 & 0.12 (61\(\%\)\(\
RA-PGD, AdvMoE still outperforms AT (S-Dense), AT (Sparse), and AT (MoE) consistently, evidenced by the **bold** numbers in the RA-AA columns.

**MoE-CNN trained by AdvMoE enjoys better router utility.** Based on the results above and the preliminary studies in Sec. 4, we next peer into the performance difference achieved by AT (Sparse), AT (MoE), and AdvMoE from the perspective of pathway diversities. We ask:

1. What is the relationship between the dynamic pathways generated by the routers trained by AdvMoE and the static mask optimized by AT (Sparse)? What is the difference between the routing decisions using AdvMoE and AT (MoE), and how does it impact the performance?

Regarding 1, we investigate the cosine similarity between the pathways generated by training methods, either AT (MoE) or AdvMoE, and the static mask found by AT (Sparse). Since the latter can be regarded as a single pathway used for all the data, we term it _'mask pathway'_ in contrast to _'MoE pathway'_. We calculate the intersection of union (**IoU**) score between the MoE pathway and the mask pathway under each testing dataset (the clean or adversarial version). **Fig. 6** presents the IoU distributions based on the clean and adversarial test datasets (**Fig. 6a** for AdvMoE and **Fig. 6b** for AT (MoE)). We remark that a smaller IoU score indicates a larger discrepancy between the MoE pathway and the mask pathway. As we can see, the IoU distribution of AdvMoE vs. AT (Sparse) in Fig. 6a shifts closer to \(0\) compared with Fig. 6b. This observation applies to both standard and adversarial evaluation and suggests that AdvMoE (our proposal) has a better capability than AT (MoE) to re-build input-specific MoE pathways, which are more significantly different from the input-agnostic mask pathway identified by the pruning-based method, AT (Sparse).

Regarding 2, we observe from Fig. 6 that the routers learned by AT (MoE) are more fragile to adversarial attacks compared to AdvMoE, as evidenced by the less intersection area of adversarial data vs. clean data. This is also aligned with **Insight 3** in Sec.4. Moreover, the routing policy learned by AdvMoE is more diverse than AT (MoE), as indicated by the latter's density-concentrated IoU scores. In contrast, the distribution of AdvMoE is dispersed with a smaller peak value. Therefore, regarding the expert utility, AdvMoE is able to assign the inputs to a larger group of pathways than AT (MoE), making better use of experts.

A coupling effect of expert number \(N\) and per-expert model scale \(r\) on AdvMoE.Recall that there exist two key parameters involved in MoE-CNN (**Fig. A1**): _(a)_ the number of experts \(N\), and _(b)_ the model scale \(r\) that defines the per-expert (or per-pathway) model capacity. Given the backbone model (_e.g._, ResNet-18 in this experiment), a larger \(N\) paired with a small \(r\) implies that each expert may only have limited model capacity, _i.e._, corresponding to a less number of channels. Regardless of \(N\), if \(r=1\), the full backbone network will be used to form the identical decision pathway.

**Fig. 7** shows the RA of MoE-CNN trained by AdvMoE vs. the model scale parameter \(r\) at different values of \(N\). Two insightful observations can be drawn. **First**, there exists an MoE regime (_e.g._, \(N<8\) and \(r\in[0.5,0.9]\)), in which AdvMoE can outperform AT (Dense) (_i.e._, \(r=1\)) by a substantial margin. This shows the benefit of MoE in adversarial robustness. However, if the number of experts becomes larger (_e.g._, \(N=10\)), the increasing diversity of MoE pathways can raise the difficulty of routers' robustification and thus hampers the performance of AdvMoE (see \(N=10\) and \(r=0.8\) in **Fig. 7**). **Second**, there exists an

\begin{table}
\begin{tabular}{l|c|c c c c|c|c c c} \hline \hline
**Method** & **Backbone** & **RA-PGD** (\%) & **RA-AA** (\%) & **SA** (\%) & **GFLOPS**(\%) & **Method** & **Backbone** & **RA-PGD** (\%) & **RA-AA** (\%) & **SA** (\%) & **GFLOPS** (\%) \\ \hline \multicolumn{10}{c|}{**AT (Dense)**} \\ \(\sim\)AT (S-Dense) & & 50.13\(\pm\)0.13 & 44.72\(\pm\)0.15 & 82.99\(\pm\)0.11 & 0.54 & \(\sim\)**4** (Dense)** & & 51.75\(\pm\)0.12 & 45.13\(\pm\)0.12 & 83.54\(\pm\)0.15 & 5.25 \\ \(\sim\)AT (S-Dense) & & 48.12\(\pm\)0.09 & 42.24\(\pm\)0.13 & 80.18\(\pm\)0.11 & 0.14 (74\%) & \(\sim\)AT (S-Dense) & & 50.66\(\pm\)0.13 & 44.14\(\pm\)0.10 & 82.24\(\pm\)0.10 & 1.31 (57\%) \\ \(\sim\)AT (Sparse) & ResNet-18 & 47.93\(\pm\)0.17 & 42.11\(\pm\)0.11 & **80.46\(\pm\)0.13** & 0.14 (74\%) & \(\sim\)AT (Sparse) & WRN-28-10 & 48.95\(\pm\)0.14 & 43.97\(\pm\)0.11 & 82.44\(\pm\)0.17 & 1.31 (57\%) \\ \(\sim\)AT (MoE) & & 45.75\(\pm\)0.14 & 40.20\(\pm\)0.18 & 78.84\(\pm\)0.15 & 0.15 (72\%) & \(\sim\)AT (MoE) & & 46.73\(\pm\)0.14 & 44.11\(\pm\)0.23 & 7.42\(\pm\)0.17 & 1.75 (67\%) \\ \multicolumn{10}{c|}{**AbvMoE**} \\ \hline \multicolumn{10}{c|}{**AT (Dense)**} \\ \(\sim\)AT (S-Dense) & & 27.23\(\pm\)0.08 & 23.11\(\pm\)0.06 & 58.21\(\pm\)0.12 & 0.54 & \(\sim\)AT (Dense) & & 27.90\(\pm\)0.13 & 23.45\(\pm\)0.11 & 57.60\(\pm\)0.09 & 5.25 \\ \(\sim\)AT (S-Dense) & & 26.41\(\pm\)0.16 & 22.11\(\pm\)0.13 & 57.02\(\pm\)0.14 & 0.14 (74\%) & \(\sim\)AT (S-Dense) & & 26.30\(\pm\)0.10 & 22.23\(\pm\)0.13 & 56.80\(\pm\)0.08 & 1.31 (75\%) \\ \(\sim\)AT (Sparse) & ResNet-18 & 26.13\(\pm\)0.14 & 21.89\(\pm\)0.11 & 57.24\(\pm\)0.12 & 0.14 (74\%) & \(\sim\)AT (Sparse) & WRN-28-10 & 25.83\(\pm\)0.16 & 21.97\(\pm\)0.09 & 57.39\(\pm\)0.14 & 1.31 (57\%) \\ \(\sim\)AT (MeE) & & 27.22\(\pm\)0.42 & 16.33\(\pm\)0.25 & 53.34\(\pm\)0.61 & 0.15 (72\%) & \(\sim\)AT (MoE) & & 22.94\(\pm\)0.55 & 17.87\(\pm\)0.24 & 53.39\(\pm\)0.49 & 1.75 (67\%) \\ \hline \multicolumn{10}{c|}{**AvMoE**} \\ \(\sim\)**10.13** & **23.33\(\pm\)**0.06 & **57.73\(\pm\)**0.11 & 0.15 (72\%) & \(\sim\)**AvMoE** & & **28.82\(\pm\)**0.14 & **23.75\(\pm\)**0.12 & **57.56\(\pm\)**0.17 & 1.75 (67\%) \\ \hline \hline \end{tabular}
\end{table}
Table 2: Robustness overview evaluated with AutoAttack [68] (**RA-AA**) on various datasets and model backbone architectures. Other settings strictly follow Tab. 1. The values of RA-PGD, SA, and GFLOPS are repeated from Tab. 1 for better comparison.

Figure 6: The distribution of the intersection of union (IoU) scores of the input-specific pathways generated by AdvMoE (a) and AT (MoE) (b) vs. the static mask found by AT (Sparse). The distribution over the clean test set and the adversarial test set is plotted for AT (MoE) and AdvMoE on setting (ResNet-18, CIFAR-100). Other settings are aligned with Tab.1.
_ineffective_ MoE regime (_e.g._, \(N\geq 8\) and \(r<0.5\)), in which the performance of AdvMoE largely deviates from that of AT (Dense). In this regime, each expert consists only of a small number of channels, which restricts its robust training ability. Accordingly, both the increasing diversity of MoE pathways (large \(N\)) and the limited capacity per pathway (small \(r\)) could impose the difficulties of AT for MoE-CNN. In our experiments, we choose \(r=0.5\) and \(N=2\), which preserves the diversity of MoE pathways (_i.e._, inference efficiency) and retains the effectiveness of robust training.

Performance with different model scales.To make sure the observations and conclusions from Tab. 1 are consistent across different values of the model scale parameter \(r\), we repeated the experiments on (CIFAR-10, ResNet-18) and (CIFAR-10, WRN-28-10) using \(r\in\{0.2,0.5,0.8\}\) to cover the {sparse, medium, dense} regimes with respect to Dense (\(r=1.0\)). **Fig. 8** summarizes the obtained experiment results. As we can see, AdvMoE yields consistent robustness improvements over all the baselines, including Dense. And the improvement rises as the model scale \(r\) increases. This is not surprising as more parameters will be used when processing one input. Yet, a clear drawback brought by the larger model scale \(r\) is the increase of inference cost, evidenced by the GFLOPS numbers. When \(r\) turns to be large (like \(r=0.8\)), the efficiency benefit brought by the pathway sparsification from MoE gradually vanishes. Thus, a medium sparsity (\(r=0.5\)) is a better choice to balance the trade-off between performance and efficiency, which is thus adopted as our default setting.

Extended study: AdvMoE for ViT.To explore the capability of our proposal AdvMoE on ViT-based MoE models (MoE-ViT), **Tab. 3** presents additional results following the recently published SOTA baseline [30] for MoE-ViT. As we can see, AdvMoE is also applicable to MoE-ViT and can boost robustness over the SOTA baseline by over \(1\%\) RA improvement, while achieving a similar level of SA. Thus, although our work focuses on robust training for MoE-CNN, it has the promise of algorithmic generality to other MoE-based architectures. We defer a more comprehensive study in the future.

Additional experiments.We conduct ablation studies on (1) robustness evaluation using AutoAttack [68] (consistent findings can be drawn as PGD attacks), (2) attacks steps used in AT, and (3) additional explorations towards the coupling effect between the number of experts and the model scales. We refer readers to Appendix C for detailed results.

## 6 Conclusion

In this work, we design an effective robust training scheme for MoE-CNN. We first present several key insights on the defense mechanism of MoE-CNN by dissecting adversarial robustness through the lens of routers and pathways. We next propose AdvMoE, the first robust training framework for MoE-CNN via bi-level optimization, robustifying routers and pathways in a cooperative and adaptive mode. Finally, extensive experiments demonstrate the effectiveness of AdvMoE in a variety of data-model setups. Meanwhile, we admit that the AdvMoE requires roughly twice the computational capacity compared to the vanilla AT baseline due to alternating optimization that calls two back-propagations per step. Addressing this efficiency concern presents a meaningful avenue for future work.

## Acknowledgement

The work of Y. Zhang, S. Chang and S. Liu was partially supported by National Science Foundation (NSF) Grant IIS-2207052 and Cisco Research Award. The work of Z. Wang is in part supported by the US Army Research Office Young Investigator Award (W911NF2010240).

Figure 8: Robustness comparison of models trained with different methods under various model scale settings. Results higher than that of AT (Dense) are marked with \(\star\). Other setups are aligned with Tab. 1. Please refer to Appendix C for exact numbers and GFLOPS comparisons.

\begin{table}
\begin{tabular}{c|c c c} \hline \hline Method & RA (\%) & SA (\%) & GFLOPS (\#) \\ \hline SOTA[30] & 44.63 & 61.72 & 0.27 \\ AdvMoE & **45.93** & 61.67 & 0.27 \\ \hline \hline \end{tabular}
\end{table}
Table 3: Performance on robust training for MoE-ViT with the setup (ImageNet, DeiT-Tiny). Other settings follow Tab. 1.

Figure 7: Performance of AdvMoE under CIFAR-10 using ResNet-18 as the backbone network for different values of expert number \(N\) and model scale \(r\). The black dash line denotes the performance of Dense (_i.e._\(r=1\)).