MomentumSMoE: Integrating Momentum into Sparse Mixture of Experts
Authors: Rachel Teo
Tan Minh Nguyen
Venue: NeurIPS 2024 poster (2024)
Keywords: Sparse Mixture of Experts
optimization
gradient descent
momentum
adam
Primary area: deep_learning_architectures
TL;DR: We first establish a connection between the dynamics of the expert representations in SMoEs and gradient descent on a multi-objective optimization problem and then integrate momentum into SMoE to develop a new family of SMoEs, named MomentumSMoE.

Abstract: Sparse Mixture of Experts (SMoE) has become the key to unlocking unparalleled scalability in deep learning. SMoE has the potential to exponentially increase in parameter count while maintaining the efficiency of the model by only activating a small subset of these parameters for a given sample. However, it has been observed that SMoE suffers from unstable training and has difficulty adapting to new distributions, leading to the model's lack of robustness to data contamination. To overcome these limitations, we first establish a connection between the dynamics of the expert representations in SMoEs and gradient descent on a multi-objective optimization problem. Leveraging our framework, we then integrate momentum into SMoE and propose a new family of SMoEs, named MomentumSMoE. We theoretically prove and numerically validate that MomentumSMoE is more stable and robust than SMoE. In particular, we verify the advantages of MomentumSMoE over SMoE on a variety of practical tasks including ImageNet-1K object recognition and WikiText-103 language modeling. We demonstrate the applicability of MomentumSMoE to many types of SMoE models, including those in the Sparse MoE model for vision (V-MoE) and the Generalist Language Model (GLaM). We also show that other advanced momentum-based optimization methods, such as Adam, can be easily incorporated into the MomentumSMoE framework for designing new SMoE models with even better performance, almost negligible additional computation cost, and simple implementations.

MomentumSMoE: Integrating Momentum into Sparse Mixture of Experts
Rachel S.Y. Teo Department of Mathematics National University Singapore rachel.. tsy@u. mus edu
Tan M. Nguyen Department of Mathematics
National University of Singapore
tanmn@nus. Gannn@nus.edu.sge
Abstract
Sparse Mixture of Experts (SMoE) has become the key to unlocking unparalleled scalability in deep learning. SMoE has the potential to exponentially increase in parameter count while maintaining the efficiency of the model by only activating small subset of these parameters for given sample. However, it has been ob- served that SMoE suffers from unstable training and has difficulty adapting to new distributions, leading to the model's s lack of robustness to data contamination. To overcome these limitations, we first establish connection between the dynamics of the expert representations in SMoEs and gradient descent multi-objective optimization problem. Leveraging our framework, we then integrate momentum into SMoE and propose new family of SMoEs, named MomentumSMoE. We the- oretically prove and numerically demonstrate that MomentumšMoF is more stable and robust than SMoE.. In particular, verify the advantages of MomentumSMoF over SMoE on a variety of practical tasks including ImageNet- K object recogni- tion and WikiText-100 103 language modeling. We demonstrate the applicability of MoMe to many types of SMoE models, including those in the Sparse MoE model for vision (V-MoE) and the Generalist Language Model (GLaM), We also show that other advanced momentum-based optimization methods, such as Adam, can be easily incorporated into the MomentumSMoE framework for design- ing new SMoE models with even better performance, almost negligible additional computation cost, and simple implementations. The code publicly available https:/ / gi thub. cow/rachtsy http///omeettttteeeoe
1 Introduction
Scaling up deep models has demonstrated significant potential for enhancing the model's performance on wide range of cognitive and machine learning tasks, ranging from large language model pre- training [13. 58, 59, 30, 5, 51, 70] and vision understanding [15, 3, 39, 1, 40] to reinforcement learning [6, 28] and scientific applications [66, 74J. However, increasing the model's size requires higher computational budget, which can be often challenging to meet. As a result, Sparse Mixture of Experts (SMoF) has been recently studied as an efficient approach to effectively scale up deep models. By modularizing the network and activating only subsets of experts for each input, SMoE maintains constant computational costs while increasing model complexity. This approach enables the development billion parlmee-paramtter models and achieves significant success various applications, including machine translation [35], image classification [61], and speech recognition [34].
1.1 Sparse Mixture of Experts
MoE replaces component in the layer of the model, for example, feed-forward convolutional layer, by a set of networks termed experts. This approach largely scales the model but increases
Please correspond to: rachel. taayu.aauuaaam nus edu
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

the computational cost. A SMoE inherits the extended model capacity from MoE but preserves the computational overhead by taking advantage conditional computation. In particular, a SMoE consists of router and E expert networks, Vi i = 1, 1.2.... ......... F.. For each input token Øt R layer i, the SMoE's router computes the affinity scores between ay and each expert as gi(a;), i 1, 2, ..... E. practice, we often choose the router g(x;) [g] (æ+), (z.( g (x+)] Wx + b, where W IR ExD and b E IR.. Then, a sparse gating function TopK is applied to select only experts with the greatest affinity scores. Here, we define the TopK function as:
gi: if go is in the K largest elements of (gp)(g:) = {{
1-00, otherwise.
The outputs from K expert networks chosen by the router are then linearly combined as
x++1 32 +
u(xx),
where softmax(g:)) == exp(g:)/ == exp(g)). We often set K = 2, i.e, top-2 routing, as this configuration has been shown to provide best trade-off between training efficiency and testing performance [35, 16, 76].
Limitations of SMoE. Despite their remarkable success, SMoE suffers from unstable training [11, 78] and difficulty in adapting to new distributions, leading to the model's lack of robustness to data contamination [55, 75]. These limitations impede the application of SMoE to many important large-scale tasks.
1.2 Contribution
In this paper, we explore the role the residual connection in SMoE and show that simple modifica- tions of this residual connection can help enhance the stability and robustness of SMoE. In particular, we develop gradient descent (GD) analogy of the SMoE, showing that the dynamics the expert representations in SMoF is associated with gradient descent step toward the optimal solution of multi-objective optimization problem. We then propose integrate heavy-ball momentum into the dynamics of SMoE, which results in the Momentum Sparse Mixture-of-E-ee -ff-Ets (MomentumSMoFF) At the core of MomentumSMou architecture of MoomentumSMoF is depicted in Fig. 1. MomentumSMoFF can be extended beyond heavy-ball momentum to integrate well with other advanced momentum-acceleratee methods such as A [33, 42] and Robust Momentum [10]. Our contribution is three-fold:
the use of momentum to stabilize and robustify the model. The
1. We incorporate heavy-ball momentum in SMoE to improve the model's stability and robustness.
We theoretically prove that the spectrum MomentumSMoF better-structured than SMoE, leading to the model's stability enhancement.
We show that the design principle of MomentumSMoE can generalized to other advanced momentum-based optimization methods, proposing AdamSMoE and Robust MomE
Our experimental results validate that our momentum--asee SMoEs improve over the baseline SMoE terms of accuracy and robustness on a variety practical benchmarks, including WikiText- 103 language modeling and ImageNet-1K K object recognition. We also empirically demonstrate that our momentum-bain based design framework is universally applicable to many existing SMoE models, including the Sparse MoE model for vision (V-MoE) [61] and the Generalist Language Model (GLaM) [16]. by changing a few lines of the baseline SMoE code.
Organization. We structure this paper as follows: In Section 2, we establish the connection between SMoE and gradient descent and derive our MomentumSMoE. In Section 3, we theoretically prove the stability advantage of MomentumSMoE over SMoE. In Section 4, we introduce AdamSMoE and Robust MomentumSMoo.. In Section we present our experimental results to justify the advantages of our momentum- based SMoE models over the traditional SMoE and other SMoE baselines. In Section 6, we empirically analyze our MomentumSMMF. We discuss related works in Section 7. The paper ends with concluding remarks. More experimental details are provided in the Appendix.

---

SMoE layer
MomentumSMoE layer
1++1 = It - ((())
Pt = - (x+) + PPt-
Dt+1 = Dt + PPt
31
1t++1
Tt
21+1
Pr
Pt
Zt
Zt
Figure 1: Illustration of SMoE (Left) and MomentumSMof layer (Right). We establish a connection between Multiple-Gradienn Descent and SMoE to introduce momentum into the model, leading to better accuracy, enhanced robustness, and faster convergence.
2
Momentum Sparse Mixture of Experts
2.1 Background: Multiple-Gradient Descent Algorithm for Multi-objective Optimization
A multi-objective optimization problem comprises of the concurrent optimization of E objective functions, (a)) = 1,,2.... E, which might be formulated as the following minimization problem
min F(x) = Σ (.))
(3)
where the feasible region and e; E R are weights representing the importance of each objective function. The optimal solution to the multi-objective optimization problem above is a Pareto-optimal point such that there no other solution that can decrease at least one of the objective functions without increasing any other objective functions. [12] shows that a necessary condition for solution to be Pareto-optimal for it to Pareto-stationary, which is defined as:
Definition (Pareto-stationary) Let x be in the interior of the feasible region, D, in which the E objective functions, F,, are smooth, and fi(x) (zF((x) be the local gradients for i 1..... E. is said to be Pareto-stationary if there exists vector a [ .rr] R such that ag 2 0, Ex' O4; = 1 and Ee= fz = 0. That is, there exists a convex combination of the gradient-vectors fi(z) that is equal to O.
Therefore, it would be intuitive extend the steepest descent algorithm a multi-objective setting by finding a descent direction, that is common to all objectives, in the convex hull of the normalized local gradients fz)
f((=)/llffxx)) We denote such set as Ü {u RM Σ a.f( a, 2 O, Vi; E a2 = 1}. Indeed, [12] developed the Gurdipe Descent Algorithm (MGDA) from such an understanding, proving that there does exist such descent direction in U, which the direction with the smallest norm in the set. Then, the update rule of MGDA is
F++1 = 24 - Σ (x))
1=1
where a= = [a]....... ....)) minimizes {|[i|||e U}.
2.2 Background: Momentum Acceleration for Gradient-Based Optimization
Among the simplest learning algorithms is gradient descent, also termed the steepest descent method. It typically starts with objective function F(x) whose minima we aim to find by modifying our iterate F4 at each time step through its gradient f(x)) = (xx,, scaled by step size 7 > 0:
I++1 = Ft F(x)) = 2, (xt))
(5)
However, following this update rule might result in slow convergence. A classical acceleration method to speed up the steepest descent, known as the heavy-ball method [53, 67], includes a momentum term in the algorithm. This takes the form of
Pt = - f(x) + PP:-13 Fe41 = + YP:
3

---

Training
Validation
MoE
SMoE
Layer
Layer
Layer
Layer
0.95
0.98
[=)()))]
0.50
0.80
0.50
0.55
0.80
0.40
0.75
20
40
60
80
0
20
40
60
80
20
40
60
20
40
60
Epochs
Figure 2: Average output norms at layers I and 6 of the MoF/SMoF during 80 training epochs on WikiText-103.
Training
Validation
MoE
SMoE
100
80
80
(x)
70
60
50-
40
Layer
Layer
Figure 3: Average output norm at each layer across IK train/validation samples of the (S)MoE trained on WikiText-1003
where y > 0 is the step size, and 2 0 is the momentum constant. Eqn. 6 can then be rewritten as:
T++1 + ((=)(())) μPt-1]] It - + +((z x1-1).
By incorporating the past gradients in each update, the descent path become smoother with fewer oscillations, resulting in a faster convergence [53, 21, 71].
2.3 (S)MoE as (Stochastic) Gradient Deseent on Multi-ohjective Optimization and
MomentumSMoF
We will now consider the SMoE model from the multi-objective optimization perspective.
MoE as Gradient Descent, In viewing each expert in an SMoE as specializing in optimizing an objective function, are going to establish connection between MoE and GD, and further leverage momentum to enhance MoE and SMoE. We rewrite Eqn. 2 of MoE as follows:
L++1 = Dt
(8)
If we regard (2;) as the local "gradient" V Fx(;)) at the f-th iteration and softmax (g: (xx)) ) as to learned to approximate oz, then we can consider the MoE Eqn. system which updates æ; using the MGDA to minimize the multi-objective optimization in Eqn. 3.
as the dynamical
SMoE Stochastic Gradient Descent. Given the analogy between MoE and GD, SMoE be interpreted as stochastic version of MoE, which corresponds to an SGD algorithm applied to the multi-objective optimization problem Eqn. SMoE is then reformulated as:
E1+1 Tt.
-m) æz - 7f(e))
(9)
Here, -x(x)) 7 K softmax (x: (x) is the SMoE output.
Empirical Evidences for the Gradient Descent Analogy of (S)MoF. We provide empirical justifi- cation for the connection between (S) and (S)GD in Fig. 2 and 3.

---

Gradient norm (x:)l decreases when t increases: As shown in Eqn. 8 and 9 above, the norm of the MoE and SMoE output corresponds to the gradient norm fx=) V vz(((z), respectively. It expected that this gradient norm decreases when t increases or equivalently when the number of layers in an (S)MoE model increases. Fig. 3 confirms this expectation by showing that the norm of the (S)MoE output decreases over layers in a 6-layer (S)MoE model trained on the WikiText- language modeling task. At the last layer, the norm increases might be due overshooting, common phenomenon that can occur when using gradient descent.
Gradient norm [ each layer decreases during training: According to the update rule of
MGDA in Eqn. 7, the coefficient of minimizes the norm Ei=1 0;.fell. In Eqn, 8 and 9, as discussed above,
(g: and coftmax(Topp((/g) (x:))) try learn a;, respectively. Thus,
is expected that these two terms learn to reduce the corresponding || E>= asfflll in Eqn. 8 and 9, which is the norm of the SMoE output at layer L. Fig. verifies this expectation by showing that each MoE and SMoE layer learns to reduce its output norm during training, suggesting that the routers softmax [g. (zt)) and softmax (opK(gg (g: (xr))) learn to approximate aj. We provide the full plots for all layers in Appendix C, Fig. 6 and
2.4 MomentumSMoE
We propose the new MomentumSMof layer, depicted in Fig. 1, accelerate the dynamics of 8, which is principled by the accelerated gradient descent theory (see Section 2.2):
(ax) + 11p-1; 1t++1
2t + TP:
(10)
where μ 2 0 and > 0 are hyperparameters corresponding to the momentum coefficient and step GD, respectively. The formulation of MomentumSMoE can be
size the comeenuum-aceeleree applied to MoE to derive the MomentumMoE.
3 Stability Analysis: MomentumSMoE vs. SMoE
In this section, demonstrate the theoretical advantages of MomentumSMoE over SMoE. In particular, we show that the spectrum of MomentumSMoF better-structured than that of SMoE, thus MomentumSMoE more stable than SMoE. We rewrite MomentumSMoE using the equivalent form of momentum acceleration given in Eqn. as follows:
E+++
(xx)) + p(æ; 21-1)
(11)
Taking inspiration from [57]. we then expand (æ;) around the Pareto sureto-stationayy solution xx at which f(x*) = (see Definition 1) using Taylor expansion to obtain an approximation of f(x:):
fxx)) = f(x") + f=f(xx (x) - x²) = V fx=" (== (xx - x=))
(12)
Substituting the Taylor expansion of f fxx) in Eqn. 12 into Eqn. 11, we attain
x1+1 3; v (x)) x²) + p(æ, xz-1))
Without loss of generality, we further let x = O as we can always replace x, with æ, + x=" The formula of MomentumSMoE can then be simplified as
Ø1+1 = 3; ? V~)(("") (a="((x);;; u(xx x.-1).
(13)
Suppose that V=f(=*) does not have any defective eigenvalues and hence is diagonalizable. Then, V=f(=*) QEQ for some invertible matrix Q and the diagonal matrix Σ with diagonal entries being the cigenvalues or (n), n 1, ........ N, of fx.f((")) We can then rewrite Eqn. 13 as
Et+1 = Cc Ex; + f( - x.-)))
(14)
Since we have decoupled the N features in Dt, we can consider each feature (n()), separately. Introducing a dummy equation Φ; 2t, we rewrite Eqn. 14 as follows:
(n) (
(- (1+ - -o(o))
( = A (
(15)
The convergence of (n) then depends on the eigenvalues X1 (A) and (2 of A. In particular, we require max ([> (A)l, (A(())] < 1. should be noted that omitting the momentum parameter, i.e., 0, recovers the standard, unaccelerated SMoE layer.
5

---

Lemma Given the matrix A (-p
(1 + μd) ra(n) and A1(A), X²(A) are eigenvalues of A, (Mx))))()))))A((()) < 1 if and only if μ (-1,1) and vo(n) E (0, 2p)..
Proposition 1 (Convergence of MomentumSMoE) The MomentumSMoE defined in Eqn. 10 con- verges if and only if E (-1,1) and (n) E (0, +
The proofs the results above are provided Appendix A. is worth noting that in both the MomentumSMoE and standard SMoE, the convergence of æ depends the eigenvalues of the Jacobian V V of the SMoE layer. Since the step size y > O, we require that Væf to be positive definite for its eigenvalues a (n) be positive. Furthermore, conditions of MomentumSMoFi is that E (-1,1), in practice, le is chosen to be positive.
though among the convergence
Proposition 1 implies that the spectrum of MomentumSMoFi is better-structured than that of SMoE. Thus, MomentumSMME is more stable than SMoE. We summarize this finding in Corollary I below.
Corollary (MomentunSMoF is more stable than SMoE) Without momentum, μ 0, the range of values that yo can take for the system to be stable is limited to 0 < yo(n) < 2. The addition of momentum expands this margin, almost doubling it, providing larger parameter range for the network to converge stably to a good output in the forward pass.
4 Beyond Heavy-ball Momentum: AdamSMoE and Robust MomentumSMoE
In addition to heavy-ball momentum, there are several advanced momentum-based algorithms in optimization that can be utilized for SMoE design. In this subsection, propose two additional variants of MomentumSMoE, AdamSMoE and Robust MomentumSMoE, which are derived from the AdamW [33, 42] and Robust Momentum [10], respectively.
4.1 Adam Sparse Mixture of Experts (AdamSMoE)
Adam [33] accelerates the gradient dynamics by utilizing the moving average of historical gradients and element- wise squared gradients. Adam with decoupled weight decay regularization (AdamW) more commonly used in practice thanks to its better generalization over Adam. We employ AdamW to derive the AdamSMoF as follows:
P: 1P::11 + (1 - [x==((x;));
m, 3mmm1 + (1 (xx, o f(x))
Ft+1 at Vm
where is small constant to prevent numerical instability, the weight decay parameter, 7 the step size, and and 3 are the decay parameters for the moment estimates.
4.2 Robust Momentum Sparse Mixture of Experts (Robust MomentumSMoE)
Deep learning models, including SMoE, are known to not be robust distribution shifts and data distortions [60, 19, 14]. Utilizing the connection between (S)MoE and (S)GD in Section 2.3, we develop the new Robust MomentumSSME from the Robust Momentum Method [10].
The Robust Momentum Method proposed by [10] has the following update rule
1/2 _4 + axx xt-1); Fe41 x. (1)) u(x, xz-1),,
(16)
where 7, y and parameterized by additional hyperparameter p follows:
kp k<1 (10-))))111 p) == 1 a
p²
(17)
(k: a= 1)(1 - p)
Here, L/m is condition ratio the objective function assuming that is m-strongly convex and L-smooth. Compared with the heavy-ball momentum in Eqn. 7, there is an additional variable y that can interpreted as feedback signal steer the J, toward a robust solution. The parameters 7, u, and designed such that the new system is robust.

---

Table Perplexity (PPL) of momentum-bant SMoE vs. SMoE baseline on clean/attacked WikiText-
103.
Clean WikiTexx-1033
Attacked WikiText-103
Model/Metric
Parameters
Valid PPL +
Test PPL
Valid PPL +
Test PPL 4
SMoE- (baseline)
216M
33.76
35.55
42.24
44.19
MomentumSMoo--medii
216M
32.29
33,46
40,94
42.33
AdemSMoE-mediuu
216M
31.59
33.25
39.27
41.11
SMoE SMoE-large (baseline)
388M
29.31
30.33
36.77
37.83
MomentunSSof large
388M
27.58
29.03
35.21
36.78
GLaM-medine (baseline)
220M
36,37
37.71
45.83
47.61
ComnentumGLaM-mellne AdamGLaM-mediun
220M
33.87
35.29
42.15
43.64
220M
32,99
34.32
41.09
42.81
We incorporate the Robust Momentum Method above in our SMoE optimization framework developed in Section 2.3 and formulate the novel Robust MomentumSMoE as follows:
y: It + a(xx xz-1);; I = I, - 734.)) + u(x; - xx-1).
(18)
where 7r μl and are defined in Eqn. 17, and - (V)) Σ softmax (TopK((TTo(o))) (g:)) (y:) is the SMoE output given the input y:- Equivalently, at each layer I, we update the input 2e and momentum vector pt as
Y = P
Pt --f(w) pP--1: L++1 xt YPo.
Remark 1 We provide an interpretation of robust momentum in Appendix B.
5
Experimental Results
In this section, we numerically justify the advantages of our momentum-based SMoE over the baseline SMoE on both WikiText-103 language modeling and ImageNet- object recognition tasks. We aim to show that: (i) MomentumSMof improves model performance across both language and vision tasks; (ii) AdamSMoE significantly outperforms the baseline and accelcrates convergence in language models, even surpassing MomentumSMoE; (iii) Robust MomentumSMME is highly effective at improving robustness of vision models to data corruption; (iv) MomentumSMoE is universal and can be easily integrated into many sate-oo-theeart SMoE and MoE models.
Throughout the experiments, compare momentum- based SMoE with the baseline SMoE of the same configuration, replacing SMoE layers with our momentum- based SMoE. For Adam-based SMoE models, we use AdamSMoE in the first layer of the model and MomentumSMoE for the subsequent layers. We provide an explanation for this implementation in Appendix D.I. We find that implementing AdamSMoE in the first layer is enough to significantly improve the model's 's overall performance and accelerate its convergence. Our results are averaged over 5 runs. Details on datasets, models, and training are provided in Appendix D, along with Table 4 summarizing the momentum methods implemented on SMoE/MoE models for different tasks in our experiments. More results can be found in Appendix E. All experiments are conducted on a server with A100 GPUs.
5.1 1ik Language Modeling
We use the Switch Transformer [18), referred to as SMoE in our tables and figures below, and GLaM [16] baselines. We consider configurations: medium layers) and large (12 layers). We report the perplexity (PPL) of MomentumsMME and AdamSMoE in comparison with the baseline SMoE on word-level WikiText-100 validation and test datasets for both model sizes in Table 1. We also include experiments on the medium-sizee GLaM. A lower PPL indicates a better performance of the model. To further demonstrate the robustness of our method, we test the models on word swap attacked WikiTexx- 103 data and present their results. Across all metrics, AdamSMoE and AdamGLaM outperform the baseline by a significant margin, verifying the strength of our method. Additionally, in Figure 4(Left), we provide the training and validation PPL. during the first 5 training epochs of MomentumSMoF. and AdamSMoR compared to the baseline SMoE to illustrate the accelerated convergence of our momentum-Be based models.
7

---

WikiText- 103
ImageNet- 1K
SMoE
MomentumSMcE
Soft MoE
Momentum MoE
AdamSMoE
Train PPL
Validation PPL
Training loss
Top-1 validation accuracy
438
600
422
100
332
305
20-
Epoch
Epoch
Epoch
Epoch
Figure 4: Left: WikiText- 103 train/validation perplexity (PPL) curves during the first training
epochs for MomentumSMoE, AdamSMoE, and SMoE. AdamSMoE has significantly faster conver- gence compared to SMoE. Right: Training loss/top-1 accuracy (%) of Momentum-SS MoE vs. Soft MoE baseline ImageNet-1K across 120 epochs of training. Momentum-See MoE has faster convergence and improved accuracy.
Table 2: Top-1 accuracy (%) and mean corruption error (mCE) of Momentum V-MoE and Robust Momentum V-MoE vs. the V-MoE baseline on ImageNet-1k and popular robustness benchmarks for image classification.
Train IN-1K
Valid IN-IK
IN-R Top-1 f
IN-A Top-1 f Top-1 f mCE $
IN-C
Model
Params
Top-1 f
Top-5
Top-1 f Top-5 f
V-MoE (baseline)
297M
76.49
92.27
73.16
90.42
36.10
5.25
46.98
67.14
MomentummVMoo
297M
76.92
92.19
73.26
90.30
37.45
6.48
48.11
65.77
Robust Momentum V-MoE
297M
76.66
92.27
73.20
90.36
37.57
6.37
48.82
64.92
An important advantage of Momentum.SMoBi is its simplicity, which allows easy implementation with negligible computational overhead. We provide comparison of the time/sample, memory, number parameters, and computation time between models in Table and 12 in Appendix E.5.
5.2 ImageNet-1K 1K Object Recognition Task
In this section, we investigate our momentum- based models on two popular vision models, Vision MoE (V-MoE) [61] and Soft MoE [56] on the ImagetNet-IK (IN-1K) object recognition task. We focus on i) improving the robustness of V-MoE using Robust MomentumSMMf (18) and ūi) demonstrating that our momentum method not limited to sparse models but can be generalized to MoE models such as Soft MoE. To benchmark robustness to data corruptions, use the standard datasets, ImageNet-R (IN-R) [24], ImageNet-A (IN-A) [26]. and ImageNet-C (IN-C) [25].
Vision Mixture of Experts (V-MoE). We use a V-MoE (small) model as the baseline. This V-MoE consists of 8 Vision Transformer (ViT) blocks [15] with every odd block's MLP being replaced by SMoE layer. In Table 2, we provide the top-1 accuracy (%) on the training and validation set of IN-1K, IN-R, IN-A, and IN-C, as well as the mean Corruption Error (mCE) for IN-C. While Momentum V-MoE and Robust Momentum V-MoE have marginal gains on clean IN-IK data, we see significant improvement on IN-R, IN-A, and IN-C with at least a 1% increase accuracy across these metrics. Specifically, Robust Momentum V-MoF has an almost 2% increase and 2 mCE decrease on IN-C. justifying the advantage of our method. Furthermore, we visualize the top-1 accuracy and mCE across increasing severity of two corruption types in Fig. 14 in Appendix E.3 to illustrate the increasing effectiveness of our method with escalating data corruption. The results of Robust MomentumSMoE on WikiText-100 103 can also be found in Appendix E.4, Table 9.
Soft Mixture Experts (Soft MoE). We use the Soft MoE-tiny with layers. The first 6 layers consist of standard VIT blocks, and the last 6 layers replace the MLP in those blocks with Soft MoE layer. We train Momentum-Soon MoE and a baseline Soft MoE model on ImageNet- 1K and present their results in Table 3. In addition, we plot the training loss and top-1 accuracy of both models for 120 training epochs Fig. 4(Right). Notably, there is a considerable increase in the accuracy of Momentum-Sode MoE over the bascline, as well clear acceleration to good solution during

---

Table Top-1/top-5 accuracy (%) of Momentum-Se SSft MoE vs. Soft MoE baseline on ImageNet-Il (IN-IK).
Valid IN-IK
Model
Params
Top-1 1
Top-5
Sofi MoE (baxeline)
231M
73.52
90.94
Momentum-Soon MoE
231M
75.47
92.34
Tuning 7
SMoE
MomentumSMoE
AdamSMoE
0.01
(A)
(B)
(C)
6.1
0.1
Proportion of times each
0.5
expert is chosen
4.5
Tuning /
1 0.9
5.5
0.7
0.5
0.3
10
Experts, ordered from largest smallest norm
Epoch
Figure 5: Left: Proportion of each expert chosen, ordered from the largest norm of each expert output to the smallest norm, in layers 3 and 5 of SMoE, Mom,ntumSMoE, and Adam SMoE, averaged over the WikiText-1003 103 validation set. Right: Log validation perplexity (PPL. during the finetuning of hyperparameters, μ and 7, for 40 training epochs in MomentumSMoB. When tuning y keep = 0.7 and vice versa with y = 1.0.
training. These findings justify the benefits and universality of our momentum based method, that extends beyond SMoE to MoE.
Empirical Analysis
We conduct empirical analysis based on the SMoE-medium trained on WikiText- 103 Section 5.1.
Norm-basee Load Imbalance. Load imbalance in SMoE occurs when only a small subset of experts consistently selected [63, 77]. Numerous methods have been developed counter this common phenomenon, such as introducing buffer capacity for each expert and load balancing loss [65, 18). Orthogonal these, in line with GD framework for SMoEs, we examine the choice of experts determined by the size of the norm of their outputs, (z)]]-
From multi-objective optimization perspective, the optimal descent direction is that minimizes the norm in the convex hull of the normalized gradients U (see Section 2.1). l). the gradients are not normalized, the minimum norm direction is then expected to be mainly influenced by the gradients with the smallest norms [12]. From our GD analogy of SMoE in Section 2.3, the gradients correspond to the experts, whose outputs not normalized. We then visualize the proportion of times the experts in a SMoE are chosen according to their norms during inference in Fig. 5(Left, A). We exactly connection observe the corresponding phenomenon in the SMoE, further empirically justifying between SMoE and GD. The full plots for all layers and all models are provided in Appendix C.
The direction with the smallest norms frequently related to the objectives that have already had substantial degree of convergence and insufficient for a balanced minimization of the multi-objective criteria. In this light, an ideal SMoE output would have norm-based balanced choice of experts and should translate to improved model performance. Indeed, in Section 5. we established the superior performance of MomentumSMoE and AdamSMoE large-scale WikiText- 103 language modeling task, and in Figure ((Left, B, C), this directly correlates with significantly more balanced selection of experts with respect their norms.
Ablation Study on Momentum y and Step Size 7. To better understand the effects of the momentum parameter and step size on the performance of the trained Mo models, we do an ablation

---

study on these two hyperparameters and include results in Fig. 5(Rgght) which contains plot of log validation PPL during 40 training epochs. We notice that MomentumSMoE is robust to the choice of N, and we select J = 0.7 for the final comparison with the baseline SMoE. On the other hand, when the value of 7 is too small, there is an adverse effect on the model. Hence, we select y = 1.0.
Making Momentum and Step Size Learnable. We note that additional time and effort are required tune the momentum parameter and step size. Thus, we explore different methods to circumvent this limitation. We discuss the results of one such method, making 7 and learnable parameters, in this section and include results in Table in Appendix E.1. We leave the discussion on other methods to Appendix E. As shown in Table 5 in Appendix E. MomentumSMoB (it), with learnable 7 and fixed p. significantly outperforms the baseline for both clean validation and test data by at least 1.5 PPL. even surpassing the tuned model in Table 1. On the attacked data, the benefits of MomentumSMoF are further enhanced with more than 2 PPI. improvements. These results confirm that the benefits of our model can be leveraged with minimal effort. Since the MomentumSMoF. is robust to the choice of jr, we do not consider the setting of fixing y and making μ learnable.
Other Optimizers. We study the integration of other advanced momentum and optimization meth- ods, such the Nesterov accelerated gradient [48], RMSProp [68], and sharpness-awe aware minimiza- tion [64], into our MomentumSMoE framework in Appendix E.4.
Related Work
Sparse Mixture of Experts. SMoE has been extensively studied enhance the training efficiency of large language models (LLMs), with various stable routing strategies proposed, including (i) allowing tokens to select the top-k experts [35, 18, 79, 11], (ii) allowing experts select the top-k tokens [77], and (iii) globally determining expert assignment [36, 9]. Recent works have also tried to enhance the robustness of SMoE. [55] study the robustness of SMoE for ViTs [15] while [75] investigates the robustness of SMoE for CNNs. Furthemore, [37] explores the potential of SMoE for domain generalization, and [22] employs SMoE for robust multi-task learning. Various works have also focused on addressing load imbalance in SMoE, including [63, 77, 7]. Our momentum-basee SMoE can be casily incorporated into these methods above further improve their performance.
Deep Learning Models with Momentum. Momentum has been utilized in the design of deep neural network (DNN) architectures [72, 38). [23] applies momentum create large and consistent dictionaries for unsupervised learning using a contrastive loss, with a momentum- based moving average of the queue encoder at the core of this approach. Many DNN-based methods for sparse coding have been designed by unfolding classical optimization algorithms, such as FISTA [4], where momentum is used in the underlying optimizer [45]. In addition, [38] introduces momentum into ResNet and DenseNet, [73, 49] integrate momentum into neural differential equations. [52] incorporates momentum into transformers, and [50] designs RNNs using cemeentum-acceleaate first-order optimization algorithms.
8 Concluding Remarks
In this paper, we propose MomentumSMoEE a new class of SMoE that utilizes heavy-ball momentum to stabilize and robustify SMoE. We theoretically justify the stability of our MomentumSMoE models compared the SMoE baseline. Furthermore, we demonstrate that our momentum-base based design framework for SMoE is universal and can incorporate advanced momentum-based optimization methods, including AdamW [33, 42] and Robust Momentum [10], into many existing SMoE models. We empirically validate the advantage of our momentum-based SMoE over the standard SMoE baseline on WikiText-101 and ImageNet-IK. As shown in Table in Appendix E.2, momentum has no positive effect on the small SMoE of only 3 layers but attains an increasing improvement with the medium and large models of 6 and 12 layers, respectively. This is expected as each layer represents an iteration of GD. A limitation of MomentumSMoE is that while beneficial for larger models, for models that have few layers, MomentumSMME has little effect. From a theoretical perspective, would be intriguing to develop a theory to explain the enhanced robustness of MomentumSMoF, Furthermore, MomentumSMoF.co can be analyzed as a fixed-point iteration. We leave these theoretical developments future work.
10

---

Acknowledgments and Disclosure of Funding
This research project is supported by the National Research Foundation Singapore under the Al This research / project is supported by the Ministry of Education, Singapore, under the Academic Research Fund Tier (FY2023) (A-8002040-00-000, A-8002033--0-000 This research / project is also supported by the NUS Presidential Young Professorship Award (A-0009807-01-00).
Singapore Programme (AISG Award No: ATSG2-TC--2023-11--SCI
References
[1] Yutong Bai, Xinyang Geng, Karttikeya Mangalam, Amir Bar, Alan L Yuille, Trevor Darrell, Jitendra Malik, and Alexei A Efros. Sequential modeling enables scalable learning for large vision models. In Proceedings of the IEFE/CVF Conference on Computer Vision and Pattern Recognition, pages 22861-22877 2024.
[2] Hangbo Bao, Li Dong, Songhao Piao, and Furu Wei, BEit: BERT pre-training of image transformers. In International Conference on Learning Representations, 2022.
[3] Hangbo Bao, Wenhui Wang, Li Dong, Qiang Liu, Owais Khan Mohammed, Kriti Aggarwal, Subhojit Som, Songhao Piao, and Furu Wei. Vlmo: Unified vision-language pre-training with (oxuur--f-moddlll---epprrr Advances in Neural Information Processing Systems, 35:32897- 32912, 2022.
[4] Amir Beck and Mare Teboulle. A fast iterative shrinkage-threshoolinng algorithm for linear inverse problems. SIAM journal imaging sciences, 2(1):183-202: 2009.
[5] Tom Brown, Benjamin Mann, Nick Ryder, Mclanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
[6] Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling, Advances in neural information processing systems, 34:15084-15097, 2021.
[7] Zewen Chi, Li Dong. Shaohan Huang, Damai Dai, Shuming Ma, Barun Patra, Saksham Singhal, Payal Bajaj, Xia Song, Xian-Ling Mao, et al. On the representation collapse of sparse mixture of experts. Advances in Neural Information Processing Systems, 35:34600-34611 2022.
[8] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. Journal of Machine Learning Research, 24(240):1- 113, 2023.
[9] Aidan Clark, Diego de Las Casas, Aurelia Guy, Arthur Mensch, Michela Paganini, Jordan Hoffmann, Bogdan Damoc, Blake Hechtman, Trevor Cai, Sebastian Borgeaud, et al. Unified scaling laws for routed language models. In International conference on machine learning, pages 4057-4086. PMLR, 2022.
[10] Saman Cyrus, Bin Hu, Bryan Van Scoy, and Laurent Lessard. A robust accelerated optimization algorithm for strongly convex functions. In 2018 Annual American Control Conference (ACC), pages 1376 1381. IEEE, 2018
[11] Damai Dai, Li Dong, Shuming Ma. Bo Zheng, Zhifang Sui, Baobao Chang, and Furu Wei. StableMoE: Stable routing strategy for mixture of experts. In Smaranda Muresan, Preslav Nakov, and Aline Villavicencio, editors, Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 7085-7095, Dublin, Ireland, May 2022. Association for Computational Linguistics.
[12] Jean-Antoine Désidéri. Multiple-gradien descent algorithm (mgda) for multiobjective opti- mization. Comptes Rendus Mathematique, 313(5)3313331, 2012.

---

[13] Jacob Devlin, Ming-Woi Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre- Pre-training of deep bidirectional transformers for language understanding. In Jill Burstein, Christy Doran, and Thamar Solorio, editors, Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Technologies, Volume I (Long and Short Papers), pages 4171-4186, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics.
[14] Samuel Dodge and Lina Karam. A study and comparison of human and deep learning recogni- tion performance under visual distortions. In 2017 26th international conference on computer communication and networks (ICCCN), pages 1-7. IEEE, 2017.
[15] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021.
[16] Nan Du, Yanping Huang, Andrew M Dai, Simon Tong, Dmitry Lepikhin, Yuanzhong Xu. Maxim Krikun, Yanqi Zhou, Adams Wei Yu, Orhan Firat, Barret Zoph, Liam Fedus, Maarten Bosma, Zongwei Zhou, Tao Wang, Emma Wang. Kellie Webster, Marie Pellat, Kevin Robinson, Kathleen Meeier-Hellstem Toju Duke, Lucas Dixon, Kun Zhang, Quoe Le, Yonghui Wu, Zhifeng Chen, and Claire Cui. GLaM: Efficient scaling of language models with ceature-of-experrs In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 5547-5569. PMLR, 17-23 Jul 2022.
[17] Gintare Karolina Dziugaite and Daniel M. Roy. Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data. In Proceedings of the 33rd Annual Conference Uncertainty in Artificial Intelligence (UAI), 2017.
[18) William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23(120):1-39, 2022.
[19] Robert Geirhos, Kantharaju Narayanappa, Benjamin Mitzkus, Tizian Thieringer, Matthias Bethge, Felix A Wichmann, and Wieland Brendel. Partial success in closing the gap between human and machine vision. Advances in Neural Information Processing Systems, 34:23885- 23899, 2021.
[20] Gauthier Gidel, Reyhane Askari Hemmat, Mohammad Pezeshki, Rémi Le Priol, Gabriel Huang, Simon Lacoste-Julien, and Toannis Mitliagkas. Negative momentum for improved game dynamics. In Kamalika Chaudhuri and Masashi Sugiyama, editors, Proceedings the /wenty-Second International Conference on Artificial Intelligence and Statistics, volume 89 of Proceedings of Machine Learning Research, pages 1802-1811. 1. PMLR, 16-18 Apr 2019.
[21] Gabricl Goh. Why momentum really works. Distill, 2017.
[22] Shashank Gupta, Subhabrata Mukherjee, Krishan Subudhi, Eduardo Gonzalez, Damien Jose, Ahmed H Awadallah, and Jianfeng Gao. Sparsely activated mixture-of-expprrs are robust multi-task learners. arXiv preprint 2rXiv:2204.076689, 2022.
[23] Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEF/CVF conference on computer vision and pattern recognition, pages 9729-9738. 2020.
[24] Dan Hendrycks, Steven Basart, Norman Mu, Saurav Kadavath, Frank Wang, Evan Dorundo, Rahul Desai, Tyler Zhu, Samyak Parajuli. Mike Guo, et al. The many faces of robustness: critical analysis of nut-of-distribution generalization. In Proceedings of the IEEF/CVF international conference computer vision, pages $340-8349, 2021.
[25] Dan Hendryeks and Thomas Dietterich. Benchmarking neural network robustness to common corruptions and perturbations. In International Conference on Learning Representations, 2019.
12

---

[26] Dan Hendrycks, Kevin Zhao, Steven Basart, Jacob Steinhardt, and Dawn Song. Natural adversarial examples. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 15262-15271, 2021.
[27] Geoffrey Hinton. Neural networks for machine learning lecture 6e, 2014.
[28] Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. Advances in neural information processing systems, 34:1273-
1 1286, 2021.
[29] Yiding Jiang*, Behnam Neyshabur*, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantastic generalization measures and where to find them. In International Conference on Learning Representations, 2020.
[30] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint SEIV:22001088366 2020.
[31] Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On brrge-batch training for deep learning: Generalization gap and sharp minima. In International Conference on Learning Representations, 2017.
[32] Nitish Shirish Keskar and Richard Socher. Improving generalization performance by switching from adam to sgd. arXiv preprint arXiv:
2017.
[33] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Interna- tional Conference on Learning Representations, 2015.
[34] Kenichi Kumatani, Robert Gmyr, Felipe Cruz Salinas, Linquan Liu, Wei Zuo, Devang Patel, Eric Sun, and Yu Shi. Building great multi-lingual teacher with sparsely-gated mixture of experts for speech recognition. arXiv preprint arXiv: 1 12.05820, 2021.
[35] Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. (GS] Scaling giant models with conditional computation and automatic sharding. In International Conference on Learning Representations, 2021.
[36] Mike Lewis, Shruti Bhosale, Tim Dettmers, Naman Goyal, and Luke Zettlemoyer. Base layers: Simplifying training of large. sparse models. In International Conference on Machine Learning. pages 6265-6274. PMLR, 2021.
[37] Bo Li, Yifei Shen, Jingkang Yang. Yezhen Wang, Jiawei Ren, Tong Che, Jun Zhang, and Ziwei Liu. Sparse mixture-of-experts are domain generalizable learners. In The Eleventh International Conference on Leurning Representations, 2023.
[38] Huan Li, Yibo Yang, Dongmin Chen, and Zhouchen Lin. Optimization algorithm inspired deep neural network structure design. In Asian Conference on Machine Learning, pages 614-629. PMLR, 2018
[39] Junnan Li, Dongxu Li, Silvio Savarese, and Steven Hoi. Blip-2: Bootstrapping language-image pre-training with frozen image encoders and large language models. In International conference on machine learning, pages 19730-19742. PMLR, 2023.
[40] Haotian Liu, Chunyuan Li, Qingyang Wu, and Yong Jae Lec. Visual instruction tuning. In seirth-seventh Conference on Neural Information Processing Systems, 2023.
[41] Jonathan Lorraine, David Acuna, Paul Vicol, and David Duvenaud. Complex momentum for optimization in games. In Gustau Camps-Val Valls. Francisco J. R. Ruiz, and Isabel Valera, editors, Proceedings of The 25th International Conference on Artificial Intelligence and Statisties, volume 151 of Proceedings of Machine Learning Research, pages 7742-7765. PMLR, 20-30 Mar 2022.
[42] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference Learning Representations, 2019. 2019.
13

---

[43] Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. In International Conference on Learning Representations, 2017.
[44] Igor Molybog, Peter Albert, Moya Chen, Zachary Ve David Esiobu, Naman Goyal, Punit Singh Koura, Sharan Narang, Andrew Poulton, Ruan Silva, et al. theory on adam instability in large-scale machine learning. arXiv preprint 2304 2023.
[45] Thomas Moreau and Joan Bruna. Understanding the learned iterative soft thresholding algorithm with matrix factorization. arXiv preprint urXiv:1706.0 3338, 2017.
[46] John Morris, Eli Lifland, Jin Yong Yoo, Jake Grigsby, Di Jin, and Yanjun Qi. Textattack:
A framework for adversarial attacks, data augmentation, and adversarial training in nip. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, pages 119-126, 2020.
[47] Yu. Nesterov. Gradient methods for minimizing composite functions. Mathematical Program- ming, (20(11:125-1611) 2013.
[48] Yurii Nesterov. A method for solving the convex programming problem with convergence rate of (1/K2). Proceedings of the USSR Academy of Sciences, 269:543-547, 1983.
[49] Nghia Nguyen, Tan Minh Nguyen, Huyen K. Vo, Stanley Osher, and Thicu Vo. Improving neural ordinary differential equations with nesterov's s accelerated gradient method. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[50] Tan Nguyen, Richard Baraniuk, Andrea Bertozzi, Stanley Osher, and Bao Wang. Momentum- mn: Integrating momentum into recurrent neural networks. Advances in Neural Information Processing Systems, 33:1 13:1994-1936, 2020.
[51] Tan Nguyen, Tam Nguyen, Hai Do. Khai Nguyen, Vishwanath Saragadam, Minh Pham, Khuong Duy Nguyen, Nhat Ho, and Stanley Osher. Improving transformer with an admixture of attention heads. Advances in neural information processing systems, 35:27937-27952, 2022.
[52] Tan Minh Nguyen, Richard Baraniuk, Robert Kirby, Stanley Osher, and Bao Wang. Momentum transformer: Closing the performance gap between self-attention and its linearization. In Mathematical and Scientific Machine Learning, pages 189-204. PMLR, 2022.
[53] Boris Polyak. Some methods of speeding up the convergence of iteration methods. Ussr computational mathematics and mathematical physics, 4(5):1-17, 1964.
[54] Ofir Press, Noah A. Smith, and Omer Levy. Improving transformer models by reordering their sublayers. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 2996-3900000 -005, Online, July 2020. Association for Computational Linguisties.
[55] Joan Puigcerver, Rodolphe Jenatton, Carlos Riquelme, Pranjal Awasthi, and Srinadh Bhojana- palli. On the adversarial robustness of mixture of experts. Advances in Neural Information Processing Systems, 35:9660 1, 2022.
[56] Joan Puigcerver, Carlos Riquelme Ruiz, Basil Mustafa, and Neil Houlsby. From sparse soft mixtures of experts. In The Twelfth International Conference on Learning Representations, 2024.
[57] Ning Qian. On the momentum term in gradient descent learning algorithms. Neural Networks, 1251),145-151, 1999.
[58] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 18):9, 2019.
[59] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research, 21(140):1-67, 2020.
14

---

[60] Benjamin Recht, Rebecca Roelofs, Ludwig Schmidt, and Vaishaal Shankar. Do imagenet classifiers generalize to imagenet? In International conference on machine learning, pages 1389.5500. PMLR, 2019.
[61] Carlos Riquelme, Joan Puigcerver, Basil Mustafa, Maxim Neumann, Rodolphe Jenatton, André Susano Pinto, Daniel Keysers, and Neil Houlsby. Scaling vision with sparse mixture of experts. Advances in Neural Information Processing Systems, 11888338899 2021.
[62]
Carlos Riquelme, Joan Puigcerver, Basil Mustafa, Maxim Neumann, Rodolphe Jenatton, André Susano Pinto, Daniel Keysers, and Neil Houlsby. Scaling vision with sparse mixture of experts. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 8583 8595. Curran Associates, Inc., 2021.
[63] Clemens Rosenbaum, Ignacio Cases, Matthew Riemer, and Tim Klinger. Routing networks and the challenges of modular and compositional computation. arXiv preprint arXiv: 127744 2019.
[64] Vincent Roulet and Alexandre d'Aspremont. Sharpness, restart and acceleration. In I. Guyon, U. Von Luxburg. S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors. Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
[65] Noam Shazeer, AAzalia Mirhoseini, *Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gate gated mixture-of- experts layer. In International Conference on Learning Representations, 2017.
[66] Shashank Subramanian, Peter Harrington, Kurt Keutzer, Wahid Bhimji, Dmitriy Morozov, Michael W Mahoney, and Amir Gholami. Towards foundation models for scientific machine learning: Characterizing scaling and transfer behavior. Advances in Neural Information Pro- cessing Systems, 36, 2024.
[67] Ilya Sutskever, James Martens, George Dahl, and Geoffrey Hinton. On the importance of initialization and momentum in deep learning. In International conference on machine learning. pages 1139-1147, PMLR, 2013.
[68] Tijmen Tieleman and Geoffrey Hinton. Rmsprop: Divide the gradient by running average of its recent magnitude. coursera: Neural networks for machine learning. COURSERA Neural Networks Mach. Learn, 17, 2012.
[69] Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Herve Jegou. Training data-efficient image transformers & distillation through attention. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 10347-10357. PMLR, 18 28-2411 Jul 2021.
[70] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timo- thée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302. /3971, 2023.
[71] Bao Wang, Tan Nguyen, Tao Sun, Andrea L Bertozzi, Richard G Baraniuk, and Stanley J Osher. Scheduled restart momentum for accelerated stochastic gradient descent. SIAM Journal on Imaging Sciences, 15(2):338-76) 2022.
[72] Bao Wang, Hedi Xia, Tan Nguyen, and Stanley Osher. How does momentum benefit deep neural networks architecture design? a few case studies. Research in the Mathematiall Sciences, 9(3):57, 2022.
[73] Hedi Xia, Vai Suliafu, Hangjie Ji, Tan Nguyen, Andrea Bertozzi, Stanley Osher, and Bao Wang. Heavy ball neural ordinary differential equations. Advances in Neural Information Processing Systems, 34:18646-18659 2021.
[74] Liu Yang, Siting Liu, Tingwei Meng, and Stanley J Osher. In-context operator learning with data prompts for differential equation problems. Proceedings of the National Academy of Sciences, 120039)):22101122222 2023.
15

---

[75] Yihua Zhang, Ruisi Cai, Tianlong Chen, Guanhua Zhang. Huan Zhang. Pin- Yu Chen, Shiyu Chang, Zhangyang Wang, and Sijia Liu. Robust -ixture-of-expen training for convolutional neural networks. In Proceedings of the IEEF/CVF International Conference Computer Vision, pages 90-101, 2023.
[76] Yanqi Zhou, Nan Du, Yanping Huang, Daiyi Peng. Chang Lan, Da Huang. Siamak Shakeri, David So, Andrew M. Dai. Yifeng Lu, Zhifeng Chen, Quoc V Le. Claire Cui, James Laudon, and Jeff Dean. Brainformers: Trading simplicity for efficiency. Proceedings of the 40th International Conference Machine Learning, volume 202, pages 42531 42542.. PMLR. 23-29 Jul 2023.
[77] Yanqi Zhou, Tao Lei, Hanxiao Liu, Nan Du, Yanping Huang, Vincent Zhao, Andrew M Dai, Quoe V Le, James Laudon, et al. Mixture-of-expetts with expert choice routing. Advances in Neural Information Processing Systems, 35:7103-7114, 2022.
[78] Barret Zoph, Irwan Bello, Sameer Kumar, Nan Du, Yanping Huang, Jeff Dean, Noam Shazeer, and William Fedus. St-moe: Designing stable and transferable sparse expert models. arXiv preprint arXiv:2202.08906 2022.
[79] Simiao Zuo, Xiaodong Liu, Jian Jiao, Young Jin Kim, Hany Hassan, Ruofei Zhang, Jianfeng Gao, and Tuo Zhao. Taming sparsely activated transformer with stochastic experts. In International Conference Learning Representations, 2022.
16

---

Supplement to "MomentumSMoE: Integrating Momentum into Sparse Mixture of Experts"
Table of Contents
A Proof of Lemma 1
17
Interpretation of Robust Momentum
19
Empirical Evidence
19
D Experiment Details
20
WikiText-103 Language Modeling
D.2 ImageNet-IK Object Recognition
22
D.3 Pseudocode
23
Additional Experimental Results
24
E.1
Hyperparameters
24
E.2
WikiText-103 Language Modeling
25
E.3
ImageNet-C Full Results
25
E.4 Further Extensions Beyond Adam and Robust Momentum
25
E.5
Comparison of Computational Efficiency
29
F Broader Impacts
29
Proof of Lemma 1
We can find the eigenvalues of matrix A explicitly as
(1 1 (1+ + - 4u
(0(n))²
[1)
2
If (n) < O,
A (1+1-70(1)) + / ((1 + (0())) 44
2
(1+ + + 21² 4p
2
(1 1)) + |1 p
2
max(11
Then, must have pa (n)
0. Letting the expression in the square root be A, expanding it,
A = (1 + μ - (n)))) - 4u
= (1 - (= + - (n(n)) + 4²² 4p
7001))) + (1 (a(
and finding the roots of the equation yiclds, μr (1 2 Then, consider two cases.
17

---

Case 1: A 2 0 when p 2 (1 v or S (1 vna(n))
la: If f 2 (1+ (n)))) then 1 + 1 2 1++++ 8 2 2+ 2 2 0.
Then, 21,2 > 0 and À 2 X2. For the system to converge, we need
<1 = - \ ( V
(19)
2
However, by assumption, we have a contradiction
(n))) + - <0. - - -111-+t00mm 1 0 5 V
1 (1+
Hence, we cannot have je 2 (1 + (T)))
1b: If S (1- - ( we further divide into two more subcases,
1bi: If μl - (n) 2 0, again have, 21. and x 2 22- By Eqn. 19, require
5 VĀ= < 11 (n). As 1 (n) 2 (1
= 2 2 >
is enough to check that A (1 - - (n())² can be satisfied.
A < (1 (1-1+70(n)) +
(1 + μ - r0(1))² - 4u < - (1-
(1 + p -
< <Ap
(1+
(n)(()) +1-1+7((1l)
- 1-
(1 - -L (o()) < 4u
4(p - (n)) < 4/
0 < ya (n)
Therefore, the conditions for convergence are i) S (1 (1- (n))" ii) 1+ 2 0 and i (n)) > 0. Then by it), / 2 (n) - N -1. Suppose for contradiction that
2 1, then by t). 1 S μ S (1 - (n)(2)) which implies that yo(n) 2 (n) and
hence, (n) 2 4. Combining this with condition i), we have vu S (na(n) - 1 which leads to the following contradiction with ii)
# S (1 - Vro(n))2
= μ 1 1 + va(n)
= μ 1+2((/7(n)) - 1) S (o(n) -
2+2 S 500 -
1+ + 2/P S 0.
Then, we must have a last condition for convergence, ir) |u| <1
1bii: If 1 + / - (n) S O, then |>2| N |u and
(1 + μ - (o()) VA
A
<0
2
For the system to be stable, we need d2 > -1 + 3
> VA 2 0. Hence, we
require i) μ S (1 ())(,)) 1i) + (n) s 0, (i) - > 0 and iv)
(3 (3+)~ A (3 (L==)
4u 41=>14=--ooa33330----111 44
4p> (1+ (1+ 3
+/ 17--00h (n())
- p +
ro(n)
4u -4(2 /1 va
+ 2/ - (n)) > 0
Combining ii and iv), we have, 2+ 2p - 1 > 0 = p> 1. Similarly, we assume for
contradiction that 2 1. Then, ya (n) 2 1+ p 2 and by condition i), v S. (no(n)- -1
18

---

from which follows the contradiction
S Vio(n) Vi Vra(n) + V S (Vro(n) vT) (V + vh) = 70(n) - -p
Vna(n) + v (a(n) < (by iii)
V + (n) <<2 2 < (by i) again).
Therefore, we need v) |x| < for convergence as well.
To summarize, for Case 1, in order to have stable system, require:
a) P S (1 (n)(n))
b) I μ - (o 20, (m >
i |p| <1
c)
If + (o(n)) 50
i)
+ 1 - (o(n) >
2 24 (n)) > i | <1
ii)
+ 2p 2o > 0
ini)
| <
Case 2: A < O when (1 + (n)) p (1 (n))
Then, the eigenvalues are complex and given by
+ L M (n)) =i 4p (1 11+p-1 -0
For convergence, we require A1 = Vi- < 1 or equivalently, μ < 1. Hence, the necessary condition
becomes 1 >> μ > (1-
yo (n) and to avoid contradiction, we also require > - (1))
This is satisfied when (n) + 2p and |r| < 1.
Therefore, from all the considered cases, to have stable system, we require 0 < < < 2 24 and |u|| 1, concluding the proof.
B Interpretation of Robust Momentum
To provide intuition behind robust momentum, we follow [10] and view the update rule in Eqn. 16 as Lur'e feedback control system. For simplicity, we consider the case where, Dt R and write the update in matrix form
Ftt Fe+1
= (-p
(1
xt-1 24
+ (2)
(xt-1 f(v)); yt = 1 a)
(20)
1 1
In the frequency domain, in order for the system 20 to be stable, we require that for all |2| =1
Re((1 p2 ((:: 1)G(pz) 1)) 0.
where G(pz) is a transformed transfer function, and p, a scaling factor. The imaginary axis, where Re the stability boundary, and the specific construction of the parameters V, ll and pushes this boundary into the negative real axis to -v = (1 + (1(p))11 k + 2kp kp²)/2p /2p, hence achieving robust stability through the design of y. p and a.
C Empirical Evidence
In this section, we provide the full plots presented in Section 2.3 and 6. We visualize the average norm the SMoE and MoE outputs respectively, for 80 epochs of training in all layers in Fig. and 7. Notably, in all plots, less minor exception in layer of SMoE, there is decrease in the norm of each SMoE and MoE layer output throughout training. This consistent with our expectation that the gate learns the optimal a" in multi-objective optimization problem (Eqn. 8 and 9).
19

---

Training
Validation
Layer
Layer 2
Layer
0.48
0.34
0.90
0.30
0.92
[[ ||
0.58
20
40
60
80
20
40
60
60
80
Layer
Layer
Layer
0.92
0.90
0.93
0.88
0.05
0.85
0.58
0.00
0.84
0.53
20
40
50
80
20
40
60
22
40
50
80
Epochs
Figure 6: Average norm of the SMoE outputs at all layers during 80 epochs of training for the baseline SMoE.
Training
Validation
Layer
Layer 2
Layer
0.55
0.00
1.00
0.80
0.80
0.85
0:75
0:20
0.70
Il (a:
40
60
80
20
40
Layer
Layer
Layer
1.00
1.00
0.80
0.80
0.50
0.70
0.60
0.60
0:40
0.0
Epochs
Figure Average norm of the MoE outputs at layers during 80 epochs of training for the baseline MoE.
Fig. 8, and 10, we present the proportion of time each expert is chosen decreasing magnitude of the norm of the expert outputs for SMoE, MomentumSMoE and AdamSMoE. As discussed in Section 6, in accordance with our connection between GD and SMoE, a more even distribution of expert selection based the size of the norm of the expert outputs should yield improved performance of the model. We demonstrate Section 5.1 that both MomentumSMoF. and AdamSMoE, significantly exceed the baseline performance and correspondingly, flattens the normed-based load distribution among experts as observed Fig. 8, and 10. These strong empirical evidences serve to reinforce our optimization framework for SMoE.
D Experiment Details
For clarity, we summarize the new models developed in this paper for each task Table 4 and address the lacking implementations here. First, as the introduction of AdamSMoF into V-MoE leads to unstable training, we defer this challenge to future work. Second, our primary goal when studying the Momentum-Son Soft MoE model is to showcase the universality of our MomentumSMoE method
20

---

SMoE
Layer
Layer 2
Layer
30
30
30
Proportion of times each expert chosen
20-
10-
10-
10
5
10
15
5
10
15
10
15
Layer
Layer 5
Layer 6
30-
20
10
Experts, ordered from largest to smallest norm
Figure 8: Proportion each expert chosen, ordered from the largest norm of each expert output to the smallest norm, in all layers of baseline SMoE.
MomentumSMoE
Layer
Layer
Layer 3
30
Proportion of times each expert chosen
20
10
10
15
5
10
15
5
10
15
Layer 4
Layer
Layer
30
30
30
10
15
Experts, ordered from largest to smallest norm
Figure 9: Proportion each expert chosen, ordered from the largest norm of each expert output to the smallest norm, in all layers of MomentumSMoF.
Table 4: A summary of the new models developed for WikiText- 103 language modeling and ImageNet- 1K object recognition tasks.
Task
Model/Method
MomentumSMoB
AdamSMoE
Robust MomentumSMoF
SMoE
WikiText-103
GLaM
ImageNet- 1K
ImageNet-R
V-MoE
ImageNet-A
ImageNet-C
ImageNet-IK
Soft MoE
across both SMoE and MoE, hence we did not integrate Adam W and Robust Momentum into Soft MoE. We leave these implementations for future work.
D.1
WikiText-1033 Language Modeling
Dataset: The WikiText- 103 dataset [43] derived from Wikipedia articles and is designed to capture long-range contextual dependencies. The training set contains about 28,000 articles, with
21

---

AdamSMoE
Layer
Layer
Layer
30
30
transformer and Proportion of times each expert chosen
10-
10-
10-
10
5
1.5
5
15
Layer 4
Layer
Layer 6
20
20
20
Experts, ordered from largest to smallest norm
Figure 10: Proportion of each expert chosen, ordered from the largest norm of each expert output to the smallest norm, in all layers of AdamSMoE.
total 103 million words. Each article is divided into text blocks with approximately 3,600 words. The validation and test sets have 218,000 and 246,000 words, respectively, with both sets comprising 60 articles and totaling about 268,000 words. On the attacked dataset, we corrupt the both validation and test data demonstrate the robustness of MomentumSMoE and AdamSMoE using TextAttack's 's word swap attack [46]. This adversarial attack randomly replaces words in the dataset with generic "AAA" for evaluation making it difficult for the model to predict the next word in the sequence correctly.
Model and baselines: We use the Switch Transformer [18], referred as SMoE in our tables and figures, and GLaM [16] baselines, which replaces each multilayer perceptron (MLP) layer and every other MLP layer in a vanilla language modeling transformer with SMoE layer, respectively. For MomentumSMoF, we replace each MLP layer with MomentumSMoF layer and initialise each momentum vector, Po
We the same for MomentumGLaM with every other MLP layer.
For consistency, we define the number of layers in each model as the number of SMoE layers. The default model used in each experiment is medium sized with layers, but we include comparison between a smaller model with 3 layers as well as a larger one with 12 layers in Appendix E.2. Each model has 16 experts every SMoE, MomentumSMoE and AdamSMoE layer and selects experts (K = 2) per input. All models use the same sparse router function consisting of a linear network receiving the input data followed by the TopK, then the Softmax function. The small models train for 60 epochs, the medium and large SMoE models train for 80 epochs and the GLaM models train for 120 epochs without any additional load balancing loss. Our implementation is based on the code base developed by [54], publicly available at http:///ghhuubbooooooeeessee http://gihuub.com/iingggg wtt.r/aiaaaiiiiit1oooeeeeeettt D/CCompeteSMo///re/eeeneen
AdamSMoF/AdamGLaaM It is observed that Adam has certain divergent behavior during large- scale training leading to unstable loss curves [8, 44). In line with this observation, during the initial implementations of AdamSMoE,
experience similar instability. A widely used solution to this, proposed by [32], is to switch from Adam to gradient descent at a suitable point during training. We follow suit and AdamSMoE the first layer of the model and MomentumSMoE for the subsequent layers. We find that implementing AdamSMoE in the first layer is enough to significantly improve the model's overall performance and accelerate its convergence.
D.2 ImageNet-1K Object Recognition
We use the ImageNet-11 IK dataset that contains 1.28M training images and 50K validation images. There are 1000 classes of images and the model learns an object recognition task. For robustness to common corruptions, we use ImageNet-C (IN-C) [25] which consists of 15 different types of corruptions applied the ImageNet- validation set with levels of severity. We provide breakdown of our results on each corruption type and the mean Corruption Error (mCE) across
Datasets:
22

---

Hyperpa aaperparametees: mu
def MomentumSMoE(x, momentum):
momentum SMoE(x) + mu momentum
X = X + gamma * momentum
return
Figure 11: Pseudocode for MomentumSMof. implementation mmn in python with hyperparameter μ.
Hyperparameters: mu, beta, eps 1e-8
AdamSMoE (x, gradient, squared squared_gradiente grl_ggrrdddn::: gradient squared_gradient beta * squared gradient
mu) * SMoE(x) + mu momentum
squared_gradient
beta) * SMoE(x) **
x = x + gamma / storch..srt(ssqurrdd_oo (rraiiieeddrrdettte + eps) gradient
return
Figure 12: Pseudocode for AdamSMoE implementation python with hyperparameters ja and B.
escalating levels of severity for corruption types in Appendix E.3. To test robustness to input data
distribution shifts, we use ImageNet-A (IN-A) [26]. IN-A contains 200 class subset of ImageNet-1K K
classes with adversarially filtered images. Finally, test our model ImageNet-R (IN-R) [24] which contains various artistic renditions of images. This evaluates the model's generalization ability abstract visual renditions.
Metries: On ImageNet-IK, ImageNet-A and ImageNet-R, we report the top-1 accuracies for all experiments. On ImageNet-C, the standard metric for evaluation the mCE. To calculate this, we average the top-1 error rate for each corrupiion type across the 5 levels of severity and divide them by AlexNet's average errors, then take the final average across all corruption types. The direction of increasing decreasing values of these metrics signifying greater robustness will be indicated in the table with an arrow.
Model and baselines: We use small Vision Mixture of Experts (V-MoE) [62] model the SMoE baseline for ImageNet-IK K object recognition task well as the standard robustness benchmarks. This variant of V-MoE consists of Vision Transformer (ViT) blocks [15] with every odd block's MIP being replaced by a SMoE layer. In turn, in Momentum V-MoF, we replace every other MIP layer with a MomentumSMoE layer and similarly for Robust Momentum V-MoE. For all vision SMoE models, we select 2 experts (K = 2) at every SMoE layer for each patch. We follow the training configurations and setting as in [62] and their code base available here
For our MoE baseline on clean ImageNet-1K 1K data, we use Soft Mixture of Experts (Soft MoE) [56]. Soft MoE model is designed to side step the challenging discrete optimization problem assigning each token to an expert, as in SMoE. through a soft token assignment. Instead of each token being routed to one expert, in a Soft MoE layer, each expert is assigned certain number of slots and each slot is allocated weighted average of all tokens. These weights are dependant on both the tokens and the experts. In this case, Soft MoE not considered as a SMoE, but instead a MoE. We use the smallest model, Soft MoE-tiny with 12 layers. The first 6 layers consists of standard ViT blocks and the last layers replace the MLP those blocks with a Soft MoE layer. In Momentum- Soft MoE, we implement the momentum parameter into each Soft-MoE layer as in SMoE layer.
As there were no training details provided for Soft MoE on ImageNet- IK, we follow the training proce- dure in [69] for 120 epochs, using their published code base htt.ggii..o//cccbboooooooooooooe We train Momentum-Sonn MoE and the baseline model using the PyTorch implementation of Soft MoE at http:://ggthub.oom/wwwoo//////d
D.3 Pseudocode
In this section, we provide the pseudocode written in python for MomentunSMoF, AdamSMoE, and Robust MomentumSMMFF for clarification on our implementation. These are found Figures 11, 12 and 13 respectively.
23

---

Hyperparameters: P, L,
def RobustNomentuunSMoEKK momentum):
= LLL/ / m
gamma k * (11 p) ** 2) * (1 + p) /
mu = k * p ** 3 / (k 1
alpha x* 3 / ((k 1) ((1 p) ** 2) (1 + p))
= alpha * gamma * momentum
momentum SMoE(y) + mu * momentum
gamma * momentum
return x
Figure 13: Pseudocode for Robust MomentumSMof implementation in python with hyperparame- ters p. L and m.
Table 5: Perplexity (PPL) results on clean and attacked WikiText- 103 validation and test data for standard MoE with tuned hyperparameters μr and y and MomentumSMoE trained with different learning settings for and that do not require tuning: i) Both μ and yy are scalar learnable
parameters Pytorch, ūi) Only y is learned with fixed μ = 0.7.
Model
WikiText-103 Valid PPL +
Test PPL
Clean
33.76
35.55
SMoE (buseline)
Attacked
42.24
44.19
Clean Attacked
32.29
33.46
MomentumSMoE
0.7
1.0
40.94
42.33
Clean
32.28
33.87
MomentumSMoE (i)
Learnale
Learnale
Attacked
40.41
42.32
Clean Attacked
32.28
33.69
MomentumSMoB (ii)
0.7
Learnable
39.96
41.84
E Additional Experimental Results
E.1
Hyperparameters
In this section, we discuss two additional methods to avoid tuning MomentumSMoE hyperparameters for ease of implementation and efficiency.
Time-varying momentum: Another form of the classical momentum proposed by [48] replaces the constant momentum parameter with a time-varying one t 1/t + 2, which removes the need to choose an appropriate momentum hyperparameter. We adopt this modification into MomentumSMoE to replace μ while keeping y fixed 1.0, and the MomentumSMoE time-varying (TV) layer as follows
P: -(((x,) (or)++ + t-1
[ P:-15 Øt+1 3t + PPt
Zero-order hold and 7: Recall Eqn. 10, the proposed accelerated MomentumSMoE layer
p. = - f(x) + up:-1i Ft++ 3; + -p.
When replace MLP layers with MomentumSMoF layers, as do in language model Mo- MentumSMoE, each expert function a linear network. Explicitly expressing the MomentumSMoE layer results in the following expression
P: = Σ g( W. Jz, + PP:-1 Bx, + PP:-1i F+41
x, + YP:
and can be interpreted state space representation with constant state and output matrices, l and ty From this perspective, we experiment with learning discretized and by applying the Zero-order hold (ZOH) such that they become adaptive parameters is a common practise optimization algorithms, j and y are then parameterized as follows
H4 = = Ah
(Ap) ( Ac 1)4y
24

---

Table 6: Perplexity (PPL) results on clean and attacked WikiText-100 103 validation and test data for standard MomentumSMoE with tuned hyperparameters with different learning settings for ja and 7 that not require tuning: i) Both lr and ? are scalar learnable parameters in Pytorch. iř) Only 7 is learned with fixed 0.7. üii) Only is learned as linear network then composed with a sigmoid function with fixed μ.. Included are the time-varying MomentumSMoE (TV) and Zero-order hold MomentumSMoE (ZOH).
and and MomentumSMof trained
Model
WikiText-100 103 Valid PPI. +
Test PPI.
Clean
33.76
35.55
SMoE (baseline)
Attacked
42.24
44.19
Clean Attacked
32.29
33.46
MomentunsMME
0.7
1.0
40.94
42.33
Clean
32.28
33.87
MomentumSMoE (i)
Learnable
Learnablc
Attacked
40.41
42.32
Clean
32.28
33.69
MomentumSMoE (ii)
0.7
Learnablc
Attacked
39.96
41.84
Clean
32.30
33.96
MomentumSMoF.cii
0.7
Linear network
Attacked
40.35
42.39
Clean
33.02
35.08
MomentumSMME (TV)
1/t
1.0
Attacked
41.44
43.97
Clean
32.21
34.00
MomentumSMoFF (ZOH)
Ax
(eAr 1)47/AH
Attacked
40.74
42.70
where A the Softplus function of learned scalar parameter, μa learned scalar parameter and y linear network with scalar outputs.
We report the PPL for MomentumSMoE (TV) and MomentumSMoE (ZOH) clean and word swap attacked WikiText- 103 validation and test data in Table 6, an extended version of Table 5. We also include an additional setting of learning y the sigmoid of linear network. In this setting, every input has a different learning rate. While these models do improve over the baseline, there is no advantage the standard MomentumSMME. and the other learnable μ and 7 settings. Hence, do not recommend implementations in this section over those and keep the results here for reference.
Results:
E.2 WikiText-100 Language Modeling
We present the results of all three sizes of SMoE and MomentumSMoE models in Table 7 and observe that with increasing model depth, the effectiveness of the momentum parameter in improving model performance increases well. This aligns with the analogy of each layer being GD step and with higher number of iterations, the momentum term becomes more effective. While beneficial for large models, implementing MomentumSMoE is computationally cost efficient, and hence, minimally affected by model depth, we note that there is an adverse effect small models such as SMolE-sall with only 3 layers. We hypothesize that the primary reason for this the insufficient layers for the momentum term to make positive impact and is limitation of our method.
E.3 ImageNet-CC C Full Results
We provide the full results of the top-1 accuracy and mCE of all 15 corruption types in ImageNet-C for V-MoE, Momentum V-MoE, Robust Momentum
and Sharpness Aware ComenntuVV-Mo (SAM-V-MoE), developed in Appendix E.4, in Table 8. Included in Figure 14 is plot of the mCE with escalating severity of impulse noise and gaussian noise corruption, which illustrates the advantages of our methods with increasing data corruption. We observe that across all 15 corruption types, except for motion blur, Robust Momentum V-MoE outperforms the baseline V-MoE, with as high 6.5% increase in top-1 accuracy and 8 mCE decrease fog corruption.
E.4 Further Extensions Beyond Adam and Robust Momentum
The advantage of an optimization perspective for SMoFs are the countably many descent algorithms available that can be used to improve the model. In this section, we elaborate on six more extensions
25

---

Table 7: Perplexity (PPL)
baseline sooo=-smal/mmeduu//argg.. MomentumSMoE.
small/medium/large. MlamSMoE-mediuse baseline GLaM-medium, MomentumGLaM M-MMimm and AdamGlaMMmmeding on clean and attacked WikiText- 103 validation and test data.
Clean WikiText-103
Attacked WikiText-103
Model/Metric
Valid PPL
Test PPL
Valid PPL +
Test PPL
SMoE-small (baseline)
84.26
84.81
98.60
99.29
MomentumSMoF smalll small
85.71
86.65
100.26
101.18
SMoE-medium (baseline)
33.76
35.55
42.24
44.19
MomentumSMot Mooeiummm---------um-m medium
32.29
33,46
40.94
42.33
AdamSMoE medium
31.59
33.25
39.27
41.11
GLaMM-edium (baseline)
36.37
37.71
45.83
47.61
MomentumGi
33.87
35.29
42.15
43.64
AdamGI aM-medium
32.99
34.32
41.09
42.81
SMoE-larg (baseline)
29.31
30.33
36.77
37.83
MomentumSMoFF.
27.58
29.03
35.21
36.78
Table 8: Top-1 I accuracy (%) and mean corruption error (mCE) of V-MoE, Momentum V-MoE, Robust MomentumV-Mo V-MoE and SAM-V-MoE on each corruption type in ImageNet-C.
Model/
V-MoE
Robust MomentumV-MoE
Corruption Type
Momentum V-MoB
SAM-V-MoE
Metric
(Baseline)
Top-1 1
67.25
67.77
67.79
67.34
Brightness
mCE +
58.01
57.08
57.06
57.85
Top-1
36.94
40.94
42.71
41.64
Contrast
mCE 4
73.91
69.22
67.15
68.41
Top-1
39.89
41.01
41.02
40.26
Defocus Blur
mCE I
73.31
71.95
71.94
72.86
Top-1
53.14
53.08
$3.19
52.89
Elastic Transform
mCE 4
72.53
72.63
72.45
72.92
Top-1
38.23
42.00
44.85
44.03
Fog
mCE1 4
75.39
70.79
67.31
68.31
Top-1
50.12
51,51
51,83
51.45
Frost
mCE
60.35
58.66
58.27
58.73
Top-1
49.38
50.92
52.08
$0.92
Gaussian Noise
mCE/
57.11
55.37
54.06
55.36
Top-1
36.65
36.74
37.30
36.43
Glass Blur
mCE +
76.67
76.56
75.88
76.93
Top-1
47.72
49.30
50.59
49.15
Impulse Noise
mCEJ 1
56.66
54.95
53,56
55.11
Top-1
60.21
60.44
60.53
60.23
JPFG Compression
mCE 1
65.61
65.23
65.08
65.58
Top-1
43.59
43.10
43.35
42.08
Motion Blur
mCE 1
71.77
72.40
72.08
73.69
Top-1
61.66
62.61
63.14
62.30
Pixelate
mCE! 1
53.41
$2.09
51.35
$2.52
Top-1
47.96
49.03
50.51
49.66
Shot Noise
mCE +
58.18
56.98
55.33
56.28
T
36.91
37.31
37.32
37.12
Snow
mCE +
72.78
72.32
72.31
72.54
Top-1
35.03
35.88
36.14
35.20
Zoom Blur
mCE 1
81.38
80.31
79.99
81.17
to MomentumSMoE, namely, Nesterov accelerated gradient [48], time varying momentum with scheduled restart [64, 47], RMSprop [27]. sharpness aware minimization [64] (SAM), negative momentum [20], and complex momentum [41]. We report their results on clean and word swap attacked WikiText-I10s language modeling task Table The nature SAM is to find local minima where the geometry of the loss landscape is flat to improve the generalization ability of the model. This aligns with the objective improving the robustness of models distribution shifts. Hence, implement SAM in V-MoE well and provide comparison with V-MoE, Momentum V-MoE
26

---

V-MoE
MomentumV-Mot
Robust MomentumV-MoE
SAM-V-MoE
Impulse noise
Gaussian noise
80
80
MCE
70
70
60
50
50
5
2
Severity
Severity
Figure 14: mCE (lower is better) of baseline V-MoE, Momentum V-MoE, Robust NomentumVVMol and SAM-V-MoE on increasing severities of impulse noise and gaussian noise corruption. As the severity increases, the effect of momentum, robust momentum and SAM becomes increasingly apparent.
Table 9: Perplexity (PPL) results on clean and word swap attacked WikiText-100 103 validation and test data for baseline SMoE, NAG-SMoE, SR- RM-E, rms-SMoE, SAM-SMoE, Robust MomentumSMoE, Negative-MomenumSSMe.. and Complex-No Seaaple--MomentumS.MM MoE.
Clean Wiki 103
Attacked WikiText- 103
Model/Metric
Valid PPL |
Test PPL
Valid PPL
Test PPL J
SMoE (baseline)
33.76
35.55
42.24
44.19
NAG-SMoR
33.83
35.46
41.94
43.97
SR-SMoE
32.96
35.01
41.21
43.72
rms-SMoE
32.43
34.25
40.58
42.60
SAM-SMoE
33.39
35.05
41.47
43.44
Robust MomentumSMoE
33.22
34.45
41.49
42.78
Pegativ--MomenuunSMoE
33.48
35.09
41.68
43.62
Complex- Mompllx-MMmennunSMMtt
32.08
33.34
40.24
41.66
and Robust Momentum V-MoE in Table 10. In all models, we replace all SMoE layers with their corresponding newly derived layer unless stated otherwise.
Nesterov accelerated gradient (NAG): Nesterov accelerated gradient (NAG) [48] takes the mo- mentum method step further by looking ahead to where the momentum term will move the parameters and providing a correction. The foresight of the update prevents the parameters from moving an undesirable direction too quickly, preventing large oscillations especially when the learning rate is high. We implement NAG in our SMoE gradient framework as
Pt ff= /PP:11) + PPt-1; E++1
+
where (x,) is the SMoE output and μ E (0, 1) and > O are two hyperparameters corresponding to the momentum coefficient and step size respectively. refer to this implementation as NAG- SMoE.
Time-varying momentum with scheduled restart: An enhancement of the time-varying momen- tum covered in Section E. is to include a scheduled restart for the momentum parameter, 1/t + 2. Such modification can help to recover an optimal convergence rate as accelerated methods do not necessarily maintain monotonic decrease in objective value [64, 47). As our model has 6 layers, we choose restart after layer 3 and the RS-SMoE update
t mod 3
P: -f((
mod 3 + Pt-1; Lt+1 = Dz + YP
where f(x)) is the SMoE output and
0 is the step size.
RMSprop: RMSprop an adaptive step size algorithm that scales the gradient by an exponentially decaying average of the squared gradients [27]. The algorithm replaces global learning rate, which
27

---

Table 10: Top-1 accuracy (%) and mean corruption (mCE) of V-MoE, MomentumVVMooe
Robust MomentumV-Min and SAM-V-MoE on validation and test ImageNet-IK data and popular standard robustness benchmarks for image classification.
Valid IN-IK
Test IN-IK
IN-R
IN-A
IN-C
Model
Top-1 f
Top-1 f
Top-1
Top-1 f
Top-1 f
mCE +
V-MoE (baseline)
76.49
73.16
36.10
5.25
46.98
67.14
Momentum V-MB
76.92
73.26
37.45
6.48
48.11
65.77
Robust Momentum V-MoE
76.66
73.20
37.57
6.37
48.82
64.92
SAM-V-MoE
76.26
72.84
36.64
6,27
48.05
65.88
difficult to tune due to the widely differing magnitudes of the gradients of each parameter, with one that adapts to the individual parameter gradients. In addition, by keeping rumning average of the squared gradients, we avoid extreme fluctuations in the step size due to stochastic batch updates. Incorporating rmsprop into the SMoE optimization perspective, have the following update in rms-SMoE layer.
Pt = PP-1 + (4)(2))) Et+1 Zt
7
f(x))
VP:
where f(x) is the SMoE output, μ E (0, 1) is the moving average parameter and, 7 > the step size. Similar to AdamSMoF, we experience some instability when using rms-SMoE to replace all SMoE layers. Hence, we follow suit, and only replace the first layer with rms-SMoE and replace subsequent SMoE layers with MomentumSMoE.
Sharpness-Aware Minimization: In a standard machine learning approach, training models with the usual optimization algorithms minimize the training empirical loss, which is typically non-convex. This results in multiple local minima that may have similar loss values but lead to vastly different generalization abilities. As such, models that perform well on training data, may still have inferior validation performance. From studies relating the geometry of the loss landscape to generalization, specifically, and intuitively, flatter minima should yield greater generalization ability [17, 29, 31]. Further, this would increase the model's 's robustness to distribution shifts in input data or labels.
The sharpness aware minimization (SAM) algorithm [64] aims to take advantage of this relationship and seek not just any local minimum during training, but a minima whose neighbourhood also has uniformly low loss values, in other words, lower sharpness, to improve robustness. Implementing the algorithm with the 1z-norm in the SMoE optimization framework yields the SAM-SMoB layer
{(x;)= (f)))//ff
p: f(x, + ((x;)); Et+1 x₁ 7/P
where fxx) is the SMoE output, >> 0 is a hyperparameter controlling the size of the neighbourhood and, y the step size.
Complex and Negative Momentum: Classic momentum works well in a gradient-basee optimiza- tion when the Jacobian of our update step, as a Exed-point operator, has real eigenvalues. However, this might not always be the case. Negative momentum, which simply chooses negative momentum parameters, is preferred by operators with complex eigenvalues [20]. In situations where the spectrum is purely imaginary or has mixtures of complex and imaginary eigenvalues, complex momentum is more robust than both classic and negative momentum [41]. This is due to oscillations at fixed frequencies between adding and subtracting the momentum term. We oscillate the sign of the momen- tum term by choosing a complex momentum parameter and then updating our weights by only the real part of the momentum term. This translates to the following update in the Complex-Momenturn layer:
P: -- PP:-1 F++1 3, RPt:),
(21)
where μ E C and > 0 are hyperparameters corresponding to the momentum coefficient and step size, respectively and Rt extracts the real component the momentum term.
Results: On the language modeling task, though all models, except NAG-SMoE clean WikiText- 103 validation data, do outperform the baseline across all metrics, most have only marginal gains that fall short of the performance gap achieved with MomentumSMoE and AdamSMoE. This is
28

---

Table 11: Run time per sample, memory and number of parameters of MomentumSMoE and AdamSMoE as compared to the baseline SMoE during training and test time.
Sec/Sample (Training)
Sec/Sample (Test)
Memory
Memory (Test)
Model
Parameters
(Training)
SMoE (baseline)
0.0315
0,0303
22168MB
17618MB
216M
MomentumSMoF
0.0317
0.0304
22168MB
17618MBB
216M
AdamSMoF 0.0321
0.0307
22168MB
17618MB
216M
Table 12: Total computation time for SMoE, MomentumSMoF and AdamSMoF to reach 38 PPI. on WikiText-103 validation data.
Model
Time (minutes)
SMoE (baseline)
85,56
MonentumSMoF
81.77
AdamSMoE
84.23
with the exception of Complex-MomenntnSSSttE which has an even greater improvement over the baseline than MomentumSMoE. Complex-MomentumsSe
enhanced results further verify the power and promise of our momentum- -baed framework in designing SMoE. On the computer vision task, we find that SAM-V-ME does perform relatively well on corruption benchmarks, exceeding the baseline by more than 1% on ImageNet-A and ImageNet-C.
E.5
Comparison of Computational Efficiency
common consideration when introducing modified layers into deep models is the potential increase in computational overhead. We aim to alleviate that concern by providing the run time per sample, memory and number of parameters of MomentumSMof: and AdamSMoE as compared the baseline SMoE during both training and test time in Table 11. We also provide the total computation time required for all models to reach the same PPL level on WikiText- 103 validation data in Table 12. We observe that MomentumSMoE and AdamSMoE are comparable to the baseline SMoE across all metries at both training and test time with negligible computational cost.
F Broader Impacts
Our research enhances both clean data handling and robust performance, particularly in socially impactful domains. Notably, we demonstrate improved results in object recognition, benefiting self-driving cars, and language modeling, enhancing Al chatbot assistants. We show significant advancements in resisting data perturbation, aiming to protect critical AI systems from malicious actors. Furthermore, achieve competitive performance in language modeling with contaminated data, reflecting weal-world scenarios where data often imperfect. While the potential for Al misuse exists, work provides substantial improvements in fundamental architectures and theory, which we hope will lead to further socially beneficial outcomes.
29

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The claims made in the abstract and introduction are clearly stated in the Contribution in the Introduction. These claims accurately reflect the paper's contributions and scope.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
The limitations are discussed in the Conclusion.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-specfiiatiinn asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
30

---

Answer: [Yes]
Justification: All theoretical results in the paper are given together with the full set of assumptions and complete/correct proofs (See Section 3 and Appendix A in our manuscript).
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: We provide the experiment details in the Experiment Details Section in the Appendix of our manuscript. We also provide the source code so that the results in the paper can be easily reproduced.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
31

---

5. Open access to data and code
Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Question:
Answer: [Yes]
Justification: We provide the source code so that the results in the paper can be easily repro- duced. We verify our proposed methods using public benchmarks (See the Experimental Results Section in our manuscript)
Guidelines:
The answer NA means that paper does not include experiments requiring code. Please see the NeurIPS code and data submission guidelines (ttips//niip.cc/public
Quides/CoodeSbbmissioooooiici for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to to reproduce the results. See the NeurIPS code and data submission guidelines (htips: hetppuulllgguiieeSoeeeuuuuieeeee
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/eetails
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We specify all the training and test details necessary understand the results in the Experiment Details Section in the Appendix of our manuscript.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification:
We report error bars suitably and correctly defined of the experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emoor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
32

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative еπτог rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We provide sufficient information on the computer resources for all experi- ments in our Experimental Results Section.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics http:://euuupscccccuuuucettteete
Answer: [Yes]
Justification: The research conducted in the paper conforms, in every respect, with the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g., s if there is special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
We discuss broader impacts in Appendix F.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
33

---

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), faimess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: The paper poses no such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but encourage authors to take this into account and make best. faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification:
We cite the githubs we and the baselines compare with in manuscript.
Guidelines:
The answer NA means that the paper does existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
34

---

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswitheode. com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: The paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moode as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: The paper does not involve research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
35

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
36


Reviews:
[{"review_id": "kSw5XLnjw8", "reviewer": "Reviewer_qEGD", "summary": "This paper addresses the instability problem of training SMoE models. By establishing a relationship between SMoE and multi-objective optimization, the authors integrate momentum into SMoE and propose MomentumSMoE. Experimental results show that MomentumSMoE is more stable than SMoE during training.", "questions": "please refer to weaknesses", "limitations": "The authors have adequately discussed the limitations.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper tackles a critical issue in the training of SMoE models.\n\n2. The proposed method is generalizable and can be applied to various SMoE models such as V-MoE and GLaM.\n\n3. Experimental results demonstrate that this method is more stable than SMoE during the training process.", "weaknesses": "1. This method has little effect on models with few layers.\n\n2. The largest models for evaluation only have  388M parameters, which are much smaller than mainstream MoE LLMs.\n\n3. From a theoretical standpoint, developing a framework to explain the enhanced robustness of MomentumSMoE would be interesting.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Q1. This method has little effect on models with few layers.**\n\n**Answer:** A momentum-based approach like MomentumSMoE needs more layers to show its advantages, just like a heavy-ball momentum, or Adam needs a couple of iterations to start showing its faster convergence compared to gradient descent. Also, models with few layers are not common in practice due to their worse performance compared to models with more layers. For example, as can be seen in Table 7 in our manuscript, the SMoE-small baseline has much lower validation and test perplexity (PPL) compared to SMoE-medium and large baselines, 84.26/84.81 PPL vs. 33.76/35.55 and 29.31/30.33 PPL, respectively. The poor PPL of the SMoE-small baseline renders its performance on the WikiText-103 language modeling tasks negligible.\n\n**Q2. The largest models for evaluation only have 388M parameters, which are much smaller than mainstream MoE LLMs.** \n\n**Answer:** Thanks for your comment. During the rebuttal, we have tried our best to conduct an additional experiment on a Sparse Mixture of Experts (SMoE) backbone with more than 1B parameters. However, due to the shortage of computing resources and the short time of the rebuttal, our trainings have not finished yet. We will include these results of our momentum-based SMoE on a larger SMoE backbone in the revised manuscript and during the discussion period if possible. \n\n**Q3. From a theoretical standpoint, developing a framework to explain the enhanced robustness of MomentumSMoE would be interesting.**\n\n**Answer:** Thanks for your suggestion. We agree with the reviewer that developing a framework to explain the enhanced robustness of MomentumSMoE would be interesting. The approach to theoretically proving the robustness of MomentumSMoE is to consider a minimizer $x^*$ of the objective function $F(x)$ in Eqn. (3) in our manuscript, which is found by heavy-ball updates when starting at the initial point $x_0$. Then, we prove that by tuning the momentum $\\mu$ and step size $\\gamma$, the heavy-ball iterations will find a point $x_k$ such that $\\||x_k - x^*\\|| \\le c\\rho^{k}$ for a constant c, $\\rho \\in [0, 1)$, and $k \\ge 1$. This result implies that the output of MomentumSMoE is close to $x^{*}$ as long as the input data is close to $x_0$, thus verifying the robustness of MomentumSMoE. This proof can be extended from the proof of Proposition 1 and Corollary 1 in our manuscript. We will include the detailed proposition and its proof for the robustness of MomentumSMoE in our revision.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "x39QRlktc9", "reviewer": "Reviewer_oJG3", "summary": "The paper introduces a novel approach to enhancing the robustness and stability of Sparse Mixture of Experts (SMoE) models. Inspired by the analogy of gradient descent and SMoE, the authors develop a family of models by incorporating momentum into the training process. The key idea is that training SMoE is a multi-objective optimization problem where the monument-based gradient descent method is more stable and robust than the vanilla one. They proposed the AdamSMoE and Robust MomentumSMoE, which demonstrate improved performance across a variety of tasks, including language modeling and object recognition.", "questions": "(1) Please explain more of line 140 (\"Thus, it is expected that these two terms learn to reduce ...).", "limitations": "N/A", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "(1) The integration of momentum into SMoE is a non-trivial innovation that addresses instability and inefficiency issues in existing models.\n\n(2) The paper provides convincing empirical evidence showing the effectiveness of MomentumSMoE across multiple benchmarks.\n\n(3) The proposed method's compatibility with other momentum-based optimizers, like Adam, suggests it can be broadly applied to various SMoE architectures.", "weaknesses": "(1) Formulating SMoE as a multi-objective optimization problem is doubtful to me. Every expert network is continually changing during the model training, which makes each objective nonstatic, which violates the basic assumption of multi-objective optimization, whose objectives should be very clear and stable. \n\n(2) It is unconvincing to use ||f(x)|| as the key metrics to measure the efficacy of SMoE or MoE. This confuses me a lot. Please explain why the output norm represents the goodness/badness of the model.\n\n(3) There are some grammar issues. Please use `` instead of \" in the paper (line 665).\n\n(4) There is no sufficient discussion of computation overhead. Training efficiency is a critical issue for current foundation model training. Does computation significantly increase by applying momentum over the SMoE? Keeping an additional copy weight (p in Fig 1) would take additional memory and may decrease the throughput.\n\nI'd like to hear a more insightful discussion regarding all the points above from the authors.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Q1. Formulating SMoE as a multi-objective optimization problem is doubtful to me.** \n\n**Answer:** We believe there is a misunderstanding of our formulation of SMoE as a multi-objective optimization problem. Please allow us to clear this misunderstanding by clarifying the role of expert networks in our multi-objective optimization (MOO) framework for Sparse Mixture of Experts (SMoE). We rewrite the minimization problem in Eqn. (3) as follows:\n\n$$\n\\min_{x\\in D, \\theta \\in \\Theta} F(x, \\theta) := \\sum_{i=1}^E c_i F_i(x, \\theta^{i}).\n$$\nwhere $\\theta^{i}$ is the parameters of the expert network $i$, $\\theta = \\{\\theta^{1},\\dots,\\theta^{E}\\}$, and $\\Theta$ is the parameter space. Also, $D$ is the feasible region, and $c_i \\in \\mathbb{R}$, $i=1,\\dots,E$, are weights representing the importance of each objective function. In order to solve this optimization problem, we employ the alternating minimization approach as follows:\n\n*Step 1 - fix $\\theta$ and minimize $F$ w.r.t. $x$*: \n\n$$\nx_{t+1} = x_{t} - \\gamma \\sum_{i=1}^{E}\\alpha_{i}^*\\nabla_xF_i(x_t,\\theta^{i}_{t}) \n$$\n\n$$\n= x_{t} - \\gamma \\sum_{i=1}^{E}\\alpha_{i}^* f_{i}(x_{t}, \\theta_{t}^i),\n$$\n\nwhere $\\alpha^*=(\\alpha_1^*,\\dots,\\alpha_E^*)$ satisfy the Pareto-stationary condition in Definition 1 in our manuscript, $\\gamma$ is the step size, and $f_{i}=\\nabla_xF_i$.\n\n*Step 2 - fix $x$ and minimize $F$ with respect to $\\theta$* \n$$\n\\theta_{t+1}^{i} = \\min_{\\theta^{i} \\in \\Theta^{i}}  c_i F_i(x_{t+1}, \\theta^{i}), i=1,\\dots,E.\n$$\nwhere $\\theta^{i}$ lies in the parameter space $\\Theta^{i}$. \n\nIn our formulation of SMoE as a multi-objective optimization problem, we regard $-u_i(x_t)$, where $u_i$ is the $i^{\\text{th}}$ expert network,  as the gradient $\\nabla_{x}F_i(x_t,\\theta^{i}_{t})$ in Step 1, and the score $\\text{softmax}(g_i(x_t))$ from the router $g(x_t)$ is learned to approximate $\\alpha_i^*$. Step 1 in the alternating minimization algorithm above corresponds to forward pass of an SMoE, and step 2 corresponds to an update step of the model's parameters (i.e., similar to an M-step in an EM algorithm). Note that this parameter update step is implicitly performed at each training iteration via optimizing the objective loss at the final layer of the model using backpropagation and gradient descent. Results in Fig 2. and the discussion from line 137 to line 144 in our manuscript provide supporting evidence for this implicit parameter update.\n\nIn Step 1, the gradient $\\nabla_{x}F_i(x_t,\\theta_{t}^{i})$ is a function of $\\theta_{t}^{i}$ and is supposed to change during the model training, which matches the observation that expert network is continually changing during the model training. Also, the objective function of our MOO problem is $F(x, \\theta)$ in which we want to find $x^*$ and $\\theta^*$ to minimize $F$. This objective function is static, clear, and stable.\n\n**Q2. Why $||f(x)||$ as the key metrics to measure the efficacy of (S)MoE?** \n\n**Please explain more of line 140.** \n\n**Answer:** In Fig. 2 and 3 in our manuscript, $||f(x)||$ is not used as a key metric to measure the efficacy of SMoE/MoE. Instead, the results in Fig. 2 and 3 are to justify the connection between SMoE/MoE and gradient descent. In particular, from Eqn. (8) and (9), the output of (S)MoE corresponds to $-f(x_t)$, where $f(x_t) = \\nabla_xF_(x_t)$ and is the gradient of the objective function $F$ with respect to $x$ at $x_t$. If (S)MoE indeed performs (stochastic) gradient descent in its forward pass, then the norm of the (S)MoE output must decrease when t (i.e., layer index) increases since the gradient norm $\\||f(x_t)\\||$ decreases when gradient descent updates are applied. Fig. 3 confirms this by showing that the norm of the (S)MoE output decreases over layers in a 6-layer (S)MoE model trained on the WikiText-103 language modeling task. At the last layer, the norm increases might be due to overshooting, a common phenomenon that can occur when using gradient descent.\n\nAdditionally, from line 116 to line 119 in our manuscript, we hypothesize that the scores $\\text{softmax}(g_{i}(x_t))$ for MoE and $\\text{softmax}(\\text{TopK}(g_{i}(x_t)))$ for SMoE from the router $g(x_t)$ are learned to approximate $\\alpha_i^*$. If this is indeed true, then in order to satisfy the Pareto-stationary condition in Definition 1 in our manuscript, the scores should be learned to minimize the norm $\\|\\sum_{i=1}^E \\alpha_i \\tilde{f_i}\\|$ (see line 93 - 99 in our manuscript), which is equivalent to the norm of the (S)MoE output. Fig. 2 verifies this expectation by showing that each MoE and SMoE layer learns to reduce its output norm during training, suggesting that the scores $\\text{softmax}(g_i(x_t)) $ and $\\text{softmax}(\\text{TopK}(g_{i}(x_t)))$ are learned to approximate $\\alpha_i^{*}$. These explain line 140.\n\n**Q3. Grammar issues:** We have addressed the grammar issues in our revision.\n\n**Q4. Insufficient discussion of computation overhead.** \n\n**Answer:** The reviewer might have overlooked Appendix E.5 of our manuscript. Table 11 in E.5 reports the training/inference runtime, training/inference memory usage, and the number of parameters of our MomentumSMoE/AdamSMoE vs. the SMoE baseline. Our MomentumSMoE/AdamSMoE have the same training/inference memory usage and number of parameters as SMoE. Furthermore, our momentum models are on par with the SMoE baseline in terms of training/inference runtime. Compared to the SMoE baseline, our MomentumSMoE (see Fig. 1) only needs an additional momentum state $p$, which is accumulated from layer to layer. Thus, MomentumSMoE only uses an additional tensor shared between layers to store this momentum state $p$. Similarly, AdamSMoE only needs two additional tensors shared between layers to store $p$ and $m$. The memory overhead for introducing these additional tensors is minimal, thus explaining why memory increase and throughput decrease in our MomentumSMoE/AdamSMoE compared to the SMoE baseline are negligible.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "unwj5JN3td", "reviewer": "Reviewer_puL2", "summary": "This paper proposes a variant of sparse mixture of experts, MomentumSMoE, by incorporating momentum into the traditional sparse mixture of experts framework. The authors provide both theoretical proofs and empirical evidence demonstrating that MomentumSMoE offers greater stability and robustness compared to the standard sparse mixture of experts. Experiments on language modeling and object recognition tasks are conducted to verify the effectiveness of the proposal.", "questions": "1. Why the MomentumV-MoE and Robust MomentumV-MoE have marginal gains on clean IN-1K data, is there any in-depth analysis available on this?\n2. In the ImageNet-1K Object Recognition experiment, why was the popular top-5 accuracy metric not used, as it was in the Soft Mixture of Experts experiment?\n3. As stated in the weaknesses, the authors could provide pseudocode to better clarify their proposal.", "limitations": "The authors adequately addressed the limitations.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The idea of integrating momentum into sparse mixture of experts  is interesting. \n2. Both the theoretical proof and extensive empirical results are provided to demonstrate that the proposed MomentumSMoE is more stable and robust than SmoE; the experimental results are appealing.\n3. The code is provided.", "weaknesses": "The pseudocode may be provided to better illustrate the implementation of the proposal.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Q1. Why the MomentumV-MoE and Robust MomentumV-MoE have marginal gains on clean IN-1K data, is there any in-depth analysis available on this?** \n\n**Answer:** V-MoE's result reported in Table 2 in our manuscript is among the state-of-the-art results on clean IN-1k data for the models that have around 45M effective parameters without pretraining. To further improve this result is challenging, and the almost 0.43\\% improvement of our MomentumV-MoE vs. the V-MoE baseline is already nontrivial. Notice that both our MomentumV-MoE and Robust MomentumV-MoE significantly improve over the V-MoE baseline on robustness benchmarks, such as IN-R, IN-A, and IN-C, as shown in Table 2.  \n\n\n**Q2. In the ImageNet-1K Object Recognition experiment, why was the popular top-5 accuracy metric not used, as it was in the Soft Mixture of Experts experiment?**\n\n**Answer:** Thanks for your comment. We provide the top-5 accuracy for ImageNet-1K object recognition task in Table 1 in the attached PDF.\n\n**Q3. As stated in the weaknesses, the authors could provide pseudocode to better clarify their proposal.** \n\n**Answer:** Thanks for your suggestion. We include the pseudocode for our MomentumSMoE, AdamSMoE, and Robust MomentumSMoE below. \n~~~\nHyperparameters: mu\n\ndef MomentumSMoE(x, momentum):\n\tmomentum = - SMoE(x) + mu * momentum \n\tx = x + gamma * momentum\n\treturn x\n~~~\n\n~~~\nHyperparameters: mu, beta, eps = 1e-8\n\ndef AdamSMoE(x, gradient, squared_gradient):\n\tgradient = - (1 - mu) * SMoE(x) + mu * momentum \n\tsquared_gradient = beta * squared_gradient + (1 - beta) * SMoE(x) ** 2\n\tx = x + gamma / (torch.sqrt(squared_gradient) + eps) * gradient - k * x\n\treturn x\n~~~\n\n~~~\nHyperparameters: p, L, m\n\ndef RobustMomentumSMoE(x, momentum):\n\tk = L / m\n\tgamma = k * ((1 - p) ** 2) * (1 + p) / L\n\tmu = k * p ** 3 / (k - 1)\n\talpha = p ** 3 / ((k - 1) * ((1 - p) ** 2) * (1 + p))\n\ty = x + alpha * gamma * momentum\n\tmomentum = - SMoE(y) + mu * momentum \n\tx = x + gamma * momentum \n\treturn x\n~~~", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "OEZ0H0aYZ0", "reviewer": "Reviewer_GFpG", "summary": "The paper introduces MomentumSMoE, a novel integration of heavy-ball momentum into Sparse Mixture of Experts (SMoE) to enhance stability and robustness. It establishes a connection between SMoE and gradient descent on multi-objective optimization problems.\nThe paper demonstrates theoretical and empirical improvements of MomentumSMoE over standard SMoE across various tasks. The method is universally applicable to many SMoE models, including V-MoE and GLaM, with minimal additional computational cost.", "questions": "How can the assumptions about the fixed point operator's spectrum and the Jacobian's conservativeness be validated or relaxed in practical scenarios?\n\nAre there more general acceleration tools than momentum you might want to use for this problem?", "limitations": "The limitations are discussed briefly, but a delineated section elaborating on all the limitations would be valuable.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "To the best of my knowledge, attempting to accelerate the fixed point iteration in SMoE is an original idea.\n\nIt seems like there is comprehensive empirical evidence for the method, but I am not an expert on metrics for the SMoE, and will have to rely on other reviews to be confident in this strength.\n\nThe paper is fairly clear, with well-organized sections and figures.", "weaknesses": "My largest negative for this paper is the largely unfounded connection between the SMoE and gradient descent.  If the authors had made a connection to accelerating fixed-point iterations in general, I would want to accept this paper.  Essentially, the authors are assuming that $\\nabla_x f$ has strictly real eigenvalues when they should just work with truly, potentially complex, eigenvalues, ex., using tools as in Azizian et. al.  For example, when performing this analysis, various other acceleration schemes are often better, like negative momentum (Gidel et. al.) or complex momentum (Lorraine et. al.).  I would be curious to see some empirical investigation (or theoretical) or what the eigenvalues of $\\nabla_x f$ are – ex., as in Figure 7 of https://arxiv.org/pdf/2102.08431 -- to validate any theoretical claims about what acceleration schemes should be used.\n\nBut, of course, the spectrum is only known in small-scale problems, leading to the second weakness, which is that some of the methods – ex., RobustSMoE – seem to rely on knowing the spectrum to set various parameters, which we won’t have access in real settings.  Th\n\nThe theoretical results are also largely just reproductions of known theoretical results for momentum once you assume that the update from the SMoE is a gradient. This makes them not much of a contribution from my point of view other than leveraging existing tools. I think these results could be easily substituted for analogous techniques from Azizian.\n\nAzizian, Waïss, et al. \"Accelerating smooth games by manipulating spectral shapes.\" International Conference on Artificial Intelligence and Statistics. PMLR, 2020.\nLorraine, Jonathan P., et al. \"Complex momentum for optimization in games.\" International Conference on Artificial Intelligence and Statistics. PMLR, 2022.\nGidel, Gauthier, et al. \"Negative momentum for improved game dynamics.\" The 22nd International Conference on Artificial Intelligence and Statistics. PMLR, 2019.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Q1. Unfounded connection between the SMoE and gradient descent. Connection to accelerating fixed-point iterations. The authors should work with complex eigenvalues [Azizian et. al.] using negative (Gidel et. al.) or complex momentum (Lorraine et. al.). Empirical (or theoretical) investigation and an analysis of the eigenvalues of $\\nabla_x f$  is needed.**\n\n**Answer:** We respectfully disagree with the reviewer’s comment that the connection between SMoE and gradient descent (GD) is largely unfounded. In Section 2.3, from line 128 to line 144, we discuss the empirical evidences (Fig. 2, 3) for connection between SMoE and GD.\n\nGD updates can be considered as fixed-point iteration: at optimality where the gradient is 0, GD finds a fixed-point solution. Thus, techniques in accelerating fixed-point iterations can be used to find the equilibrium point of SMoE. Following the reviewer's suggestion, we have conducted additional experiments that incorporate negative momentum (Gidel, 2019) and complex momentum (Lorraine, 2022) into SMoE using our momentum-based design framework. Table 2 in the attached PDF compares the PPL of NegativeMomentumSMoE and ComplexMomentumSMoE with MomentumSMoE, AdamSMoE, and SMoE on the clean/attacked WikiText-103. Our AdamSMoE achieves the best PPL. ComplexMomentumSMoE obtains slightly better PPL than MomentumSMoE (positive momentum coefficient), while both of these methods significantly outperform NegativeMomentumSMoE and the SMoE baseline.\n\nFurthermore, when analyzing the eigenvalues of $\\nabla_x f$, where $f$ is the output of an SMoE layer in the trained SMoE above, we observe that the mean absolute real and imaginary parts of the eigenvalues are 4.06 and 0.34, averaged over all SMoE layers in the model, respectively. (Gidel, 2019) and (Lorraine, 2022) suggest that positive momentum converges when the eigenvalues are real, while negative momentum has better convergence when the eigenvalues have large imaginary parts. Also, the convergence of the complex momentum method is robust to the wider range of complex eigenvalues, including purely real and imaginary ones. Since the imaginary part of the eigenvalues of $\\nabla_x f$ is small compared to their real part (0.34 vs. 4.06), NegativeMomentumSMoE should not work well, and MomentumSMoE with positive momentum should be a better design choice. Since ComplexMomentumSMoE can handle a wider range of complex eigenvalues, it helps improve the performance of MomentumSMoE. However, due to the small imaginary part of eigenvalues of $\\nabla_x f$, the improvement is small. These explain our results in Table 2 in the attached PDF and we also provide a plot of the spectrum in Figure 1.\n\nWe would like to thank the reviewer for your suggestion on negative and complex momentum. The significant improvement of MomemtunSMoE, AdamSMoE, and ComplexMomentumSMoE over the SMoE baseline further verifies the power and promise of the momentum-based framework for designing SMoE. Moreover, given our framework, analyzing the eigenvalues of $\\nabla_x f$ offers a principled way to select the right momentum method for SMoE design.\n\n**Q2. Methods like RobustSMoE seem to rely on knowing the spectrum to set various parameters, which we won’t have access in real settings.**\n\n**Validating/Relaxing assumptions about the fixed point operator's spectrum and the Jacobian's conservativeness**\n\n**Answer:**  The motivation behind the design of complex momentum is to work over a large range of eigenvalues, as it is common that, in practice, we do not know the spectrum [Lorraine, 2022]. Similarly, for MomentumSMoE and AdamSMoE, we need to tune the hyperparameters, i.e., $\\mu$, $\\gamma$, as well as $\\beta$ in AdamSMoE, in order to achieve convergence.\n\nIn contrast, the Robust Momentum Method in [Cyrus, 2018] that we use to develop Robust MomentumSMoE **does not require the spectrum information**, and the convergence of the algorithm **does not depend on the spectrum values**. This Robust Momentum Method can be considered as a Lur’e feedback control system, and the hyperparameters $\\gamma$, $\\mu$, and $\\alpha$ (see Eqn. (18), our manuscript) are designed to push the stability boundary into the negative real axis. This allows an additional margin for the stability conditions of the system to hold.\n\n**Q3. The novelty of theoretical results?**\n\n**Answer:** As mentioned in line 156 in our manuscript, we are inspired by the techniques in [Qian, 1999] and adapt these techniques to prove the convergence and stability results of Momentum SMoE (Proposition 1 and Corollary 1). The techniques we use are not new, and other techniques, such as those in [Azizian, 2020] or [Wilson, 2018], can be leveraged to prove similar results. Our objective is to provide theoretical guarantees on the convergence and stability of MomentumSMoE to prove that MomentumSMoE is more stable than the SMoE baseline, which justifies the empirical advantages of MomentumSMoE over SMoE shown in our manuscript.\n\n**Q4. More general acceleration tools than momentum?**\n\n**Answer:** Momentum methods, such as Nesterov’s accelerated gradient method, are simple approximations of the proximal point method [Ahn, 2022]. For future work, we can employ the proximal point method to develop better SMoE models. In our manuscript, we also explore the sharpness-aware minimization (SAM) [Roulet, 2017] approach to enhance the robustness of SMoE (lines 805 - 821 in Appendix E.4).\n\n**References**\n\nO’donoghue, B, et al. Adaptive restart for accelerated gradient schemes. 2015.\n\nLucas, J., et al. Aggregated momentum. 2018.\n\nCyrus, S., et al. A robust accelerated optimization algorithm for strongly convex functions. 2018\n\nQian, N. On the momentum term in gradient descent learning algorithms. 1999.\n\nWilson, A. C., et al. A Lyapunov analysis of accelerated methods in optimization. 2021.\n\nAhn, K., et al.Understanding nesterov's acceleration via proximal point method. 2022.\n\nRoulet, V., et al. Sharpness, restart and acceleration. 2017.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
## Global Rebuttal

Dear AC and reviewers,

Thanks for your thoughtful reviews and valuable comments, which have helped us improve the paper significantly. We are encouraged by the endorsements that: 1) The idea of integrating momentum into Sparse Mixture of Experts (SMoE) is original, interesting, a non-trivial innovation that tackles a critical issue in the training of SMoE models (all reviewers); 2) The paper provides convincing and comprehensive empirical evidence showing the effectiveness of MomentumSMoE across multiple benchmarks (all reviewers); 3) The proposed method is compatible with other momentum-based optimizers, like Adam (Reviewer oJG3), and can be applied to various SMoE models such as V-MoE and GLaM (Reviewer qEGD). We have included the additional experimental results requested by the reviewers in the 1-page attached PDF.

One of the concerns shared by Reviewer oJG3 and GFpG is that the formulation of SMoE as a multi-objective optimization problem is not clear and the connection between SMoE and gradient descent needs justification. We address this concern here.

**Clarifying the formulation of SMoE as a multi-objective optimization (MOO) problem**

Let us start by rewriting the minimization problem in Eqn. (3) as follows:

$$
\min_{x\in D, \theta \in \Theta} F(x, \theta) := \sum_{i=1}^E c_i F_i(x, \theta^{i}).
$$
where $\theta^{i}$ is the parameters of the expert network $i$, $\theta = \{\theta^{1},\dots,\theta^{E}\}$, and $\Theta$ is the parameter space. Also, $D$ is the feasible region, and $c_i \in \mathbb{R}$, $i=1,\dots,E$, are weights representing the importance of each objective function. In order to solve this optimization problem, we employ the alternating minimization approach as follows:

*Step 1 - fix $\theta$ and minimize $F$ w.r.t. $x$*: 

$$
x_{t+1} = x_{t} - \gamma \sum_{i=1}^{E}\alpha_{i}^*\nabla_xF_i(x_t,\theta^{i}_{t}) 
$$

$$
= x_{t} - \gamma \sum_{i=1}^{E}\alpha_{i}^* f_{i}(x_{t}, \theta_{t}^i),
$$

where $\alpha^*=(\alpha_1^*,\dots,\alpha_E^*)$ satisfy the Pareto-stationary condition in Definition 1 in our manuscript, $\gamma$ is the step size, and $f_{i}=\nabla_xF_i$.

*Step 2 - fix $x$ and minimize $F$ with respect to $\theta$* 
$$
\theta_{t+1}^{i} = \min_{\theta^{i} \in \Theta^{i}}  c_i F_i(x_{t+1}, \theta^{i}), i=1,\dots,E.
$$
where $\theta^{i}$ lies in the parameter space $\Theta^{i}$. 

In our formulation of SMoE as a multi-objective optimization problem, we regard $-u_i(x_t)$, where $u_i$ is the $i^{\text{th}}$ expert network,  as the gradient $\nabla_{x}F_i(x_t,\theta^{i}_{t})$ in Step 1, and the score $\text{softmax}(g(x_t))_i$ from the router $g(x_t)$ is learned to approximate $\alpha_i^*$. Step 1 in the alternating minimization algorithm above corresponds to the forward pass of an SMoE, and step 2 corresponds to an update step of the model's parameters (i.e., similar to an M-step in an EM algorithm). Note that this parameter update step is implicitly performed at each training iteration via optimizing the objective loss at the final layer of the model using backpropagation and gradient descent. Results in Fig 2. and the discussion from line 137 to line 144 in our manuscript provide supporting evidences for this implicit parameter update.

As can be seen in Step 1, the gradient $\nabla_{x}F_i(x_t,\theta_{t}^{i})$ is a function of $\theta_{t}^{i}$ and is supposed to change during the model training, which matches the observation that expert network is continually changing during the model training. Also, the objective function of our MOO problem is $F(x, \theta)$ in which we want to find $x^*$ and $\theta^*$ to minimize $F$. This objective function is static, clear, and stable.

**Empirical evidences to justify the connection between the SMoE and gradient descent**

In Section 2.3 of our manuscript, from line 128 to line 144, we discuss the empirical evidences (see Figure 2 and 3 in our manuscript) for connection between SMoE and gradient descent. In particular, from Eqn. (8) and (9), the output of (S)MoE corresponds to $-f(x_t)$, where $f(x_t) = \nabla_xF(x_t)$ and is the gradient of the objective function $F$ with respect to $x$ at $x_t$. If (S)MoE indeed performs (stochastic) gradient descent in its forward pass, then the norm of the (S)MoE output must decrease when t (i.e., layer index) increases since the gradient norm $\||f(x_t)\||$ decreases when gradient descent updates are applied. Fig. 3 confirms this by showing that the norm of the (S)MoE output decreases over layers in a 6-layer (S)MoE model trained on the WikiText-103 language modeling task. At the last layer, the norm increase might be due to overshooting, a common phenomenon that can occur when using gradient descent.

Additionally, from line 116 to line 119 in our manuscript, we hypothesize that the scores $\text{softmax}(g_i(x_t))$ for MoE and $\text{softmax}(\text{TopK}(g_{i}(x_t)))$ for SMoE from the router $g(x_t)$ are learned to approximate $\alpha_i^*$. If this is indeed true, then in order to satisfy the Pareto-stationary condition in Definition 1 in our manuscript, the scores should be learned to minimize the norm $\||\sum_{i=1}^E \alpha_i \tilde{f_i}\||$ (see line 93 - 99 in our manuscript), which is equivalent to the norm of the (S)MoE output. Fig. 2 verifies this expectation by showing that each MoE and SMoE layer learns to reduce its output norm during training, suggesting that the scores $\text{softmax}(g_{i}(x_t))$ and $\text{softmax}(\text{TopK}(g_{i}(x_t)))$ are learned to approximate $\alpha_i^{*}$.

-----

We hope that our rebuttal have cleared your concerns about our work. We are glad to answer any further questions you have on our submission, and we would appreciate it if we can get your further feedback at your earliest convenience.

Meta-review: This paper introduces several momentum-based algorithms designed to enhance the stability and robustness of Sparse Mixture of Experts (SMoE). As many reviewers have noted, this approach is novel and holds significant potential for impact within the community. The paper effectively demonstrates both theoretical and empirical improvements of MomentumSMoE over standard SMoE across various tasks.

After personally reading the paper, I observed that MomentumSMoE aligns more closely with fixed-point iteration acceleration techniques rather than traditional momentum-based gradient descent. Additionally, as discussed in Section 3, the assumptions that the sequence $x_0, \dots, x_t $ lies in the neighborhood of the optimal point $ x^*$, and that $ x^* = 0$, are not convincingly justified in practical scenarios. Furthermore, as Reviewer GFgG pointed out, a detailed analysis of the spectrum of $\nabla_x f(x)$ would provide a deeper understanding of the effectiveness of MomentumSMoE.

Despite these concerns, the novelty and potential impact of the proposed methods are significant. Therefore, I  recommend accepting this paper.
Decision: Accept (poster)