On the Adversarial Robustness of Mixture of Experts
Anonymous Author(s)
Abstract
Adversarial robustness is a key desirable property of neural networks.
It has beenempirically shown to be affected by their sizes, with larger networks being typicallymore robust.
Recently, Bubeck and Sellke [4] proved a lower bound on the Lipschitzconstant of functions that ﬁt the training data in terms of their number of parameters.
This raises an interesting open question, do—and can—functions with moreparameters, but not necessarily more computational cost, have better robustness?
We study this question for sparse Mixture of Expert models (MoEs), that make itpossible to scale up the model size for a roughly constant computational cost.
Wetheoretically show that under certain conditions on the routing and the structureof the data, MoEs can have signiﬁcantly smaller Lipschitz constants than theirdense counterparts.
The robustness of MoEs can suffer when the highest weightedexperts for an input implement sufﬁciently different functions.
We next empiricallyevaluate the robustness of MoEs on ImageNet using adversarial attacks and showthey are indeed more robust than dense models with the same computational cost.
We make key observations showing the robustness of MoEs to the choice of experts,highlighting the redundancy of experts in models trained in practice.
Introduction
Adversarial robustness refers to prediction robustness of a given machine learning model to adversarial,but bounded, changes to the input.
Neural networks trained with standard classiﬁcation objectives havebeen shown to have poor adversarial robustness [24], a property attributed to their overparametrization.
Conversely, in practice, larger models, that have more parameters and higher computation cost, haveshown better robustness [19, 25, 1, 12].
In a recent work
Bubeck and Sellke [4] studied this phenomenon from a theoretical perspective,by analyzing the relationship between model size and its Lipschitz constant - which measures thesensitivity of a function to changes in the input.
In particular the authors proved that any functionwith P parameters, that memorizes N input data points in D dimensions, has a Lipschitz constantof at least O ( (cid:113) N · DP ) .
This shows that, on a given dataset, larger models can have better robustness(smaller Lipschitz constant).
Note that as this is only a lower bound, it does not guarantee that largermodels will indeed have a smaller Lipschitz constant.
Interestingly, the result is agnostic to otherproperties of the function, such as its computation cost, or speciﬁc architecture.
Given the above result it is natural to wonder: can one increase the model size, without increasingthe computation cost, and enjoy better robustness?
An increasingly popular class of such models are
Mixture of Experts (MoE) [22].
MoE models have multiple expert sub-models, and a routing function that selects for each input a small subset of the experts and routes the input only to them.
Neuralnetworks with MoE layers, such as Switch Transformer [10] in NLP, and V-MoE [21] in Computer
Vision, have been shown to achieve superior performance in comparison to their dense counterparts,by allowing one to scale the model size without increasing the computation cost.
In this paper westudy the following question: are MoE models more adversarially robust than dense models?
In general, sparse MoE models are not continuous, and hence not even smooth.
Small changes tothe input can result in the router selecting a different expert for a given example.
Unfortunately, insome cases these experts may be very different, resulting in large changes to the output.
Anotherfactor that affects the robustness of MoE models is the geometry of input data and its routing –howthe data is divided between experts.
This decides what data an expert is trained on and hence theirrobustness.
Finally, it has been observed that –unless certain auxiliary losses are also applied toencourage balancedness– the number of used experts tends to collapse to very few, and the remainingones are just ignored [21, 10].
Given these stability issues, it is a priori not clear if MoEs, despitehaving more parameters, are more robust than dense models.
To theoretically study the effect of both of these factors, (1) router stability and (2) data routing,on robustness we consider the smooth MoE models, where each expert output is weighted by theirrouting probability.
In particular, MoE models where all experts are simultaneously applied to everyinput.
We analyze models with ﬁxed routing proving that MoEs with linear experts can achieve betterrobustness if the data is well separated and routed accordingly.
In the extreme case, MoEs with Eexperts can have a smaller Lipschitz constant by a factor of 1 / √ E compared to equivalent densemodels.
We also characterize the effect of the difference between experts in terms of robustness,showing that MoEs can have a high Lipschitz constant when two experts are very different forinputs that weigh the experts similarly and highly.
We show that both these factors, data routing anddifference between relevant experts, characterize the Lipschitz constant of MoE models.
We next evaluate robustness experimentally by using adversarial attacks [19].
We experimentallyshow that on the ImageNet dataset [6], MoEs are more robust than dense models that have thesame computation cost.
Interestingly, we observe that adversarial attacks result in signiﬁcant routingchanges of the inputs, but do not result in lower robustness than dense models.
This suggests thatin practice, MoEs are robust to the choice of the experts to an extent.
We next perform standardadversarial training of both the dense and the MoE models and again observe that MoEs are morerobust against adversarial attacks.
The main contributions of the current work are as follows:
We propose a simple theoretical framework to understand the robustness of Mixture of
Experts (MoE) models.
We provide general and non-trivial sufﬁcient conditions under which
MoEs are provably more robust than their dense counterparts for the linear experts setting.
We perform extensive experiments to demonstrate that in practice MoEs indeed enjoy
better robustness than dense models to norm bounded adversarial attacks.
We also uncover
intriguing properties of adversarial attacks for MoE models and show that even for robust MoE models, the attacks very often change the routing of data points, thereby pointing to ahigh degree of redundancy in such models.
Preliminaries2.1 Mixture of Experts
Mixture of Experts combine the outputs of (sub-)models, i.e., the experts , via a weighted sum of theiroutputs [14, 15, 27, 9].
Sparse Mixture of Experts condition the weights of the sum in the inputs, andactivate only K (out of E ) experts, where K is typically a very small integer compared to E (suchas K = 1 or K = 2 ) [22, 10].
This form of conditional computation models allows one to easilyincrease the number of parameters in the model (roughly) independently of its compute cost [20].
This approach has been recently applied to signiﬁcantly increase the model size and quality of modelsused in Natural Language Processing [22, 17, 18, 10] and Computer Vision applications [21, 26].
More concretely, we deﬁne the sparse and smooth (or dense) versions of mixtures of experts below.
While previous empirical works have sometimes selected more than one expert per input ( K > 1 ),for simplicity we constrain ourselves to the case where only one expert is selected for each input( K = 1 ).
Recently, practical models also tend to this setup to match the FLOPs of dense models [10].
Sparse MoEs.
In these models only one expert is selected for a given example x based on its routingprobabilities p i ( x ) .
These models are not continuous functions of the input.
A small change in theinput can lead to large changes in the output, due to changes in selected experts.
A sparse MoE layer is deﬁned as:
[Equation]
Here f i ( x ) are the individual expert functions.
1 is the indicator function that takes the value of1 if the condition is satisﬁed and 0 else.
Note that to do tie breaking in the above deﬁnition, incase multiple experts have the same maximum probability, we simply sample uniformly one of theexperts with the maximum probability.
In practice, MoE models can have multiple sparse layers incombination with dense layers.
Further, random noise is sometimes added to the p
i ’s
Smooth MoEs.
To theoretically analyze robustness, we consider a smooth mixture of expert models,as they are continuous functions.
In this case, the output of each expert f i , is weighted according toits routing probabilities:
[Equation]
The routing probabilities p
i ( x ) are usually computed by a routing layer, e.g., p i ( x )
= σ ( Sx ) i , whereσ is the softmax and S ∈ R E
× D is a trainable variable.
Such linear routing layers are the commonchoice for MoE models in practice [10].
Though these models have recently been shown to achieve state of art performance for tasks in
NLP [10] and vision [21], there have not been many works analyzing their robustness.
Recently,
Allingham et al.
[2] empirically studied robustness of MoEs to natural perturbations in data (e.g.,evaluation across different corrupted versions of Cifar10 [16] and ImageNet [6]), and showed MoEsare more robust for these natural perturbations than the corresponding dense models.
To the best ofour knowledge ours is the ﬁrst work theoretically analyzing robustness of MoEs, and empiricallyexploring it in the context of adversarial perturbations.
2.2 Adversarial robustness
Adversarial robustness models the susceptibility of a function to adversarial perturbations to itsinput [24, 19].
More concretely, given a function f , loss function (cid:96) , and an input x ∈ R D , we canwrite the adversarial loss incurred at x with label y as follows.
[Equation]
Here (cid:15) is the attack radius per input.
For the norm constraint ( (cid:107) z (cid:107) ), popular choices in practice, are(cid:96) ∞ and (cid:96) 2 norms [19, 5].
Adversarial accuracy is the accuracy of the model on data with perturbationsthat satisfy the above equation, i.e., when the loss function (cid:96) is the 0 / 1 classiﬁcation loss.
Since optimizing equation 3 for neural networks is computationally hard, popular methods to ﬁndadversarial examples use local search approaches such as gradient ascent.
The Fast Gradient Signed
Method (FGSM) is perhaps the simplest one, with only one gradient update step [11].
For (cid:96) ∞ boundedperturbations the FGSM update is the following:
[Equation]
sgn() is the sign function.
The Projected Gradient Descent (PGD) method is a stronger attack thatdoes multiple ( τ ) gradient ascent steps to ﬁnd the perturbation [19].
The update rule for (cid:96) ∞ boundedperturbations is as follows, starting with x 0 = x .
[Equation]
α is chosen to be (cid:15)τ .
Here C is the constraint set { z : (cid:107) z (cid:107) ∞ ≤ (cid:15) } and Π is the projection operator.
Existing works have established that neural networks are sensitive to adversarial attacks resultingin a signiﬁcant drop in their accuracy [24].
Interestingly, increasing the size of neural networks,thereby increasing their capacity, leads to an improvement in adversarial accuracy, showing that largerneural networks are more robust [19, 25, 1, 12].
In a recent work, Bubeck and Sellke [4] studied thisphenomenon theoretically.
They proved a lower bound on the Lipschitz constant of any function thatﬁts the training data that scales inversely 1 √ P with the number of function parameters P .
However it was not clear if functions that have more parameters, but not necessarily more computation cost (e.g.
MoEs), can achieve better robustness.
In this paper we analyze adversarial robustness of MoEs, boththeoretically and experimentally.
Notation.
We use small bold letters x to denote vectors and capital bold letters to denote matrices
W .
Scalars are denoted with plain letters.
[ N ] denotes the integer set from 1 to N .
(cid:107) .
(cid:107) denotes the (cid:96) 2norm unless speciﬁed otherwise.
Robustness analysis
In this section, we present our main results regarding the robustness of the mixture of expert models.
We ﬁrst present a bound on the Lipschitz constant of MoEs for general expert functions f i .
Whilethis is an upper bound for general functions it still highlights the two key components that affect therobustness of MoE models.
3.1 Router stability
In this section, we compute the Lipschitz constant of smooth MoEs with learnable routing.
Let,
[Equation]
p
[Equation]
x
[Equation]
f i ( x ) where p i ( x ) = [Equation] x
[Equation] [Equation] and { s i ∈ R D } i ∈ [ E ] are learnable variables that decide the routing of an example to different experts.
We then prove the following upper bound on the Lipschitz constant of the MoE models with learnablerouting.
Lemma 1.
Let { f i } i ∈
[ E ] be smooth functions with Lipschitz constants { L f i } i ∈
[ E ] and let L f be the
Lipschitz constant of f .
Let ¯ s ( x )
=
(cid:80) j p j ( x ) s j .
Then f ( x ) in equation 6 satisﬁes,
[Equation]
and hence
[Equation]
The above Lemma bounds the Lipschitz constant of MoE models with two terms that depend on (1)data routing, and (2) router stability.
The ﬁrst term above is from the individual Lipschitz constant ofthe experts.
This depends on how the data is partitioned or routed to different experts, which we willdiscuss in more detail in the next section.
The second term arises from the router used to computeprobabilities for different experts.
One may wonder if we can make this term arbitrarily large byincreasing the norm of x .
This is not the case as increasing the norm of x usually results in collapse ofthe routing probabilities to a single expert.
This makes p i ( x )( s i − ¯ s ( x )) zero for all experts i ∈
[ E ] .
To study the above more concretely, we consider the case of E = 2 experts.
In this setting the secondterm above reduces to (cid:107) p 1 ( x ) f 1 ( x )( s 1 − ¯ s ( x ))
+ (1 − p 1 ( x ))
f 2 ( x )( s 2 − ¯ s ( x ))
(cid:107) .
Now, in theextreme case, where p 1 ( x ) is either 0 or 1, the above term collapses to
as p i ( x ) f i ( x )( s i − ¯ s ( x ))
is0 in that case.
This is also expected as for examples that are routed to an expert with high probability,the main dominating factor is the individual expert Lipschitz constant, as small perturbations do notcause much change to the probabilities.
A more interesting setting is when p 1 ( x ) is 12 .
In this setting, small perturbations to the input cancause the model to drastically change the weight between the experts, and the above term reduces to(cid:107) 1 4 · ( f 1 ( x ) − f 2 ( x )) ·
( s 1 − s 2 ) (cid:107) .
Hence MoE models can suffer from large Lipschitz constants ifthe two experts are very different for points on the boundary, i.e., if | f 1 ( x ) − f 2 ( x ) | (cid:29) 0 .
Alternatelyif f 1 ( x ) ≈ f 2 ( x ) for points on the boundary (i.e. p 1 ( x ) ≈ 12 ) then this term is small.
We will later see in the experiments that MoEs trained in practice often have good router stability andchanges to routing do not result in much degradation of accuracy.
3.2 Data routing
In this section we will study the effect of data routing on the Lipschitz constant of the individualexperts – the ﬁrst term in RHS of equation 7 (lemma 1).
In particular we will try to address howlarge can max i ∈
[ E ] { L f i } be in comparison to the Lipschitz constant of a dense model trained on 170 the same data.
Towards this end, we will theoretically analyze the robustness of MoE models with apre-determined ﬁxed routing and linear experts.
3.2.1 Setup
We consider the setting of ﬁxed routing, where we assume the routing of individual examples toexperts is pre-determined and ﬁxed.
In particular we also consider linear models as experts.
More speciﬁcally, we consider N input points of dimension D stacked in the matrix X ∈ R N × D ,together with the corresponding vector of targets y ∈ R N .
In the dense case (i.e., no experts), wewould like to ﬁnd a linear model w ∗ ∈ R D that minimizes
[Equation] [Equation]
The least-squares solution [13] for this problem that minimizes the MSE loss is w ∗ =
( X
[Equation]
Here † denotes the pseudo inverse.
On the other hand, for mixture of experts, let the dataset be splitinto E subsets S 1 , · · · , S E , with X i , y i denoting the data in set S i .
Now let that expert i be routedthe data from set S i .
Below we write the objective for each expert:
[Equation]
Similarly, let w ∗ i =
( X [Equation] i X i ) † X (cid:62) i
y i be the optimal solution for expert i that minimizes the MSE.
3.2.2 Analysis
Before introducing our analysis relating the Lipschitz constant of the dense and MoE models, we ﬁrstpresent two examples to better understand how routing affects Lipschitz constants.
In particular weconsider a simple setting with 2 experts.
Let X (cid:62) =
[ X (cid:62) 1 , X (cid:62) 2 ] , be the data routed to the two experts.
Recall that the Lipschitz constant of the function f ( X ) =
Xw is (cid:107) w (cid:107) .
Case X 1 ⊥ X 2 .
In this setting we get
[Equation]
as X 1 and X 2 are in orthogonal subspaces.
Hence, experts have smaller Lipschitz constant whendata routed to different experts lies in orthogonal subspaces.
Case X 1 = − X
and y 1 =
y 2 .
In this case, we see that:
[Equation]
Hence experts have worse Lipschitz constant when the data routed to different experts is aligned.
These two simple examples illustrate under which conditions experts have an advantage over a singledense model, and when they do not.
We now present our main result.
To capture this relation between the data geometry and the routing,we introduce the following quantities.
Let { U i } i ∈
[ E ] be the projection matrices onto orthogonalsubspaces in R D .
One way to construct them is by ﬁrst taking the singular vectors of X (cid:62) X andassigning them to the set i with the largest projection.
This guarantees that U
i is orthogonal to U jwith j (cid:54) =
i , and they have greatest alignment with subspace spanned by X i .
We ﬁrst deﬁne a quantity to capture how well U i captures the span of the data subset X i .
Deﬁnition 1 (In-subspace distance: (cid:15) 1 ) .
∃ (cid:15) 1 ≥ 0 , (cid:107) U i U (cid:62) i
( X (cid:62) X ) † − ( X (cid:62) i X i ) † (cid:107) 2 ≤ (cid:15) 1 , ∀ i ∈
[ E ] .
Here (cid:107) .
(cid:107) 2 for a matrix denotes the spectral norm.
(cid:15) 1 is small when, ∀ i, U i captures X i perfectly.
Next we deﬁne the projection distance between data from two different subsets.
Deﬁnition 2 (Cross-subspace distance: (cid:15) 2 ) .
∃ (cid:15) 2 ≥ 0 such that for any z in the span of X (cid:62) j , we have
[Equation]
Here (cid:15) 2 is small if the data in different subsets X i lies in orthogonal subspaces.
Theorem 1.
Let w ∗ be the minimizer of equation 8 and { w ∗ i } be the minimizers of equation 9 .
Then,
[Equation]
where (cid:98)·(cid:99) + denotes the projection onto non-negative numbers.
The above result lower bounds the Lipschitz constant of the dense model (cid:107) w ∗ (cid:107) in terms of Lipschitzconstants of the experts (cid:107) w ∗ i (cid:107) in the MoE model.
We present the proof of this theorem in appendix A.
In the case where (cid:15) 1 = (cid:15) 2 = 0 , we obtain the following bound (cid:107) w ∗ (cid:107) ≥ (cid:113)(cid:80)
Ei =1 (cid:107) w ∗
i (cid:107) 2 .
Assuminga balanced setting where all experts have the same norm parameters, we get the scaling (cid:107) w ∗ (cid:107) ≈
O ( √ E ) · max i ∈
[ E ] {(cid:107) w ∗ i (cid:107)} .
Hence in this setting, experts have a Lipschitz constant that is smallerby a factor of 1 √ E compared to dense models.
This happens when the data lies in E orthogonalsubspaces and the data from each subspace is routed to the same expert.
This shows that MoE modelscan have signiﬁcantly smaller Lipschitz constant than their dense counterparts, while having the samecomputation cost.
For MoEs in practice, data routed to different experts does display some clusteringof the features (see Figure 7 in Riquelme et al.
As data in different partitions X i gets more aligned, (cid:15) 1 and (cid:15) 2 increase, and reduce the gap betweenthe Lipschitz constant of the dense and the MoE models.
This reduces the gap in the Lipschitzconstant of the dense model and the experts.
In the extreme case, if all the datapoints are the same,then (cid:15) 1 and (cid:15) 2 are large, eliminating this difference.
In appendix E, we give a complementary resultrelating the Lipschitz constants of the experts to that of the dense model when both (cid:15) 1 and (cid:15) 2 canpossibly be large.
Connections to Bubeck and Sellke [4].
Bubeck and Sellke [4] proved a universal lower bound on the Lipschitz constant of a function required to δ -memorize N training samples in D dimension.
[Equation]
For the linear model considered above, this reduces to ˜Ω (cid:16) δ √ N (cid:17) as the number of parameters is
D .
As MoEs have E times more parameters, they have a lower bound of ˜Ω (cid:16) δ (cid:113) NE (cid:17) , i.e., the lowerbound on the Lipschitz constant of MoEs is smaller by a factor of (cid:113) 1 E .
Our result shows that thislower bound is in fact achievable, and hence tight, when the data lies in E orthogonal subspaces anddata from each subspace is routed to the same expert.
Experiments
[Figure] 4.1 Setup
We compare the robustness of Vision Transformer (ViT) [8] and Vision MoE (V-MoE) [21] modelsagainst adversarial attacks.
In particular, we use the ViT-B/32 and V-MoE-B/32 models through allthe experiments.
These models have the same backbone architecture but the latter replaces one inevery two feedforward layers with a sparse Mixture of Experts, selecting K = 2 out of E
= 32feedforward experts that are applied on each token 1 .
Since one could argue that the router in the
V-MoE model introduces an overhead that should be accounted for, we have also trained a biggerversion of the dense ViT-B/32 (which we coin ViT-B++/32) that has roughly the same quality asthe V-MoE model, but higher GFLOP and runtime cost.
The compute needed to evaluate an imageon the dense ViT models are 8.9 and 17.9, respectively; and 12.4 GFLOPs on the V-MoE model.
Appendix B contains all the details of the architectures used in the experiments.
F a l s e d i s c
o
v e r y r a t e
JFT-300M 1010100.
C l a ss i f
i c a t i o n e rr o r
r
a t e
ImageNet
Model
ViT (8.9 GFLOP)
ViT++ (17.9 GFLOP)
V-MoE (12.4 GFLOP)
Adversarial ft.
False
True
We pre-train our models on the private dataset JFT-300M [23] for 7 epochs (517 859 steps with a batchsize of 4 096 images), using an image resolution of 224 × 224 pixels, and standard data augmentation(inception crop and horizontal ﬂips).
Since JFT-300M is a multi-label dataset, we minimize thesigmoid cross-entropy loss.
V-MoE also adds auxiliary losses to encourage a balanced load for allexperts; we used the same recipe as in [21].
In both cases we use Adam ( β 1 = 0 .
9 , β 2 = 0 .
,with a peak learning rate of 8 · 10 − 4 , reached after a linear warm-up of 10 4 steps and then linearlydecayed to a ﬁnal value of 10 − 5 .
Weight decay of 0.1 was used on all parameters.
This is the samepre-training protocol used in [8, 21].
After pre-training, the models are ﬁne-tuned on ImageNet [6], at a resolution of 384 × 384 pixels andthe same data augmentations as before, for a total of 10 4 steps, using a batch size of 4 096 images.
SGD with Momentum ( µ = 0 . 9 ) is used for ﬁne-tuning, with a peak learning rate of 0 . 03 , reachedafter a linear warm-up of 500 steps, and followed with cosine decay to a ﬁnal value of 10 − 5 .
Thenorm of the ﬂattened vector of gradients is clipped to a maximum value of 10.
Since ImageNetimages have a single label, we minimize the softmax cross-entropy loss during ﬁne-tuning.
The sameﬁne-tuning protocol was adopted in [8, 21].
We evaluate the adversarial robustness of both the pre-trained and ﬁne-tuned models, by means of PGDadversarial attacks [19].
We maximize the corresponding loss (sigmoid or softmax cross-entropy),varying the (cid:96) ∞ norm constraint on the input image, for a total of τ = 40 steps.
4.2 Adversarial robustness of V-MoEs
During pre-training the ViT-B/32 model achieves a precision at 1 of 39.3%, and the V-MoE-B/32achieves a precision-at-1 of 43.5% (conversely, the false discovery rate at 1 is 60.7% and 56.5%,respectively).
After ﬁne-tuning, the classiﬁcation error achieved by each model on ImageNet is 19.3%and 17.8%, respectively.
Figure 1 shows in solid lines the false discovery rate (left) and the classiﬁcation error rate (right) as afunction of the (cid:96) ∞ constraint.
Despite the fact that the V-MoE model contains a router that makesdiscrete choices among the experts conditioned on the input, which could potentially lead to a severeweakness against the adversarial attacker, we can observe that it follows the same trend as the basedense ViT model, and it is able to preserve a lower error over a wide range of (cid:96) ∞ values.
A largerversion of the dense model matching the quality of the V-MoE has a much higher cost.
In addition, we also ﬁne-tuned the base ViT and V-MoE models on ImageNet using PGD adversarialtraining.
We use the same recipe as above but we perform a PGD attack of 10 steps on the input images,with ﬁxed (cid:96) ∞ = 8255 , before computing the gradients of the model parameters and updating them.
The classiﬁcation error on the original ImageNet dataset achieved by each model after adversarialﬁne-tuning is 51.7% and 49.8%, for both models respectively.
Figure 1 (right) shows in dashed linesthe classiﬁcation error when these models are evaluated against an adversarial attacker using different(cid:96) ∞ constraints.
As the (cid:96) ∞ increases, the beneﬁt of adversarial ﬁne-tuning to preserve accuracy is shown in both cases.
Once again, both models report similar trends, and the V-MoE model showsbetter robustness for a wide range of (cid:96) ∞ values.
4.3 Effect of the adversarial attacks on the selected experts
As described in section 3.1, in the region close to the decision boundary of the router, if two expertshave very different outputs, the Lipschitz constant of the MoE model could be much higher than thatof a similar dense model.
If the index of the selected experts changes signiﬁcantly, but the model stillshows a reasonably high accuracy under adversarial attacks (compared to a dense model), this wouldsuggest that the outputs of the two selected sets of experts are not that different.
Figure 2 shows the rate of changes in the router as a function of the (cid:96) ∞ used in the adversarial attack,on the different layers of the V-MoE-B/32 model that have a MoE layer.
For each token processed bythe model, we compute the intersection-over-union (IoU) of the selected set of experts before andafter the adversarial attack.
We average the IoU across all processed tokens and deﬁne the rate ofrouting changes as the complement of the average IoU.
In dashed lines we report the rate of changes of a V-MoE model ﬁne-tuned using PGD adversarialtraining.
Not only the model has a better accuracy against adversarial attacks when the (cid:96) ∞ increases,as reported in ﬁgure 1, but the rate of routing changes is also generally lower across all layers.
[Figure]
∞ used in the adversarial attacks on ImageNet against a V-MoE model. The dashed line depicts a V-MoE model ﬁne-tuned with PGD adversarial training.
A signiﬁcant fraction of experts change in each MoE layer .
Given that ﬁgure 1 shows the V-MoE model keeping its advantage over the dense ViT, we hypothesize that the output of different experts for inputs close to the decision boundary of the router is similar.
Despite the fact that a signiﬁcant fraction of choices change on each layer as the (cid:96) ∞ increases,ﬁgure 1 shows that the V-MoE model still keeps its advantage against adversarial attacks over the
ViT model.
This suggests that in the regions close to a decision boundary of the router, the twocorresponding experts have a similar output, hence preventing the full V-MoE model from beingless robust in practice.
Conversely, the fact that the V-MoE model has a better base quality than the
ViT counterpart, suggests that the experts are not equivalent for regions far away from the decisionboundary, otherwise it would be reduce to a dense model.
4.4 Attacking the router’s auxiliary losses
Sparse MoE models usually employ auxiliary losses to balance the load among all experts.
Inparticular, in the implementation of V-MoEs, if the load of the experts is highly unbalanced, theexperts receiving signiﬁcantly more tokens than the average could ignore all tokens that exceed theexpert’s capacity, potentially leading to a signiﬁcantly worse performance.
The question is whetheran adversarial attacker can exploit this in practice.
Figure 3 (left) shows the false discovery rate on
JFT-300M, for a V-MoE model when the router’s auxiliary losses are maximized in the adversarialattack, together with the corresponding cross-entropy loss.
We use the same weight for the auxiliarylosses as the one used to train the models.
The ﬁgure shows that attacking the auxiliary loss does notoffer any signiﬁcant advantage for the attacker.
The results are analogous in terms of the classiﬁcationerror on ImageNet (not shown here in interest of space).
Increasing the model size by increasing the total number of experts
In section 1 we asked if, given that the lower bound on the Lipschitz constant given by Bubeck and
Sellke [4] is agnostic to the cost of the function, could we make a model more robust by increasing its [Figure] model size without increasing the total cost?
Here we measure how much increasing the total numberof experts in a V-MoE model helps against adversarial attacks.
Notice that, indeed, increasing thetotal number of experts E does not make the V-MoE model more expensive.
Figure 3 (center andright) show the false discovery rate on JFT-300M and the classiﬁcation error on ImageNet for anincreasing number of experts.
All V-MoE models select K = 2 experts for each token.
On the one hand, increasing the total number of experts improves the robustness on JFT-300M up to
E = 16 experts.
The curves for E = 16 and E = 32 are highly overlapping, thus any difference ismost likely due to noise in the training and ﬁne-tuning process.
On the other hand, when the V-MoEsare ﬁne-tuned on ImageNet, all models with more than two experts achieve roughly the same accuracyunder adversarial attacks.
This shows that, although the results presented in section 3 showing better robustness of MoEs requiresome assumptions, the conclusions hold to some extent in real scenarios.
Thus, increasing the numberof parameters by growing the number of experts in a model with sparse MoE layers is an effectiveway of improving the model’s robustness.
Conclusion
In this work we analyzed the adversarial robustness of MoEs showing their advantage over densemodels, with more experts leading to better robustness, both theoretically and empirically.
We showedhow the properties of the data and its routing plays and important role in learning robust MoEs.
While there is some evidence that routing learned by MoEs in practice display some clustering ofthe features in some layers of the model (see Figure 7 in Riquelme et al.
[21]), it is currently notexplicitly encouraged during training.
Hence developing smarter routing strategies that take datageometry into account can be an interesting direction of future work.
Currently our analysis is limitedto linear models, extending this to general models and deriving the dependency of optimal routing onthem is another interesting research direction.
We have also shown that, for inputs that weigh two experts similarly, if the two expert values arevery different, then the MoEs can suffer from higher Lipschitz constant.
However, for models trainedin practice, we saw their predictions to be relatively stable, despite signiﬁcant changes in choice ofexperts, highlighting potential redundancy of learned experts.
However too much redundancy, withall experts learning similar functions is a waste of capacity and can affect model performance.
Henceit is an interesting research problem to balance robustness and accuracy of MoEs by controlling theredundancy of experts.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’scontributions and
scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes] Please see sections 1 and 5.
(c) Did you discuss any potential negative societal impacts of your work?
[N/A]
Our work
analyzes adversarial robustness of popular models
(d) Have you read the ethics review guidelines
and
ensured
that
your paper
conforms to
them?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes] Please see sections sections 3.1 and 3.2.
b) Did you include complete proofs of all theoretical results?
[Yes] Please see appendix A.
If you ran experiments...
a) Did you include the code, data , and instructions needed to reproduce the main experimental results (
either in
the supplemental material or as a URL
[No] The pre-training
data is proprietary and cannot be publicly shared.
Appendix B includes a description of
the software stack used and a link to the software that we used to run all experiments.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] Please see section 4.
(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[No] We notice fairly stable adversarial accuracy metrics acrossdifferent settings.
(d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
[Yes] The hardware speciﬁcations and the total compute used in the experiments are detailed in appendix B.
If you are using existing assets (e
g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes] Please see section 4.
(b) Did you mention the license of the assets?
[No] Part of the data that we use is propri-
etary.
The public data that we use (ImageNet) is publicly available after registration on
https://image-net.org , and broadly used by the Machine Learning and ComputerVision communities.
(c) Did you include any new assets either in the supplemental material or as a URL?
[No]Part of code we use is proprietary.
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identiﬁable
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable?
[N /A
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount
spent on participant compensation?
[N/A]