On the Adversarial Robustness of Mixture of Experts
Abstract
Adversarial robustness is a key desirable property of neural networks.
It has been empirically shown to be affected by their sizes, with larger networks being typically more robust.
Recently, Bubeck and Sellke [3] proved a lower bound on the Lipschitz constant 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 more parameters, but not necessarily more computational cost, have better robustness?
We study this question for sparse Mixture of Expert models (MoEs), that make it possible to scale up the model size for a roughly constant computational cost.
We theoretically show that under certain conditions on the routing and the structure of the data, MoEs can have signiﬁcantly smaller Lipschitz constants than their dense counterparts.
The robustness of MoEs can suffer when the highest weighted experts for an input implement sufﬁciently different functions.
We next empirically evaluate the robustness of MoEs on ImageNet using adversarial attacks and show they 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 have been 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, have shown better robustness [19, 25, 1, 11].
In a recent work
Bubeck and Sellke [3] studied this phenomenon from a theoretical perspective, by analyzing the relationship between model size and its Lipschitz constant - which measures the sensitivity of a function to changes in the input.
In particular the authors proved that any function with P parameters, that memorizes N input data points in D dimensions, has a Lipschitz constant of 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 larger models will indeed have a smaller Lipschitz constant.
Interestingly, the result is agnostic to other properties 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 increasing the 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.
Neural networks with MoE layers, such as Switch Transformer [9] 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 we study 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 to the input can result in the router selecting a different expert for a given example.
Unfortunately, in some cases these experts may be very different, resulting in large changes to the output.
Another factor that affects the robustness of MoE models is the geometry of input data and its routing –how the data is divided between experts.
This decides what data an expert is trained on and hence their robustness.
Finally, it has been observed that –unless certain auxiliary losses are also applied to encourage balancedness– the number of used experts tends to collapse to very few, and the remaining ones are just ignored [21, 9].
Given these stability issues, it is a priori not clear if MoEs, despite having 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 their routing probability.
In particular, MoE models where all experts are simultaneously applied to every input.
We analyze models with ﬁxed routing proving that MoEs with linear experts can achieve better robustness if the data is well separated and routed accordingly.
In the extreme case, MoEs with E experts can have a smaller Lipschitz constant by a factor of 1 / √ E compared to equivalent dense models.
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 for inputs that weigh the experts similarly and highly.
We show that both these factors, data routing and difference between relevant experts, characterize the Lipschitz constant of MoE models.
We next evaluate robustness experimentally by using adversarial attacks [19].
We experimentally show that on the ImageNet dataset [6], MoEs are more robust than dense models that have the same computation cost.
Interestingly, we observe that adversarial attacks result in signiﬁcant routing changes of the inputs, but do not result in lower robustness than dense models.
This suggests that in practice, MoEs are robust to the choice of the experts to an extent.
We next perform standard adversarial training of both the dense and the MoE models and again observe that MoEs are more robust 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 a high degree of redundancy in such models.
Preliminaries 2.1 Mixture of Experts
Mixture of Experts combine the outputs of (sub-)models, i.e., the experts , via a weighted sum of their outputs [13, 14, 27, 8].
Sparse Mixture of Experts condition the weights of the sum in the inputs, and activate only K (out of E ) experts, where K is typically a very small integer compared to E (such as K = 1 or K = 2 ) [22, 9].
This form of conditional computation models allows one to easily increase 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 models used in Natural Language Processing [22, 17, 18, 9] 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 [9].
Sparse MoEs.
In these models only one expert is selected for a given example x based on its routing probabilities p
i ( x ) .
These models are not continuous functions of the input.
A small change in the input 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 of 1 if the condition is satisﬁed and 0 else.
Note that to do tie breaking in the above deﬁnition, in case multiple experts have the same maximum probability, we simply sample uniformly one of the experts with the maximum probability.
In practice, MoE models can have multiple sparse layers in combination with dense layers.
Further, random noise is sometimes added to the p
i ’s
We refer to MoEs with only 1 expert as dense models.
In dense models, which are the standard in neural networks literature, same function is applied to all input examples.
In sparse MoEs introduced above, different part of network / expert is applied to each input example depending on which expert selected.
This ﬂexibility allows us to train larger sparse MoEs, that have the same computational cost as the dense models, as only a part of MoE is activated/selected for each input, with better accuracy than dense models [9, 21].
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 to its 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 common choice for MoE models in practice [9].
Though these models have recently been shown to achieve state of art performance for tasks in
NLP [9] 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 MoEs are more robust for these natural perturbations than the corresponding dense models.
To the best of our knowledge ours is the ﬁrst work theoretically analyzing robustness of MoEs, and empirically exploring it in the context of adversarial perturbations.
Load balancing loss To prevent MoEs from collapsing and always selecting the same expert for all inputs, it is customary during training to add a load balancing loss , that encourages equal fraction of inputs being routed to different experts.
We deﬁne and present these losses in detail in appendix E.
2.2 Adversarial robustness
Adversarial robustness models the susceptibility of a function to adversarial perturbations to its input [24, 19].
More concretely, given a function f , loss function (cid:96) , and an input x ∈ R D , we can write 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, 4].
Adversarial accuracy is the accuracy of the model on data with perturbations that 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 ﬁnd adversarial 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 [10].
For (cid:96) ∞ bounded perturbations the FGSM update is the following:
[Equation] [Equation]
() is the sign function.
The Projected Gradient Descent (PGD) method is a stronger attack that does multiple ( τ ) gradient ascent steps to ﬁnd the perturbation [19].
The update rule for (cid:96) ∞ bounded perturbations is as follows, starting with x 0 = x .
[Equation] [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 resulting in 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 larger neural networks are more robust [19, 25, 1, 11].
In a recent work, Bubeck and Sellke [3] studied this phenomenon 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 is 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, both theoretically and experimentally.
Notation.
We use small bold letters, e.g., x , to denote vectors and capital bold letters to denote matrices, e.g., W .
Scalars are denoted with plain letters.
[ N ] denotes the integer set from 1 to N .
(cid:107) .
(cid:107) denotes the (cid:96) 2 norm 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 .
While this is an upper bound for general functions it still highlights the two key components that affect the robustness of MoE models.
3.1 Router stability
In this section, we compute the Lipschitz constant of smooth MoEs with learnable routing.
Let, f ( x ) =
E (cid:88) i
= p
i ( x ) f i ( x ) where p
i ( x ) = exp( (cid:104) s i , x (cid:105) )
(cid:80) Ej =1 exp( (cid:104) s j , x (cid:105) )
[Equation]
and { s i } i ∈
[ E ] , with each s i in R D , 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 learnable routing.
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 of the experts.
This depends on how the data is partitioned or routed to different experts, which we will discuss in more detail in the next section.
The second term arises from the router used to compute probabilities for different experts.
One may wonder if we can make this term arbitrarily large by increasing the norm of x .
This is not the case as increasing the norm of x usually results in the collapse of the 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 second term 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 the extreme 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 )) is 0 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 not cause much change to the probabilities.
A more interesting setting is when p 1 ( x ) is 12 .
In this setting, small perturbations to the input can cause the model to drastically change the weight between the experts, and the above term reduces to (cid:107) 14 · ( f 1 ( x ) − f 2 ( x )) · ( s 1 − s 2 ) (cid:107) .
Hence MoE models can suffer from large Lipschitz constants if the two experts are very different for points on the boundary, i.e., if | f 1 ( x ) − f 2 ( x ) | (cid:29) 0 .
Alternately if 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 and changes to routing do not result in much degradation of the ﬁnal accuracy.
3.2 Data routing
In this section we will study the effect of data routing on the Lipschitz constant of the individual experts – the ﬁrst term in RHS of equation 7 (lemma 1).
In particular we will try to address how large can max
i ∈
[ E ] { L f i } be in comparison to the Lipschitz constant of a dense model trained on the same data.
Towards this end, we will theoretically analyze the robustness of MoE models with a pre-determined ﬁxed routing and linear experts, in the light of the structure of the data.
3.2.1 Setup
We consider the setting of ﬁxed routing, where we assume the routing of individual examples to experts is pre-determined and ﬁxed.
Moreover, we 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), we would like to ﬁnd a linear model w ∗ ∈ R D that minimizes
[Equation]
The least-squares solution [12] for this problem that minimizes the MSE loss is w ∗
[Equation] [Equation]
Here † denotes the pseudo inverse.
On the other hand, for mixture of experts, let the dataset be split into E subsets S 1 , · · · , S E , with X i , y i denoting the data in set S i .
Now let the data from the set S i be routed to the expert i .
Below we write the objective for each expert:
[Equation]
Similarly, let w ∗ i =
( X (cid:62) 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 ﬁrst present two examples to better understand how routing affects Lipschitz constants.
In particular we consider a simple setting with E = 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]
The second equality follows from the folowing steps - 1)
X T X
= X T 1 X 1 + X T 2 X 2 , 2) ( X
[Equation]
X
[Equation]
+
X
[Equation]
2 X
[Equation]
) † =
[Equation] [Equation]
X
[Equation]
X
[Equation]
) † , and 3) ( X T 1 X 1 ) † X T 2 = 0 , where steps 2 and 3 follow from the orthogonality of X 1 and X 2 (see [15]).
Hence, experts have smaller Lipschitz constant when data 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 single dense 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 orthogonal subspaces in R D .
One way to construct them is by ﬁrst taking the singular vectors of X (cid:62) X and assigning them to the set
i with the largest projection.
This guarantees that U
i is orthogonal to U j with 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:
[Equation]
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 (cid:107)
[Equation] [Equation]
= j .
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 }
i ∈
[ E ] 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 Lipschitz constants 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 .
Assuming a 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 smaller by a factor of 1 √ E compared to dense models.
This happens when the data lies in E orthogonal subspaces and the data from each subspace is routed to the same expert.
This shows that MoE models can have signiﬁcantly smaller Lipschitz constant than their dense counterparts, while having the same computation cost.
For MoEs in practice, data routed to different experts does display some clustering of 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 between the Lipschitz constant of the dense and the MoE models.
This reduces the gap in the Lipschitz constant 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 D, we give a complementary result relating the Lipschitz constants of the experts to that of the dense model when both (cid:15) 1 and (cid:15) 2 can possibly be large.
Connections to Bubeck and Sellke [3].
Bubeck and Sellke [3] 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 lower bound on the Lipschitz constant of MoEs is smaller by a factor of (cid:113) 1 E .
Our result shows that this lower bound is in fact achievable, and hence tight, when the data lies in E orthogonal subspaces and data from each subspace is routed to the same expert.
[Equation]
Experiments 4.1 Setup
We compare the robustness of Vision Transformer (ViT) [7] and Vision MoE (V-MoE) [21] models against adversarial attacks.
In particular, we use the ViT-B/32 and V-MoE-B/32 models through all the experiments.
These models have the same backbone architecture but the latter replaces one in every two feedforward layers with a sparse Mixture of Experts, selecting K = 2 out of
E = 32 feedforward experts that are applied on each token 1 .
One could argue that the router in the
V-MoE model introduces an overhead that should be accounted for.
Thus, we have trained a bigger version of the dense ViT-B/32 (which we coin ViT-B++/32) that reaches roughly the same predictive performance as the V-MoE model, but has higher cost.
The cost of evaluating an image on the dense
ViT models is 8.9 and 17.9 GFLOPs and runtime cost, respectively; and 12.4 GFLOPs on the V-MoE model.
Appendix B contains additional experimental details.
We pre-train our models on the private dataset JFT-300M [23] for 7 epochs (517 859 steps with a batch size 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 the sigmoid cross-entropy loss.
V-MoE also adds auxiliary losses to encourage a balanced load for all experts; 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 linearly decayed to a ﬁnal value of 10 − 5 .
Weight decay of 0.1 was used on all parameters.
This is the same pre-training protocol as that used in [7, 21].
After pre-training, the models are ﬁne-tuned on ImageNet [6], at a resolution of 384 × 384 pixels and the 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 , reached after a linear warm-up of 500 steps, and followed with cosine decay to a ﬁnal value of 10 − 5 .
The norm of the ﬂattened vector of gradients is clipped to a maximum value of 10.
Since ImageNet images have a single label, we minimize the softmax cross-entropy loss during ﬁne-tuning.
The same ﬁne-tuning protocol was adopted in [7, 21].
We evaluate the adversarial robustness of both the pre-trained and ﬁne-tuned models, by means of PGD adversarial 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/ achieves 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 a function of the (cid:96) ∞ constraint.
Despite the fact that the V-MoE model contains a router that makes discrete choices among the experts conditioned on the input, which could potentially lead to a severe weakness against the adversarial attacker, we can observe that it follows the same trend as the base dense ViT model.
It is able to preserve a lower error over a wide range of (cid:96) ∞ values.
A larger version 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 adversarial training.
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 lines the 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 shows better robustness for a wide range of (cid:96) ∞ values.
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
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 experts have very different outputs, the Lipschitz constant of the MoE model could be much higher than that of a similar dense model.
If the index of the selected experts changes signiﬁcantly, but the model still shows a reasonably high accuracy under adversarial attacks (compared to a dense model), this would suggest that the outputs of the two selected sets of experts do not differ much.
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 by the model, we compute the intersection-over-union (IoU) of the selected set of experts before and after the adversarial attack.
We average the IoU across all processed tokens and deﬁne the rate of routing 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 adversarial training.
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]
ﬁgure 1 shows the V-MoE model keeping its advantage over the dense
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 two corresponding experts have a similar output, hence preventing the full V-MoE model from being less 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 decision boundary, 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.
In particular, in the implementation of V-MoEs, if the load of the experts is highly unbalanced, the experts receiving signiﬁcantly more tokens than the average could ignore all tokens that exceed the expert’s capacity, potentially leading to a signiﬁcantly worse performance.
The question is whether an adversarial attacker can exploit this property 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 adversarial attack, together with the corresponding cross-entropy loss.
We use the same weight for the auxiliary losses as the one used to train the models.
The ﬁgure shows that attacking the auxiliary loss does not offer any signiﬁcant advantage for the attacker.
The results are analogous in terms of the classiﬁcation error on ImageNet (not shown here in interest of space)
[Figure]
4.5 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 [3] is agnostic to the computational cost of the function, could we make a model more robust by increasing its model size without increasing the total cost?
Here we measure how much increasing the total number of experts in a V-MoE model helps against adversarial attacks.
Notice that increasing the total number of experts E does not make the V-MoE model more expensive.
Figure 3 (center and right) shows the false discovery rate on JFT-300M and the classiﬁcation error on ImageNet for an increasing 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 is most likely due to noise in the training and ﬁne-tuning process.
On the other hand, when the V-MoEs are ﬁne-tuned on ImageNet, all models with more than two experts achieve roughly the same accuracy under adversarial attacks.
This shows that, although the results presented in section 3 showing better robustness of MoEs require some assumptions, the conclusions hold to some extent in real scenarios.
Increasing the number of parameters by growing the number of experts is an effective way of improving the model’s robustness.
4.6 Robustness against AutoPGD attacks
We conducted additional experiments using a more sophisticated adversarial attack, AutoPGD [5], which selects the step size to use in each update of the attack.
We also increased the number of steps performed in the attack, using τ = 100 steps.
Figure 4a shows the results achieved by two dense (ViT-B/32 and ViT-B++/32) models and a sparse model (V-MoE-B/32).
Although AutoPGD is slightly more effective as an adversarial attack against all methods (for example, with (cid:96) ∞
= 10 − 3 the false discovery rate on JFT-300M of ViT-B/32 is 0.744 using PGD and 0.757 using AutoPGD), the trend of all models is identical to that represented in ﬁgure 1.
[Figure] [Figure]
Layers 1, 5, 7, 11)
[Figure] [Figure]
Figure 4b shows the rate of routing changes in the different MoE layers in the V-MoE-B/32 model.
This ﬁgure is analogous to ﬁgure 2 which shows the results for the standard PGD attack.
Compared to it, AutoPGD is able to change a higher fraction of the selected experts.
For example, the maximum fraction of changes in Layer 1 using PGD was near 0.2, while AutoPGD increases it up to around 0.5.
AutoPGD offers the same conclusion as the PGD attacks: models using sparse MoE layers offer better robustness against adversarial attacks (per GFLOP) than dense models, despite the fact that the router itself can be quite sensible to these attacks.
Conclusion
In this work we analyzed the adversarial robustness of MoEs showing their advantage over dense models, with more experts leading to better robustness, both theoretically and empirically.
We showed how 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 of the features in some layers of the model (see Figure 7 in Riquelme et al.
[21]), it is currently not explicitly encouraged during training.
Hence developing smarter routing strategies that take data geometry into account can be an interesting direction of future work.
Currently our analysis is limited to linear models, extending this to general models and deriving the dependency of optimal routing on them is another promising research direction.
We have also shown that, for inputs that weigh two experts similarly, if the two expert values are very different, then the MoEs can suffer from higher Lipschitz constant.
However, for models trained in practice, we saw their predictions to be relatively stable, despite signiﬁcant changes in choice of experts, highlighting potential redundancy of learned experts.
However too much redundancy, with all experts learning similar functions is a waste of capacity and can affect model performance.
Hence it is an interesting research problem to balance robustness and accuracy of MoEs by controlling the redundancy of experts.
Acknowledgments and Disclosure of Funding
All authors would like to thank Sven Gowal for sharing their AutoPGD attack codebase.
We also thank Neil Houlsby, Basil Mustafa, and André Susano Pinto for providing insightful feedback while working on this project.
Checklist
For all authors... (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s contributions 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] 2.
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. 3.
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 across different settings.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes] The hardware speciﬁcations and the total compute used in the experiments are detailed in appendix B. 4.
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 proprietary.
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 Computer Vision 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] 5.
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]
A Proofs
Proof of lemma 1.
We have by chain rule ∇ x f ( x ) =
k (cid:88) i
= p
i ( x ) ∇ f i ( x ) + k (cid:88) i = ∇ p
i ( x ) f i ( x ) .
Hence, (cid:107)∇ x f ( x ) (cid:107) ≤ k (cid:88) i = p
i ( x ) L f i + (cid:107) k (cid:88) i = ∇ p
i ( x ) f i ( x ) (cid:107) , where
[Equation]
Let ¯ s ( x )
=
(cid:80) j p j ( x ) s j .
Substituting this in the above equation gives us
[Equation]
x
[Equation]
≤
[Equation]
x )
[Equation]
The last inequality follows from Holder’s inequality and (cid:80) ki =1 p i ( x ) = 1 .
Proof of theorem 1.
Since U i are orthogonal to each other, we can decompose (cid:107) w ∗ (cid:107) as follows.
[Equation]
From deﬁnition 1 we have, for any unit norm vector z ,
[Equation]
(cid:12)(cid:12)
The last step follows from triangle inequality.
This implies
[Equation]
Hence, by setting z = X (cid:62) y / (cid:107) X
(cid:62) y (cid:107) , we see that
[Equation]
− (cid:15) 1 (cid:107) X (cid:62) y
(cid:107) .
Recall that X (cid:62) y =
[Equation]
=1 X
[Equation]
j y j .
Then using triangle inequality
[Equation]
and the above inequality, we get,
[Equation]
The last step follows from deﬁnition 2.
Thus,
[Equation]
Substituting this in equation 10 gives us the result.
(cid:107) w ∗ (cid:107) 2 ≥ (cid:88) i ( (cid:98)(cid:107) w ∗
i (cid:107) − (cid:15) 1 (cid:107) X (cid:62) y (cid:107) − (cid:15) 2 (cid:88) j (cid:54) =
i (cid:107) X (cid:62) j y j (cid:107)(cid:99) + ) 2 .
B Experiments details B.1 Architectures
Table 1 contains the details of the architectures used in section 4.
All V-MoE models have replaced one in every two feedforward layers in the original ViT-B/32 architecture, with a MoE of feedforward layers, and select only K = 2 experts out of the E available experts.
All models were pre-trained on
JFT-300M, at a resolution of 224 × 224 pixels, for a total number of 7 epochs, using the same batch size and optimizer settings as used by Dosovitskiy et al.
[7] and Riquelme et al.
[21] (see details section 4).
We used the B/32 variants of ViT and V-MoE since these are the “base” conﬁgurations suggested in the respective papers.
In order to match the quality of a dense ViT and a sparse V-MoE model, we also trained a bigger version of ViT-B/32, which we refer to as ViT-B++/32.
The values for the number of layers, number of attention heads, embedding dimension and the hidden dimension of the FFN, are simply an interpolation between the ViT-B/32 and the ViT-L/32 corresponding values.
We implemented all models, training and evaluation code using JAX 2 and FLAX 3 .
Although we cannot release the models used in the experiments, since they are pre-trained on proprietary data, the code used in all of them is available at http://github.com/google-research/vmoe .
B.2 Compute resources
We use TPUv3 to train and evaluate our models.
In particular, we used 32 TPUv3 cores for pretraining and ﬁne-tuning the models.
In most of the evaluation experiments against adversarial attacks we used 8 TPUv3 cores, since the compute needed to perform the evaluation against adversarial attack is much lower.
The total training cost (in terms of total training runtime and FLOPs) of ViT-B/32 and
V-MoE-B/32 can be found in [21].
ViT-B/32 is pre-trained for a total of 27.6 TPUv3-core-days, performing 56.1 ExaFLOPs.
V-MoEB/32 (with 32 experts) is trained for a total of 54.9 TPUv3-core-days and performs a total of 76.
ExaFLOPs.
Finally, the larger dense ViT-B++/32 is trained for a total of 78.3 TPUv3-core-days and performs 113.4 ExaFLOPs.
The GFLOPs are computed automatically by JAX.
For instance, to compute the GFLOPs used for each individual image during the evaluation of the models, the code snippet in listing 1 is used.
[Table] [Table]
D A complementary result to theorem
We now present our complementary result.
In essence, in a regression setting, our result tries to capture how different the sub-problems tackled by the experts should be in order to get improved robustness compared with a dense approach.
In what follows, we use the shortcuts Z = X (cid:62) X ∈ R D × D and Z i =
X (cid:62)
i X i ∈
R D × D for i ∈
[ E ] .
Also, we recall that each X i is in R N i × D with (cid:80) Ei =1 N i = N .
Theorem 2.
Let w ∗ be the minimizer of equation 8 and { w ∗ i }
i ∈
[ E ] be the minimizers of equation 9 .
For all i ∈ [ E ] , let us assume that there exist β i ∈ R D and r i ∈
R N i such that we can write y i
=
X i β i
+ r i where the vectors r i ’s are residual terms that can for instance account for non-linear effects in X i and/or some stochastic noise.
Further assume the { Z i } i ∈
[ E ] to be invertible and deﬁne η i =
Z − 1 i X (cid:62) i r i and η = Z − 1 X (cid:62) r with r = [ r 1 , . . .
, r E ] ∈ R N .
Let us denote the (normalized) difference between any two vectors ( β i , β j ) as ∆ ij ∈ R D , that is for any i (cid:54) =
j, ∆ ij =
Z − 1 Z j ( β i − β j ) , and consider the cumulative difference with respect to the vector β i as for any i ∈
[ E ] , ∆ i = (cid:88) j (cid:54) =
i ∆ ij ∈ R D .
Consider γ ∈ (0 , 1] .
If the structures of the underlying sub-problems on { ( X i , y i ) } i ∈ [ E ] differ sufﬁciently from each other in the following sense for all
i ∈
[ E ] , (cid:107) ∆ i (cid:107) ≥
[Equation]
√
[Equation]
, then it holds that
[Equation]
In other words, the experts have smaller Lipschitz constants than the dense model by a factor γ .
As a corollary, if we can control the residual terms r i ’s such that
[Equation]
then the same conclusion is implied by the simpler sufﬁcient condition
[Equation]
Proof of theorem 2.
As a ﬁrst comment, since we assume that all the Z i ’s are invertible, then so is
Z = (cid:80)
Ei =1 Z i .
We readily have for i ∈
[ E ]
[Equation]
Moreover, for any i ∈ [ E ] , we have [Equation]
=   [Equation]
 
[Equation]
=
[Equation]
  [Equation] j [Equation]
 
[Equation] j [Equation] = β
[Equation]
−
[Equation]
i + η = w ∗ i − ∆ i + η − η i . To summarize, for any i ∈ [ E ] , it therefore holds
[Equation]
Let us consider γ ∈ (0 , 1] .
We can observe that
[Equation]
We develop each term individually:
[Equation]
with
[Equation]
and using w ∗
=
β i − ∆ i + η = β
i + η i − ∆ [Equation] + η − η i , we get
[Equation]
Gathering (A) + (B) + (C) , we obtain the expression
[Equation]
We want to understand the conditions on ∆ i to guarantee that
[Equation]
To ﬁnd a sufﬁcient condition, we lower bound the quadratic form φ (∆ i ) as follows
For any ∆ i ∈ R D , φ (∆ i ) ≥ γ (cid:107) ∆ i (cid:107) 2 −
[Equation]
γ
[Equation]
∆ i
[Equation]
β
[Equation]
η (cid:107) + γ (cid:107) β i + η (cid:107) 2 − (cid:107) β i + η i (cid:107)
[Equation]
root
[Equation] root [Equation]
) with root [Equation] = [Equation]
and root −
[Equation]
The lower bound is a standard polynomial of degree 2 and it is non-negative beyond its positive root
[Equation]
We apply the exact same rationale for all i ∈
[ E ] and take the minimum of the conditions to make them hold simultaneously.
Since γ ∈ (0 , 1] and 1 / √ γ ≥ 1 , if we further assume that the residual terms are sufﬁciently small
• (cid:107) η (cid:107) ≤ min i ∈
[ E ] (cid:107) β
i (cid:107) • for all i ∈
[ E ] , [Equation]
then we get the simpler sufﬁcient condition after using the triangle inequality
[Equation]
E Auxiliary losses for MoEs
In this section we present the Auxiliary losses used for training MoEs and attack objectives used for our experiments in section 4.
Auxiliary losses are used to ensure the data is routed in a balanced way across different experts.
In particular we use the losses from Riquelme et al.
[21] which are deﬁned below.
Recall that the router weights for a given token x and a given expert j , with routing parameters { s j } Ej =1 , were deﬁned in equation 6 as:
[Equation]
In practice, we use a noisy version of this router:
[Equation]
with (cid:15) j ∼ N ( µ = 0 , σ = 1 E ) .
The quantity (cid:104) s j , x (cid:105) is the routing logit for a given token and expert j .
Importance Loss.
The importance of expert j for a batch of tokens X is simply deﬁned as the sum of the router weights assigned to that expert over the tokens in the batch:
[Equation]
where X =
[ x 1 , . . .
, x n ] (cid:62) and p j ( x i ) is the routing weight assigned to the j -th expert for the i -th token.
We use the squared coefﬁcient of variation of the importance distribution over experts, Imp ( X ) :
= { Imp j ( X ) }
Ej =1 :
[Equation]
Load Loss.
For a given token, the load loss is based on the probability that the noisy logit of expert j is among the top- k logits when a new Gaussian sample is drawn.
Let’s deﬁne the k -th maximum noisy logit for a given token x as:
[Equation]
Then, the probability that the noisy logits of expert j are above the threshold if the noise is resampled, is given by the expression:
[Equation]
Analogous to the importance, the load of expert j is deﬁned as the sum over the tokens in a given batch:
[Equation]
And the load loss corresponds to the squared coefﬁcient of variation of the load distribution, with
Load ( X ) :
= {
[Equation]
[Equation]
Final Loss.
The ﬁnal loss used for training the models includes the classiﬁcation loss (i.e. crossentropy) and both auxiliary losses with weights 0.005:
[Equation]
F Experiment Metrics
Precision at 1 and False discovery rate In multi-label datasets (such as JFT), the traditional accuracy metric used for single-label classiﬁcation is ill-deﬁned.
Usually, Precision at k , Recall at k , and other more general metrics are used in this setting.
In particular, we use the Precision at 1 .
Let Y =
[ y 1 , . . .
, y N ] (cid:62) , y i ∈ R C be the logits output by the model for a set of N examples, with true labels ˆ Y =
[ˆ y 1 , . . .
, ˆ y N ] (cid:62) , ˆ y i ∈ { 0 , 1 } C .
Precision at 1 is deﬁned as:
[Equation]
where arg max c y i,c is essentially the index of the class with the highest logit, for a given example i , and δ is the Kronecker delta function.
That is, the non-zero elements of the sum are those examples for which the highest logit corresponds to one of the true labels of such example.
The false discovery rate is deﬁned as 1 − P @1 .
Observe than when the examples have a single label, precision at 1 is equivalent to the accuracy, and the false discovery rate is equivalent to the error rate.
Rate of routing changes Given two sets A and B , the intersection over union (IoU) is deﬁned as:
[Equation]
If the two sets are equal, IoU is 1.
If the two sets do not share any element, IoU is 0 .
We use this metric in section 4.3 to compare how much the set of selected experts change when adversarial attacks are peformed on V-MoE models.
Thus, the rate of routing changes is deﬁned as 1 − IoU.