MomentumSMoE: Integrating Momentum into Sparse Mixture of Experts

Authors: rachel s y teo
tan m nguyen
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-13 09:30:11

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.

Main paper: 
 introduction: 
 Scaling up deep models has demonstrated significant potential for enhancing the model's performance on a wide range of cognitive and machine learning tasks, ranging from large language model pretraining [13, 58, 59, 30, 5, 51, 70] and vision understanding [15, 2, 3, 39, 1, 40] to reinforcement learning [6, 28] and scientific applications [66, 74] . However, increasing the model's size requires a higher computational budget, which can be often challenging to meet. As a result, Sparse Mixture of Experts (SMoE) 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 of billion-parameter models and achieves significant success in various applications, including machine translation [35] , image classification [61] , and speech recognition [34] .A MoE replaces a component in the layer of the model, for example, a feed-forward or convolutional layer, by a set of networks termed experts. This approach largely scales up the model but increases the computational cost. A SMoE inherits the extended model capacity from MoE but preserves the computational overhead by taking advantage of conditional computation. In particular, a SMoE consists of a router and E expert networks, u i , i = 1, 2, . . . , E. For each input token x t ∈ R D at layer t, the SMoE's router computes the affinity scores between x t and each expert as g i (x t ), i = 1, 2, . . . , E. In practice, we often choose the router g(x t ) = [g 1 (x t ), g 2 (x t ), . . . , g E (x t )] ⊤ = W x + b, where W ∈ R E×D and b ∈ R E . Then, a sparse gating function TopK is applied to select only K experts with the greatest affinity scores. Here, we define the TopK function as:EQUATIONThe outputs from K expert networks chosen by the router are then linearly combined asEQUATIONwhere softmax(g i ) := exp(g i )/ E j=1 exp(g j ). We often set K = 2, i.e., top-2 routing, as this configuration has been shown to provide the 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.In this paper, we explore the role of the residual connection in SMoE and show that simple modifications of this residual connection can help enhance the stability and robustness of SMoE. In particular, we develop a gradient descent (GD) analogy of the SMoE, showing that the dynamics of the expert representations in SMoE is associated with a gradient descent step toward the optimal solution of a multi-objective optimization problem. We then propose to integrate heavy-ball momentum into the dynamics of SMoE, which results in the Momentum Sparse Mixture-of-Experts (MomentumSMoE). At the core of MomentumSMoE is the use of momentum to stabilize and robustify the model. The architecture of MomentumSMoE is depicted in Fig. 1 . MomentumSMoE can be extended beyond heavy-ball momentum to integrate well with other advanced momentum-accelerated methods such as AdamW [33, 42] and Robust Momentum [10] . Our contribution is three-fold:1. We incorporate heavy-ball momentum in SMoE to improve the model's stability and robustness.2. We theoretically prove that the spectrum of MomentumSMoE is better-structured than SMoE, leading to the model's stability enhancement.3. We show that the design principle of MomentumSMoE can be generalized to other advanced momentum-based optimization methods, proposing AdamSMoE and Robust MomentumSMoE.Our experimental results validate that our momentum-based SMoEs improve over the baseline SMoE in terms of accuracy and robustness on a variety of practical benchmarks, including WikiText-103 language modeling and ImageNet-1K object recognition. We also empirically demonstrate that our momentum-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] , just 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 MomentumSMoE. In Section 5, 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 MomentumSMoE. We discuss related works in Section 7. The paper ends with concluding remarks. More experimental details are provided in the Appendix.< l a t e x i t s h a 1 _ b a s e 6 4 = " 5 X o f M q 1 w h t g D y z a J W t M f A s s T N S R B m q r v n V 7 k Q 0 C V g I V B C l W r Y V g 5 M S C Z w K N i y 0 E 8 V i Q n u k y 1 q a h i R g y k n H 1 w / x v l Y 6 2 I + k r h D w W P 2 9 k Z J A q U H g 6 c m A w L 2 a 9 k b i f 1 4 r A f / C S X k Y J 8 B C O n n I T w S G C I + i w B 0 u G Q U x 0 I R Q y f W t m N 4 T S S j o w A o 6 B H v 6 y 7 O k f l y 2 z 8 q n t y f F y l U W R x 7 t o j 1 U Q j Y 6 R x V 0 g 6 q o h i h 6 R M / o F b 0 Z T 8 a L 8 W 5 8 T E Z z R r a z j f 7 A + P w B q N S U E w = = < / l a t e x i t >x t+1 = x t f (x t )< l a t e x i t s h a 1 _ b a s e 6 4 = " J q 7 + R G 6 G P y 9 t t v m D j V j N d Z L G P d C q j h U i h e R 4 G S t x L N a R R I 3 g y G t x O / + c i 1 E b F 6 w F H C / Y j 2 l Q g F o 2 i l 5 l M 3 w 1 N v 3 C 2 V 3 Y o 7 B V k k X k 7 K k K P W L X 1 1 e j F L I 6 6 Q S W p M 2 3 M T 9 D O q U T D J x 8 V O a n h C 2 Z D 2 e d t S R S N u / G x 6 7 p g c W 6 V H w l j b U k i m 6 u + J j E b G j K L A d k Y U BZ Q Y x h s s k p F p B 9 R y K T R v o E D J 2 7 H h V A W S t 4 L R b e 6 3 H r m x I t I P O I 6 5 r + h A i 1 A w i r n 0 1 M N S r 1 x x q + 4 U Z J l 4 c 1 K B O e q 9 8 l e 3 H 7 F E c Y 1 M U m s 7 n h u j n 1 K D g k k + K X U T y 2 P K R n T A O x n V V H H r p 9 N b J + Q k U / o k j E x W G s l U / T 2 R U m X t W A V Z p 6 I 4 t I t e L v 7 n d R I M r / 1 U 6 D h B r t l s U Z h I g h H J H y d 9 Y T h D O c 4 I Z U Z k t x I 2 p I Y y z O L J Q / A W X 1 4 m z b O q d 1 m 9 u D + v 1 G 7 m c R T h C I 7 h F D y 4 g h r c Q R 0 a w G A I z / AZ Q Y x h s s k p F p B 9 R y K T R v o E D J 2 7 H h V A W S t 4 L R b e 6 3 H r m x I t I P O I 6 5 r + h A i 1 A w i r n 0 1 M N S r 1 x x q + 4 U Z J l 4 c 1 K B O e q 9 8 l e 3 H 7 F E c Y 1 M U m s 7 n h u j n 1 K D g k k + K X U T y 2 P K R n T A O x n V V H H r p 9 N b J + Q k U / o k j E x W G s l U / T 2 R U m X t W A V Z p 6 I 4 t I t e L v 7 n d R I M r / 1 U 6 D h B r t l s U Z h I g h H J H y d 9 Y T h D O c 4 I Z U Z k t x I 2 p I Y y z O L J Q / A W X 1 4 m z b O q d 1 m 9 u D + v 1 G 7 m c R T h C I 7 h F D y 4 g h r c Q R 0 a w G A I z / A K b 4 5 y X p x 3 5 2 P W W n D m M 4 f w B 8 7 n D 6 n X j g E = < / l a t e x i t >x t < l a t e x i t s h a 1 _ b a s e 6 4 = " 2 5 P z a v i 2 u Z 9 P S C t x H 0 H w t 3 T 4 w p I = " > A A A B 7 n i c b V D L S g N B E O z 1 G e M r 6 t H L Y B A E I e y K r 2 P Q i 8 c I 5 g H J E m Y n s 8 m Q 2 d l l p l c M S z 7 C i w d F v P o 9 3 v w b J 8 k e N L G g o a j q p r s r S K Q w 6 L r f z t L y y u r a e m G j u L m 1 v b N b 2 t t v m D j V j N d Z L G P d C q j h U i h e R 4 G S t x L N a R R I 3 g y G t x O / + c i 1 E b F 6 w F H C / Y j 2 l Q g F o 2 i l 5 l M 3 w 1 N v 3 C 2 V 3 Y o 7 B V k k X k 7 K k K P W L X 1 1 e j F L I 6 6 Q S W p M 2 3 M T 9 D O q U T D J x 8 V O a n h C 2 Z D 2 e d t S R S N u / G x 6 7 p g c W 6 V H w l j b U k i m 6 u + J j E b G j K L A d k Y U BN v n C R 7 0 M S C h q K q m + 6 u I J H C o O t + O 4 W V 1 b X 1 j e J m a W t 7 Z 3 e v v H / Q N H G q G W + w W M a 6 H V D D p V C 8 g Q I l b y e a 0 y i Q v B W M 7 q Z + 6 4 l r I 2 L 1 i O O E + x E d K B E K R t F K r a S X 4 Z k 3 6 Z U r b t W d g S w T L y c V y F H v l b + 6 / Z i l E V f I J D W m 4 7 k J + h n V K J j k k 1 I 3 N T y h b E Q H v G O p o h E 3 f j Y 7 d 0 J O r N I n Y a x t K S Q z 9 f d E R i N j x l F g O y O K Q 7 P o T c X / v E 6 K 4 Y 2 f C Z W k y B W b L w p T S T A m 0 9 9 J X 2 j O U I 4 t o U w L e y t h Q 6 o p Q 5 t Q y Y b g L b 6 8 T J r n V e + q e v l w U a n d 5 n E U 4 Q i O 4 R Q 8 u I Y a 3 E M d G s B g B M / w C m 9 O 4 r w 4 7 8 7 H v L X g 5 D O H 8 A f O 5 w 8 J j I 9 j < / l a t e x i t > p t 1< l a t e x i t s h a 1 _ b a s e 6 4 = " J q 7 + R G 6 G P y 9 < l a t e x i t s h a 1 _ b a s e 6 4 = " j n T g + n c j 4 0 z m N X K h 9 j R + X b H 2 9 Y Y l c E p I = " > A A A B 6 X i c b V D L S g N B E O z 1 G e M r 6 t H L Y B C 8 G H b F 1 z H o x W M U 8 4 B k C b O T 2 W T I 7 O w y 0 y u E J X / g x Y M i X v 0 j b / 6 N k 2 Q P m l j Q U F R 1 0 9 0 V J F I Y d N 1 v Z 2 l 5 Z X V t v b B R 3 N z a 3 t k t 7 e 0 3 T J x q x u s s l r F u B d R w K R S v o 0 D J W 4 n m N A o k b w b D 2 4 n f f O L a i F g 9 4 i j h f k T 7 S o S C U b T S w 2 n Y L Z X d i j s F W S R e T s q Q o 9 Y t f X V 6 M U s j r p B J a k z b c x P 0 M 6 p R M M n H x U 5 q e E L Z k P Z 5 2 1 J F I 2 7 8 b H r p m B x b p U f C W N t S S K b q 7 4 m M R s a M o s B 2 R h Q H Z t 6 b i P 9 5 7 R T D a z 8 T K k m R K z Z b F K a S Y E w m b 5 O e 0 J y h H F l C m R b 2 V s I G V F O G N p y i D c G b f 3 m R N M 4 q 3 m X l 4 v 6 8 X L 3 J 4 y j A I R z B C X h w B V W 4 g x r U g U E I z / A K b 8 7 Q e X H e n Y 9 Z 6 5 K T z x z A H z i f P z e F j S s = < / l a t e x i t > f < l a t e x i t s h a 1 _ b a s e 6 4 = " v 7 3 / 7 i D x 2 o y t d I p U l f f 0 X R 7 x X / Y = " > A A A B 6 n i c b V D J S g N B E K 2 J W 4 x b 1 K O X x i B 4 C j P i d g x 6 8 R j R L J A M o a f T k z T p Z e j u E c K Q T / D i Q R G v f p E 3 / 8 Z O M g d N f F D w e K + K q n p R w p m x v v / t F V Z W 1 9 Y 3 i p u l r e 2 d 3 b 3 y / k H T q F Q T 2 i C K K 9 2 O s K G c S d q w z H L a T j T F I u K 0 F Y 1 u p 3 7 r i W r D l H y 0 4 4 S G A g 8 k i x n B 1 k k P X Z H 2 y h W / 6 s + A l k m Q k w r k q P f K X 9 2 + I q m g 0 h K O j e k E f m L D D G v L C K e T U j c 1 N M F k h A e 0 4 6 j E g p o w m 5 0 6 Q S d O 6 a N Y a V f S o p n 6 e y L D w p i x i F y n w H Z o F r 2 p + J / X S W 1 8 H W Z M J q m l k s w X x S l H V q H p 3 6 j P N C W W j x 3 B R D N 3 K y J D r D G x L p 2 S C y F Y f H m Z N M + q w W X 1 4 v 6 8 U r v J 4 y j C E R z D K Q R w B T W 4 g z o 0 g M A A n u E V 3 j z u v X j vq x H G i E g t C v I H D y F q 5 P V w g s O X 0 7 4 = " > A A A B 7 X i c b V D J S g N B E K 2 J W 4 x b 1 K O X x i B 4 C j P i d g x 6 8 R j B L J A M o a f T k 7 T p Z e j u E c K Q f / D i Q R G v / o 8 3 / 8 Z O M g d N f F D w e K + K q n p R w p m x v v / t F V Z W 1 9 Y 3 i p u l r e 2 d 3 b 3 y / k H T q F Q T 2 i C K K 9 2 O s K G c S d q w z H L a T j T F I u K 0 F Y 1 u p 3 7 r i W r D l H y w 4 4 S G A g 8 k i x n B 1 k n N 7 g A L g X v l i l / 1 Z 0 D L J M h J B X L U e + W v b l + R V F B p C c f G d A I / s W G G t W W E 0 0 m p m x q a Y D L C A 9 p x V G J B T Z j N r p 2 g E 6 f 0 U a y 0 K 2 n R T P 0 9 k W F h z F h E r l N g O z S L 3 l T 8 z + u k N r 4 O M y a T 1 F J J 5 o v i l C O r 0 P R 1 1 G e a E s v H j m C i m b s V k S H W m F g X U M m F E C y + v E y a Z 9 X g s n p x f 1 6 p 3 e R x F O E I j u E U A r i C G t x B H R p A 4 B G e 4 R9 J m I Q x y 4 m S 4 = " > A A A C J X i c b V D L S g M x F M 3 4 r P V V d e k m W A R F L D P i C 1 Q Q 3 b h U s C q 0 Z b i T Z t r Q Z G Z M 7 k j L M D / j x l 9 x 4 0 I R w Z W / Y l q 7 8 H U g c H L O u S T 3 B I k U B l 3 3 3 R k Z H R u f m C x M F a d n Z u f m S w u L V y Z O N e N V F s t Y 3 w R g u B Q R r 6 J A y W 8 S z U E F k l 8 H n d O + f 3 3 H t R F x d I m 9 h D c U t C I R C g Z o J b 9 0 m P h I j + h m u N b 1 c Z 1 u 0 L p K a e J n u O n l B / X b F J q 0 a 2 8 b X m 5 T N t J P t E A p s C H 0 S 2 W 3 4 g 5 A / x J v S M p k i H O / 9 F J v x i x V P E I m w Z i a 5 y b Y y E C j Y J L n x X p q e A K s A y 1 e s z Q C x U 0 j G 2 y Z 0 1 W r N G k Y a 3 s i p A P 1 + 0 Q G y p i e C m x S A b b N b 6 8 v / u f V U g z 3 G 5 m I k h R 5 x L 4 e C l N J M a b 9 y m h T a M 5 Q 9 i w B p o X 9 K 2 V t 0 M D Q F l u 0 J X i / V / 5 L r r Y q 3 m 5 l 5 2 K 7 f H w y r K N A l s k K W S MZ x q x h s s l r F u B 9 R w K R R v o E D J 2 4 n m N A o k b w W j 2 6 n f e u L a i F g 9 4 j j h f k Q H S o S C U b T S Q 9 L D X r n i V t 0 Z y D L x c l K B H P V e + a v b j 1 k a c Y V M U m M 6 n p u g n 1 G N g k k + K X V T w x P K R n T A O 5 Y q G n H j Z 7 N T J + T E K n 0 S x t q W Q j J T f 0 9 k N D J m H A W 2 M 6 I 4 N I v e V P z P 6 6 Q Y X v u Z U E m K X L H 5 o j C V B G M y / Z v 0 h e Y MEQUATIONwhere D is the feasible region and c i ∈ 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 is 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 a solution to be Pareto-optimal is for it to be Pareto-stationary, which is defined as:Definition 1 (Pareto-stationary) Let x be in the interior of the feasible region, D, in which the E objective functions, F i , are smooth, and f i (x) = ∇ x F i (x) be the local gradients for i = 1, . . . , E.x is said to be Pareto-stationary if there exists a vector α= [α 1 , . . . , α E ] ⊤ ∈ R E such that α i ≥ 0, E i=1 α i = 1 and E i=1 α i f i (x) = 0.That is, there exists a convex combination of the gradient-vectors f i (x) that is equal to 0.Therefore, it would be intuitive to extend the steepest descent algorithm to a multi-objective setting by finding a descent direction, that is common to all objectives, in the convex hull of the normalized local gradients fi (x) = f i (x)/∥f i (x)∥. We denote such a set as [12] developed the Multiple-Gradient Descent Algorithm (MGDA) from such an understanding, proving that there does exist such a descent direction in Ū , which is the direction with the smallest norm in the set. Then, the update rule of MGDA isŪ = {v ∈ R N |v = E i=1 α i fi (x); α i ≥ 0, ∀i; E i=1 α i = 1}. Indeed,EQUATIONwhereα * = (α * 1 , . . . , α * E ) minimizes {∥v∥|v ∈ Ū }. 
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 an objective function F (x) whose minima we aim to find by modifying our iterate x t at each time step t through its gradient f (x t ) = ∇ x F (x t ), scaled by a step size γ > 0:EQUATIONHowever, 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 ofEQUATIONValidation Training < l a t e x i t s h a 1 _ b a s e 6 4 = " x 0 y R R 4 h d w 1 q T 0 5 where γ > 0 is the step size, and µ ≥ 0 is the momentum constant. Eqn. 6 can then be rewritten as:1 Z j / v Z E 5 A S J W M = " > A A A B 7 3 i c b V D L T s J A F L 3 F F + I L d e l m I j H B D W k N P p Z E N y 4 x k U d C G z I d p j B h O q 0 z U y M p / I Q b F x r j 1 t 9 x 5 9 8 4 Q B c K n u Q m J + f c m 3 v v 8 W P O l L b t b y u 3 s r q 2 v p H f L G x t 7 + z u F f c P m i p K J K E N E v F I t n 2 s K G e C N j T T n L Z j S X H o c 9 r y h z d T v / V I p W K R u N e j m H o h 7 g s W M I K 1 k d r u O C g / n b r j b r F k V + w Z 0 D J x M l K C D P V u 8 c v t R S Q J q d C E Y 6 U 6 j h 1 r L 8 V S M 8 L p p O A m i s a Y D H G f d g w V O K T K S 2 f 3 T t C J U X o o i K Q p o d F M / T 2 R 4 l C p U e i b z h D r g V r 0 p u J / X i f R w Z W X M h E n m g o y X x Q k H O k I T Z 9 H P S Y p 0 X x k C C a S m V s R G W C J i T Y R F U w I z u L L y 6 R 5 V n E u K u d 3 1 V L t O o s j D 0 d w D G V w 4 B J q c A t 1 a A A B D s / w C m / W g / V i v V s f 8 9 a c l c 0 c w h 9 Y n z + a Z o + z < / l a ty R R 4 h d w 1 q T 0 5 1 Z j / v Z E 5 A S J W M = " > A A A B 7 3 i c b V D L T s J A F L 3 F F + I L d e l m I j H B D W k N P p Z E N y 4 x k U d C G z I d p j B h O q 0 z U y M p / I Q b F x r j 1 t 9 x 5 9 8 4 Q B c K n u Q m J + f c m 3 v v 8 W P O l L b t b y u 3 s r q 2 v p H f L G x t 7 + z u F f c P m i p K J K E N E v F I t n 2 s K G e C N j T T n L Z j S X H o c 9 r y h z d T v / V I p W K R u N e j m H o h 7 g s W M I K 1 k d r u O C g / n b r j b r F k V + w Z 0 D J x M l K C D P V u 8 c v t R S Q J q d C E Y 6 U 6 j h 1 r L 8 V S M 8 L p p O A m i s a Y D H G f d g w V O K T K S 2 f 3 T t C J U X o o i K Q p o d F M / T 2 R 4 l C p U e i b z h D r g V r 0 p u J / X i f R w Z W X M h E n m g o y X x Q k H O k I T Z 9 H P S Y p 0 X x k C C a S m V s R G W C J i T Y R F U w I z u L L y 6 R 5 V n E u K u d 3 1 V L t O o s j D 0 d w D G V w 4 B J q c A t 1 a A A B D s / w C m / W g / V i v V s f 8 9 a c l c 0 c w h 9 Y n z + a Z o + z < / l a t e x i t > kf (x)k Layer SMoEEQUATIONBy incorporating the past gradients in each update, the descent path can become smoother with fewer oscillations, resulting in a faster convergence [53, 21, 71] .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, we are going to establish a connection between MoE and GD, and further leverage momentum to enhance MoE and SMoE. We rewrite Eqn. 2 of MoE as follows:EQUATIONIf we regard -u i (x t ) as the local "gradient" ∇ x F i (x t ) at the t-th iteration and softmax(g i (x t )) as to be learned to approximate α * i , then we can consider the MoE in Eqn. 2 and 8 as the dynamical system which updates x t using the MGDA to minimize the multi-objective optimization in Eqn. 3.Descent. Given the analogy between MoE and GD, SMoE can be interpreted as a stochastic version of MoE, which corresponds to an SGD algorithm applied to the multi-objective optimization problem in Eqn. 3. SMoE is then reformulated as:EQUATIONHere, -f (x t ) = γ K i=1 softmax(TopK(g i (x t )))u i (x t ) is the SMoE output. Empirical Evidences for the Gradient Descent Analogy of (S)MoE. We provide empirical justification for the connection between (S)MoE and (S)GD in Fig. 2 and 3 .Gradient norm ∥f (x t )∥ 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 ∥f (x t )∥ = ∇ x F (x t ), respectively. It is 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-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.Gradient norm ∥f (x t )∥ at each layer t decreases during training: According to the update rule of MGDA in Eqn. 7, the coefficient α * i minimizes the norm ∥ E i=1 α i fi ∥. In Eqn. 8 and 9, as discussed above, softmax(g i (x t )) and softmax(TopK(g i (x t ))) try to learn α * i , respectively. Thus, it is expected that these two terms learn to reduce the corresponding ∥ E i=1 α i fi ∥ in Eqn. 8 and 9, which is the norm of the SMoE output at layer t. Fig. 2 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 i (x t )) and softmax(TopK(g i (x t ))) learn to approximate α * i . We provide the full plots for all layers in Appendix C, Fig. 6 and 7 .We propose the new MomentumSMoE layer, depicted in Fig. 1 , to accelerate the dynamics of 8, which is principled by the accelerated gradient descent theory (see Section 2.2):EQUATIONwhere µ ≥ 0 and γ > 0 are hyperparameters corresponding to the momentum coefficient and step size in the momentum-accelerated GD, respectively. The formulation of MomentumSMoE can be applied to MoE to derive the MomentumMoE. 
stability analysis: momentumsmoe vs. smoe: 
 In this section, we demonstrate the theoretical advantages of MomentumSMoE over SMoE. In particular, we show that the spectrum of MomentumSMoE is better-structured than that of SMoE, thus MomentumSMoE is more stable than SMoE. We rewrite MomentumSMoE using the equivalent form of momentum acceleration given in Eqn. 7 as follows:EQUATIONTaking inspiration from [57] , we then expand f (x t ) around the Pareto-stationary solution x * at which f (x * ) = 0 (see Definition 1) using Taylor expansion to obtain an approximation of f (x t ):EQUATIONSubstituting the Taylor expansion of f (x t ) in Eqn. 12 into Eqn. 11, we attainx t+1 = x t -γ∇ x f (x * )(x t -x * ) + µ(x t -x t-1 ).Without loss of generality, we further let x * = 0 as we can always replace x t with x t + x * . The formula of MomentumSMoE can then be simplified asEQUATIONSuppose that ∇ x f (x * ) does not have any defective eigenvalues and hence is diagonalizable. Then,∇ x f (x * ) = QΣQ -1, for some invertible matrix Q and the diagonal matrix Σ with diagonal entries being the eigenvalues σ(n), n = 1, 2, . . . , N , of ∇ x f (x * ). We can then rewrite Eqn. 13 asEQUATION)Since we have decoupled the N features in x t , we can consider each feature x t (n), separately. Introducing a dummy equation x t = x t , we rewrite Eqn. 14 as follows:EQUATIONThe convergence of x t (n) then depends on the eigenvalues λ 1 (A) and λ 2 (A) of A. In particular, we require max{|λ1 (A)|, |λ 2 (A)|} < 1.It should be noted that omitting the momentum parameter, i.e., µ = 0, recovers the standard, unaccelerated SMoE layer.A = 0 1 -µ (1 + µ) -γσ(n) and λ 1 (A), λ 2 (A) are eigenvalues of A, max{|λ 1 (A)|, |λ 2 (A)|} < 1 if and only if µ ∈ (-1, 1) and γσ(n) ∈ (0, 2 + 2µ).Proposition 1 (Convergence of MomentumSMoE) The MomentumSMoE defined in Eqn. 10 converges if and only if µ ∈ (-1, 1) and γσ(n) ∈ (0, 2 + 2µ).The proofs of the results above are provided in Appendix A. It is worth noting that in both the MomentumSMoE and standard SMoE, the convergence of x t depends on the eigenvalues of the Jacobian ∇ x f of the SMoE layer. Since the step size γ > 0, we require that ∇ x f to be positive definite for its eigenvalues σ(n) to be positive. Furthermore, even though among the convergence conditions of MomentumSMoE is that µ ∈ (-1, 1), in practice, µ is chosen to be positive.Proposition 1 implies that the spectrum of MomentumSMoE is better-structured than that of SMoE. Thus, MomentumSMoE is more stable than SMoE. We summarize this finding in Corollary 1 below.Corollary 1 (MomentumSMoE is more stable than SMoE) Without momentum, µ = 0, the range of values that γσ(n) can take for the system to be stable is limited to 0 < γσ(n) < 2. The addition of momentum expands this margin, almost doubling it, providing a larger parameter range for the network to converge stably to a good output in the forward pass. 
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, we propose two additional variants of MomentumSMoE, AdamSMoE and Robust MomentumSMoE, which are derived from the AdamW [33, 42] and Robust Momentum [10] , respectively.Adam [33] accelerates the gradient dynamics by utilizing the moving average of historical gradients and element-wise squared gradients. Adam with a decoupled weight decay regularization (AdamW) is more commonly used in practice thanks to its better generalization over Adam. We employ AdamW to derive the AdamSMoE as follows:p t = µp t-1 + (1 -µ)[-f (x t )]; m t = βm t-1 + (1 -β)f (x t ) ⊙ f (x t ) x t+1 = x t + γ √ m t + ϵ p t -κx twhere ϵ is a small constant to prevent numerical instability, κ the weight decay parameter, γ the step size, and µ and β are the decay parameters for the moment estimates.Deep learning models, including SMoE, are known to not be robust to 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 MomentumSMoE from the Robust Momentum Method [10] .The Robust Momentum Method proposed by [10] has the following update ruleEQUATIONwhere γ, µ and α are parameterized by an additional hyperparameter p as follows:EQUATIONHere, k = L/m is a condition ratio of the objective function assuming that it is m-strongly convex and L-smooth. Compared with the heavy-ball momentum in Eqn. 7, there is an additional variable y t that can be interpreted as a feedback signal to steer the x t toward a robust solution. The parameters γ, µ, and α are designed such that the new system is robust. 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:EQUATIONwhere γ, µ and α are as defined in Eqn. 17, and -f(y t ) = γ K i=1 softmax(TopK(g i (y t )))u i (y t )is the SMoE output given the input y t . Equivalently, at each layer t, we update the input x t and momentum vector p t asy t = x t + αγp t-1 ; p t = -f (y t ) + µp t-1 ; x t+1 = x t + γp t .Remark 1 We provide an interpretation of robust momentum in Appendix B. 
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-1k object recognition tasks. We aim to show that: (i) MomentumSMoE improves model performance across both language and vision tasks; (ii) AdamSMoE significantly outperforms the baseline and accelerates convergence in language models, even surpassing MomentumSMoE; (iii) Robust MomentumSMoE is highly effective at improving robustness of vision models to data corruption; (iv) MomentumSMoE is universal and can be easily integrated into many state-of-the-art SMoE and MoE models.Throughout the experiments, we compare our 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.1. We find that implementing AdamSMoE in the first layer is enough to significantly improve the model'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 8 A100 GPUs.We use the Switch Transformer [18] , referred to as SMoE in our tables and figures below, and GLaM [16] baselines. We consider 2 configurations: medium (6 layers) and large (12 layers). We report the perplexity (PPL) of MomentumSMoE and AdamSMoE in comparison with the baseline SMoE on word-level WikiText-103 validation and test datasets for both model sizes in Table 1 . We also include experiments on the medium-sized 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 WikiText-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 An important advantage of MomentumSMoE is its simplicity, which allows easy implementation with negligible computational overhead. We provide a comparison of the run time/sample, memory, number of parameters, and computation time between models in Table 11 and 12 in Appendix E.5.EQUATIONIn 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-1K (IN-1K) object recognition task. We focus on i) improving the robustness of V-MoE using Robust MomentumSMoE (18) and ii) demonstrating that our momentum method is not limited to sparse models but can be generalized to MoE models such as Soft MoE. To benchmark robustness to data corruptions, we 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 a SMoE layer. In training. These findings justify the benefits and universality of our momentum-based method, that extends beyond SMoE to MoE.x H G i E g t C v I H D y F q 5 P V w g s O X 0 7 4 = " > A A A B 7 X i c b V D J S g N B E K 2 J W 4 x b 1 K O X x i B 4 C j P i d g x 6 8 R j B L J A M o a f T k 7 T p Z e j u E c K Q f / D i Q R G v / o 8 3 / 8 Z O M g d N f F D w e K + K q n p R w p m x v v / t F V Z W 1 9 Y 3 i p u l r e 2 d 3 b 3 y / k H T q F Q T 2 i C K K 9 2 O s K G c S d q w z H L a T j T F I u K 0 F Y 1 u p 3 7 r i W r D l H y w 4 4 S G A g 8 k i x n B 1 k n N 7 g A L g X v l i l / 1 Z 0 D L J M h J B X L U e + W v b l + R V F B p C c f G d A I / s W G G t W W E 0 0 m p m x q a Y D L C A 9 p x V G J B T Z j N r p 2 g E 6 f 0 U a y 0 K 2 n R T P 0 9 k W F h z F h E r l N g O z S L 3 l T 8 z + u k N r 4 O M y a T 1 F J J 5 o v i l C O r 0 P R 1 1 G e a E s v H j m C i m b s V k S H W m F g X U M m F E C y + v E y a Z 9 X g s n p x f 1 6 p 3 e R x F O E I j u E U A r i C G t x B H R p A 4 B G e 4 RI p U l f f 0 X R 7 x X / Y = " > A A A B 6 n i c b V D J S g N B E K 2 J W 4 x b 1 K O X x i B 4 C j P i d g x 6 8 R j R L J A M o a f T k z T p Z e j u E c K Q T / D i Q R G v f p E 3 / 8 Z O M g d N f F D w e K + K q n p R w p m x v v / t F V Z W 1 9 Y 3 i p u l r e 2 d 3 b 3 y / k H T q F Q T 2 i C K K 9 2 O s K G c S d q w z H L a T j T F I u K 0 F Y 1 u p 3 7 r i W r D l H y 0 4 4 S G A g 8 k i x n B 1 k k P X Z H 2 y h W / 6 s + A l k m Q k w r k q P f K X 9 2 + I q m g 0 h K O j e k E f m L D D G v L C K e T U j c 1 N M F k h A e 0 4 6 j E g p o w m 5 0 6 Q S d O 6 a N Y a V f S o p n 6 e y L D w p i x i F y n w H Z o F r 2 p + J / X S W 1 8 H W Z M J q m l k s w X x S l H V q H p 3 6 j P N C W W j x 3 B R D N 3 K y J D r D G x L p 2 S C y F Y f H m Z N M + q w W X 1 4 v 6 8 U r v J 4 y j C E R z D K Q R w B T W 4 g z o 0 g M A A n u E V 3 j z u v X j v 3 s e 8 t e D l M 4 f w B 9 7 n D 2 C e j e A = < / l a t e x i t > µ < l a t e x i t s h a 1 _ b a s e 6 4 = " v D q c K H 9 o p 4 x F i M H H r j h m y l f X k G 8 = " > A A A B 8 n i c b V D L S g M x F M 3 U V 6 2 v q k s 3 w S K 4 K j P i a y M U 3 b i s Y B 8 w H U o m T d v Q P I b k j l C G f o Y b F 4 q 4 9 W v c + T e m 7 S y 0 9 U D I 4 Z x 7 u f e e O B H c g u 9 / e 4 W V 1 b X 1 j e J m a W t 7 Z 3 e v v H / Q t D o 1 l D W o F t q 0 Y 2 K Z 4 I o 1 g I N g 7 c Q w I m P B W v H o b u q 3 n p i x X K t H G C c s k m S g e J 9 T A k 4 K O w M i J b n x q 3 7 Q L V f c N w N e J k F O K i h H v V v + 6 v Q 0 T S V T Q A W x N g z 8 B K K M G O B U s E m p k 1 q W E D o i A x Y 6 q o h k N s p m K 0 / w i V N 6 u K + N e w r w T P 3 d k R F p 7 V j G r l I S G N p F b y r + 5 4 U p 9 K + j j K s k B a b o f F A / F R g 0 n t 6 P e 9 w w C m L s C K G G u 1 0 x H R J D K L i U S i 6 E Y P H k Z d I 8 q w a X 1 Y u H 8 0 r t N o + j i I 7 Q M T p F A b p C N X S P 6 q i B K N L o G b 2 i N w + 8 F + / d + 5 i X F r y 8 5 x D 9 g f f 5 A 8 7 e k E 0 = < / l a t e x i t > = 0.01 < l a t e x i t s h a 1 _ b a s e 6 4 = " P Z k R + c d A T k o p A m 6 G e s a I a R O J + s k = " > A A A B 8 X i c b V B N S w M x E M 3 W r 1 q / q h 6 9 B I v g a d k V q 1 6 E o h e P F e w H t k u Z T b N t a J J d k q x Q l v 4 L L x 4 U 8 e q / 8 e a / M W 3 3 o K 0 P B h 7 v z T A z L 0 w 4 0 8 b z v p 3 C y u r a + k Z x s 7 S 1 v b O 7 V 9 4 / a O o 4 V Y Q 2 S M x j 1 Q 5 B U 8 4 k b R h m O G 0 n i o I I O W 2 F o 9 u p 3 3 q i S r N Y P p h x Q g M B A 8 k i R s B Y 6 b E 7 A C H g 2 n O r v X L F c 7 0 Z 8 D L x c 1 J B O e q 9 8 l e 3 H 5 N U U G k I B 6 0 7 v p e Y I A N l G O F 0 U u q m m i Z A R j C g H U s l C K q D b H b x B J 9 Y p Y + j W N m S B s / U 3 x M Z C K 3 H I r S d A s x Q L 3 p T 8 T + v k 5 r o K s i Y T F J D J Z k v i l K O T Y y n 7 + M + U 5 Q Y P r Y E i G L 2 V k y G o I A Y G 1 L J h u A v v r x M m m e u f + F W 7 8 8 r t Z s 8 j i I 6 Q s f o F P n o E t X Q H a q j B i J I o m f 0 i t 4 c 7 b w 4 7 8 7 H v L X g 5 D O H 6 A + c z x 9 k D p A X < / l a t e x i t > = 0.5 < l a t e x i t s h a 1 _ b a s e 6 4 = " 9 I o X g M R M f h v n j m G W 1 K K Q 2 U L x D g I = " > A A A B 8 X i c b V D L S g N B E O y N r x h f U Y 9 e F o P g a d k V X x c h 6 M V j B P P A Z A m 9 k 0 k y Z G Z 2 m Z k V w p K / 8 O J B E a / + j T f / x k m y B 4 0 W N B R V 3 X R 3 R Q l n 2 v j + l 1 N Y W l 5 Z X S u u l z Y 2 t 7 Z 3 y r t 7 D R 2 n i t A 6 i X m s W h F q y p m k d c M M p 6 1 E U R Q R p 8 1 o d D P 1 m 4 9 U a R b L e z N O a C h w I F m f E T R W e u g M U A i 8 C j y / W 6 7 4 n j + D + 5 c E O a l A j l q 3 / N n p x S Q V V B r C U e t 2 4 C c m z F A Z R j i d l D q p p g m S E Q 5 o 2 1 K J g u o w m 1 0 8 c Y + s 0 n P 7 s b I l j T t T f 0 5 k K L Q e i 8 h 2 C j R D v e h N x f + 8 d m r 6 l 2 H G Z J I a K s l 8 U T / l r o n d 6 f t u j y l K D B 9 b g k Q x e 6 t L h q i Q G B t S y Y Y Q L L 7 8 l z R O v O D c O 7 s 7 r V S v 8 z i K c A C H c A w B X E A V b q E G d S A g 4 Q l e 4 N Xi p U U l l P g y z a H W A N h g D V Y Y 2 U = " > A A A B 8 X i c b V D L S g N B E O y N r x h f U Y 9 e F o P g a d k V X x c h 6 M V j B P P A Z A m 9 k 0 k y Z G Z 2 m Z k V w p K / 8 O J B E a / + j T f / x k m y B 4 0 W N B R V 3 X R 3 R Q l n 2 v j + l 1 N Y W l 5 Z X S u u l z Y 2 t 7 Z 3 y r t 7 D R 2 n i t A 6 i X m s W h F q y p m k d c M M p 6 1 E U R Q R p 8 1 o d D P 1 m 4 9 U a R b L e z N O a C h w I F m f E T R W e u g M U A i 8 8 r 2 g W 6 7 4 n j + D + 5 c E O a l A j l q 3 / N n p x S Q V V B r C U e t 2 4 C c m z F A Z R j i d l D q p p g m S E Q 5 o 2 1 K J g u o w m 1 0 8 c Y + s 0 n P 7 s b I l j T t T f 0 5 k K L Q e i 8 h 2 C j R D v e h N x f + 8 d m r 6 l 2 H G Z J I a K s l 8 U T / l r o n d 6 f t u j y l K D B 9 b g k Q x e 6 t L h q i Q G B t S y Y Y Q L L 7 8 l z R O v O D c O 7 s 7 r V S v 8 z i K c A C H c A w B X E A V b q E G d S A g 4 Q l e 4 N X R z r P z 5 r z P W w t O P r M P v + B 8 f A N d / p A T < / l a t e x i t > = 0.1 < l a t e x i t s h a 1 _ b a s e 6 4 = " 9 e Z G f f k 1 C w J m J G g B R G + B d R o s 8 B A = " > A A A B 7 n i c b V D L S g N B E O y N r x h f U Y 9 e B o P g a d k V 4 + M g B L 1 4 j G A e k C x h d j K b D J m Z X W Z m h b D k I 7 x 4 U M S r 3 + P N v 3 G S 7 E E T C x q K q m 6 6 u 8 K E M 2 0 8 7 9 s p r K y u r W 8 U N 0 t b 2 z u 7 e + X 9 g 6 a O U 0 V o g 8 Q 8 V u 0 Q a 8 q Z p A 3 D D K f t R F E s Q k 5 b 4 e h u 6 r e e q N I s l o 9 m n N B A 4 I F k E S P Y W K n V F e m N 5 1 7 3 y h X P 9 W Z A y 8 T P S Q V y 1 H v l r 2 4 / J q m g 0 h C O t e 7 4 X m K C D C v D C K e T U j f V N M F k h A e 0 Y 6 n E g u o g m 5 0 7 Q S d W 6 a M o V r a k Q T P 1 9 0 S G h d Z j E d p O g c 1 Q L 3 p T 8 T + v k 5 r o K s i Y T F J D J Z k v i l K O T I y m v 6 M + U 5 Q Y P r Y E E 8 X s r Y g M s c L E 2 I R K N g R / 8 e V l 0 j x z / Q u 3 + n B e q d 3 m c R T h C I 7 h F H y 4 h B r c Q x 0 a Q G A E z / A K b 0 7 i v D j v z s e 8 t e D k M 4 f w B 8 7 n D z v 1 j t w = < / l a t e x i t > µ = 0.9 < l a t e x i t s h a 1 _ b a s e 6 4 = " / f D n U t F o 0 F o U 2 g P 7 M T b 8 A F v L G 0 g = " > A A A B 7 n i c b V B N S w M x E J 2 t X 7 V + V T 1 6 C R b B 0 7 I r 1 n o R i l 4 8 V r A f 0 C 4 l m 2 b b 0 C S 7 J F m h L P 0 R X j w o 4 t X f 4 8 1 / Y 9 r u Q V s f D D z e m 2 F m X p h w p o 3 n f T u F t f W N z a 3 i d m l n d 2 / / o H x 4 1 N J x q g h t k p j H q h N i T T m T t G m Y 4 b S T K I p F y G k 7 H N / N / P Y T V Z r F 8 t F M E h o I P J Q s Y g Q b K 7 V 7 I r 3 x 3 F q / X P F c b w 6 0 S v y c V C B H o 1 / + 6 g 1 i k g o q D e F Y 6 6 7 v J S b I s D K M c D o t 9 V J N E 0 z G e E i 7 l k o s q A 6 y + b l T d G a V A Y p i Z U s a N F d / T 2 R Y a D 0 R o e 0 U 2 I z 0 s j c T / / O 6 q Y m u g 4 z J J D V U k s W i K O X I x G j 2 O x o w R Y n h E 0 s w U c z e i s g I K 0 y M T a h k Q / C X X 1 4 l r Q v X v 3 Kh E 9 E e U = " > A A A B 7 n i c b V B N S w M x E J 2 t X 7 V + V T 1 6 C R b B 0 7 I r t n o R i l 4 8 V r A f 0 C 4 l m 2 b b 0 C S 7 J F m h L P 0 R X j w o 4 t X f 4 8 1 / Y 9 r u Q V s f D D z e m 2 F m X p h w p o 3 n f T u F t f W N z a 3 i d m l n d 2 / / o H x 4 1 N J x q g h t k p j H q h N i T T m T t G m Y 4 b S T K I p F y G k 7 H N / N / P Y T V Z r F 8 t F M E h o I P J Q s Y g Q b K 7 V 7 I r 3 x 3 G q / X P F c b w 6 0 S v y c V C B H o 1 / + 6 g 1 i k g o q D e F Y 6 6 7 v J S b I s D K M c D o t 9 V J N E 0 z G e E i 7 l k o s q A 6 y + b l T d G a V A Y p i Z U s a N F d / T 2 R Y a D 0 R o e 0 U 2 I z 0 s j c T / / O 6 q Y m u g 4 z J J D V U k s W i K O X I x G j 2 O x o w R Y n h E 0 s w U c z e i s g I K 0 y M T a h k Q / C X X 1 4 l r Q v X r 7 n V h 8 t K / T a P o w g n c A r n 4 M M V 1 O E e G t A E A m Nv u r j D h T B v P + 3 Y K K 6 t r 6 x v F z d L W 9 s 7 u X n n / o K n j V B H a I D G P V T v E m n I m a c M w w 2 k 7 U R S L k N N W O L q b + q 0 n q j S L 5 a M Z J z Q Q e C B Z x A g 2 V m p 1 R X r j u e e 9 c s V z v R n Q M v F z U o E c 9 V 7 5 q 9 u P S S q o N I R j r T u + l 5 g g w 8 o w w u m k 1 E 0 1 T T A Z 4 Q H t W C q x o D r I Z u d O 0 I l V + i i K l S 1 p 0 E z 9 P Z F h o f V Y h L ZQuestion: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer: [Yes]Justification: We provide the source code so that the results in the paper can be easily reproduced. 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. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
empirical analysis: 
 We conduct empirical analysis based on the SMoE-medium trained on WikiText-103 in Section 5.1.Norm-based Load Imbalance. Load imbalance in SMoE occurs when only a small subset of experts are consistently selected [63, 77] . Numerous methods have been developed to counter this common phenomenon, such as introducing a buffer capacity for each expert and a load balancing loss [65, 18] . Orthogonal to these, in line with our GD framework for SMoEs, we examine the choice of experts determined by the size of the norm of their outputs, ∥f i (x)∥.From a multi-objective optimization perspective, the optimal descent direction is one that minimizes the norm in the convex hull of the normalized gradients Ū (see Section 2.1). If 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 are 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 observe the corresponding phenomenon in the SMoE, further empirically justifying our connection between SMoE and GD. The full plots for all layers and all models are provided in Appendix C.The direction with the smallest norms are frequently related to the objectives that have already had a substantial degree of convergence and is insufficient for a balanced minimization of the multi-objective criteria. In this light, an ideal SMoE output would have a norm-based balanced choice of experts and should translate to improved model performance. Indeed, in Section 5.1, we established the superior performance of MomentumSMoE and AdamSMoE on large-scale WikiText-103 language modeling task, and in Figure 5 (Left, B, C), this directly correlates with a significantly more balanced selection of experts with respect to their norms.Step Size γ. To better understand the effects of the momentum parameter and step size on the performance of the trained MomentumSMoE models, we do an ablation study on these two hyperparameters and include results in Fig. 5 (Right), which contains a plot of log validation PPL during 40 training epochs. We notice that MomentumSMoE is robust to the choice of µ, and we select µ = 0.7 for the final comparison with the baseline SMoE. On the other hand, when the value of γ is too small, there is an adverse effect on the model. Hence, we select γ = 1.0.Step Size γ Learnable. We note that additional time and effort are required to 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 γ and µ learnable parameters, in this section and include results in Table 5 in Appendix E.1. We leave the discussion on other methods to Appendix E.1. As shown in Table 5 in Appendix E.1, MomentumSMoE (ii), with learnable γ and fixed µ, 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 MomentumSMoE are further enhanced with more than 2 PPL improvements. These results confirm that the benefits of our model can be leveraged with minimal effort. Since the MomentumSMoE is robust to the choice of µ, we do not consider the setting of fixing γ and making µ learnable.Other Optimizers. We study the integration of other advanced momentum and optimization methods, such as the Nesterov accelerated gradient [48] , RMSProp [68] , and sharpness-aware minimization [64] , into our MomentumSMoE framework in Appendix E.4.Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, 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 to 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 of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
related work: 
 Sparse Mixture of Experts. SMoE has been extensively studied to 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 to 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-based SMoE can be easily incorporated into these methods above to 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 to 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 momentum-accelerated first-order optimization algorithms.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 error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.• 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. • If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration 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.• 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), fairness 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 not tied to particular applications, let alone deployments. However, if there is a 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 as 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 or 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 a system learns from feedback over time, improving the efficiency and accessibility of ML). 
concluding remarks: 
 In this paper, we propose MomentumSMoE, 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 to the SMoE baseline. Furthermore, we demonstrate that our momentum-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-103 and ImageNet-1K. As shown in Table 7 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, MomentumSMoE has little effect. From a theoretical perspective, it would be intriguing to develop a theory to explain the enhanced robustness of MomentumSMoE. Furthermore, MomentumSMoE can be analyzed as a fixed-point iteration. We leave these theoretical developments as future work.Supplement to "MomentumSMoE: Integrating Momentum into Sparse Mixture of Experts" We can find the eigenvalues of matrix A explicitly asEQUATIONIf γσ(n) ≤ 0,EQUATION≥ (1 + µ) + (1 + µ) 2 -4µ 2 = (1 + µ) + |1 -µ| 2 = max{1, µ} ≥ 1Then, we must have µσ(n) > 0. Letting the expression in the square root be ∆, expanding it,∆ := (1 + µ -γσ(n)) 2 -4µ = (1 -γσ(n)) 2 + 2(1 -γσ(n))µ + µ 2 -4µ = µ 2 -2(1 -γσ(n))µ + (1 -γσ(n)) 2and finding the roots of the equation yields, µ = (1 ± γσ(n)) 2 . Then, we consider two cases.Case 1: ∆ ≥ 0 when µ ≥ (1 + γσ(n)) 2 or µ ≤ (1 -γσ(n)) 2 1a: If µ ≥ (1+ γσ(n)) 2 , then 1+µ-γσ(n) ≥ 1+(1+ γσ(n)) 2 -γσ(n) ≥ 2+2 γσ(n) ≥ 0.Then, λ 1,2 > 0 and λ 1 ≥ λ 2 . For the system to converge, we needEQUATIONHowever, by assumption, we have a contradiction0 ≤ √ ∆ < 1 -µ + γσ(n) ≤ 1 -(1 + γσ(n)) 2 + γσ(n) = -2 γσ(n) < 0.Hence, we cannot have µ ≥ (1 + γσ(n)) 2 .1b: 2 which implies that γσ(n) ≥ 2 γσ(n) and hence, γσ(n) ≥ 4. Combining this with condition i), we have √ µ ≤ γσ(n) -1 which leads to the following contradiction with ii)If µ ≤ (1 -γσ(n))≤ √ ∆ < 1-µ+γσ(n). As 1-µ+γσ(n) ≥ 1-(1-γσ(n)) 2 +γσ(n) = 2 γσ(n) > 0, it is enough to check that ∆ < (1 -µ + γσ(n)) 2 can be satisfied. ∆ < (1 -µ + γσ(n)) 2 ⇐⇒ (1 + µ -γσ(n)) 2 -4µ < (1 -µ + γσ(n)) 2 ⇐⇒ (1 + µ -γσ(n)) 2 -(1 -µ + γσ(n)) 2 < 4µ ⇐⇒ (1 + µ -γσ(n) + 1 -µ + γσ(n)) (1 + µ -γσ(n) -1 + µ -γσ(n)) < 4µ ⇐⇒ 4(µ -γσ(n)) < 4µ ⇐⇒ 0 < γσ(n) Therefore, the conditions for convergence are i) µ ≤ (1 -γσ(n)) 2 , ii) 1 + µ -γσ(n) ≥ 0 and iii) γσ(n) > 0. Then by ii), µ ≥ γσ(n) -1 ≥ -1. Suppose for a contradiction that µ ≥ 1, then by i), 1 ≤ µ ≤ (1 -γσ(n))µ ≤ (1 -γσ(n)) 2 =⇒ µ ≤ 1 -2 γσ(n) + γσ(n) =⇒ µ + 2( γσ(n) -1) ≤ γσ(n) -1 =⇒ µ + 2 √ µ ≤ γσ(n) -1 =⇒ 1 + µ -γσ(n) + 2 √ µ ≤ 0.Then, we must have a last condition for convergence, iv) |µ| < 1.1bii:If 1 + µ -γσ(n) ≤ 0, then |λ 2 | ≥ |λ 1 | and λ 2 = (1 + µ -γσ(n)) - √ ∆ 2< 0For the system to be stable, we needλ 2 > -1 ⇐⇒ 3 + µ -γσ(n) > √ ∆ ≥ 0. Hence, we require i) µ ≤ (1 -γσ(n)) 2 , ii) 1 + µ -γσ(n) ≤ 0, iii) 3 + µ -γσ(n) > 0 and iv) (3 + µ -γσ(n)) 2 > ∆ ⇐⇒ (3 + µ -γσ(n)) 2 > (1 + µ -γσ(n)) 2 -4µ ⇐⇒ 4µ > (1 + µ -γσ(n)) 2 -(3 + µ -γσ(n)) 2 ⇐⇒ 4µ > (1 + µ -γσ(n) + 3 +µ -γσ(n)) (1 + µ -γσ(n) -3 -µ + γσ(n)) ⇐⇒ 4µ > -4(2 + µ -γσ(n)) ⇐⇒ 2 + 2µ -γσ(n) > 0Combining ii) and iv), we have, 2 + 2µµ -1 > 0 =⇒ µ > -1. Similarly, we assume for a contradiction that µ ≥ 1. Then, γσ(n) ≥ 1 + µ ≥ 2 and by condition i), √ µ ≤ γσ(n) -1 from which follows the contradiction1 ≤ γσ(n) - √ µ =⇒ γσ(n) + √ µ ≤ ( γσ(n) - √ µ)( γσ(n) + √ µ) = γσ(n) -µ =⇒ γσ(n) + √ µ ≤ γσ(n) -µ < 3 (by iii)) =⇒ √ µ + γσ(n) -1 < 2 =⇒ 2 √ µ < 2 (by i) again).Therefore, we need v) |µ| < 1 for convergence as well.To summarize, for Case 1, in order to have a stable system, we require:a) µ ≤ (1 -γσ(n)) 2 b) If 1 + µ -γσ(n) ≥ 0, i) γσ(n) > 0 ii) |µ| < 1 c) If 1 + µ -γσ(n) ≤ 0, i) 2 + 2µ -γσ(n) > 0 ii) |µ| < 1 =⇒ i) 3 + µ -γσ(n) > 0 ii) 2 + 2µ -γσ(n) > 0 iii) |µ| < 1 Case 2: ∆ < 0 when (1 + γσ(n)) 2 > µ > (1 -γσ(n)) 2Then, the eigenvalues are complex and given byλ { 1 2 } = (1 + µ -γσ(n)) ± i 4µ -(1 + µ -γσ(n)) 2 2 .For convergence, we require |λ 1,2 | = √ µ < 1 or equivalently, µ < 1. Hence, the necessary condition becomes 1 > µ > (1γσ(n)) 2 and to avoid a contradiction, we also require 1 > (1γσ(n)) 2 . This is satisfied when γσ(n) < 2 + 2µ and |µ| < 1.Therefore, from all the considered cases, to have a stable system, we require 0 < γσ(n) < 2 + 2µ and |µ| < 1, concluding the proof.To provide intuition behind robust momentum, we follow [10] and view the update rule in Eqn. 16 as a Lur'e feedback control system. For simplicity, we consider the case where, x t ∈ R and write the update in matrix formEQUATIONIn the frequency domain, in order for the system 20 to be stable, we require that for all |z| = 1,Re((1 -pz -1 )((k -1) G(pz) -1)) < 0,where G(pz) is a transformed transfer function, and p, a scaling factor. The imaginary axis, where Re = 0, is the stability boundary, and the specific construction of the parameters γ, µ and α pushes this boundary into the negative real axis to -v = (1 + p)(1k + 2kpkp 2 )/2p, hence achieving robust stability through the design of γ, µ and α.In this section, we provide the full plots presented in Section 2.3 and 6. We visualize the average norm of the SMoE and MoE outputs respectively, for 80 epochs of training in all 6 layers in Fig. 6 and 7 . Notably, in all plots, less a minor exception in layer 3 of SMoE, there is a decrease in the norm of each SMoE and MoE layer output throughout training. This is consistent with our expectation that the gate learns the optimal α * in a multi-objective optimization problem (Eqn. 8 and 9). In Fig. 8 , 9 and 10, we present the proportion of time each expert is chosen in 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 on the size of the norm of the expert outputs should yield improved performance of the model. We demonstrate in Section 5.1 that both MomentumSMoE and AdamSMoE significantly exceed the baseline performance and correspondingly, flattens the normed-based load distribution among experts as observed in Fig. 8 , 9 and 10. These strong empirical evidences serve to reinforce our optimization framework for SMoE.T 0 5 1 Z j / v Z E 5 A S J W M = " > A A A B 7 3 i c b V D L T s J A F L 3 F F + I L d e l m I j H B D W k N P p Z E N y 4 x k U d C G z I d p j B h O q 0 z U y M p / I Q b F x r j 1 t 9 x 5 9 8 4 Q B c K n u Q m J + f c m 3 v v 8 W P O l L b t b y u 3 s r q 2 v p H f L G x t 7 + z u F f c P m i p K J K E N E v F I t n 2 s K G e C N j T T n L Z j S X H o c 9 r y h z d T v / V I p W K R u N e j m H o h 7 g s W M I K 1 k d r u O C g / n b r j b r F k V + w Z 0 D J x M l K C D P V u 8 c v t R S Q J q d C E Y 6 U 6 j h 1 r L 8 V S M 8 L p p O A m i s a Y D H G f d g w V O K T K S 2 f 3 T t C J U X o o i K Q p o d F M / T 2 R 4 l C p U e i b z h D r g V r 0 p u J / X i f R w Z W X M h E n m g o y X x Q k H O k I T Z 9 H P S Y p 0 X x k C C a S m V s R G W C J i T Y R F U w I z u L L y 6 R 5 V n E u K u d 3 1 V L t O o s j D 0 d w D G V w 4 B J q c A t 1 a A A B D s / w C m / W g / V i v V s fV + w Z 0 D J x M l K C D P V u 8 c v t R S Q J q d C E Y 6 U 6 j h 1 r L 8 V S M 8 L p p O A m i s a Y D H G f d g w V O K T K S 2 f 3 T t C J U X o o i K Q p o d F M / T 2 R 4 l C p U e i b z h D r g V r 0 p u J / X i f R w Z W X M h E n m g o y X x Q k H O k I T Z 9 H P S Y p 0 X x k C C a S m V s R G W C J i T Y R F U w I z u L L y 6 R 5 V n E u K u d 3 1 V L t O o s j D 0 d w D G V w 4 B J q c A t 1 a A A B D s / w C m / W g / V i v V s fFor clarity, we summarize the new models developed in this paper for each task in ✓ ✓ ✓ GLaM ImageNet-1K V-MoE ✓ × ✓ ImageNet-R ImageNet-A ImageNet-C ImageNet-1K Soft MoE ✓ × ×across both SMoE and MoE, hence we did not integrate AdamW and Robust Momentum into Soft MoE. We leave these implementations for future work.Dataset: The WikiText-103 dataset [43] [46] . This adversarial attack randomly replaces words in the dataset with a generic "AAA" for evaluation making it difficult for the model to predict the next word in the sequence correctly.We use the Switch Transformer [18] , referred to 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 a SMoE layer, respectively. For MomentumSMoE, we replace each MLP layer with a MomentumSMoE layer and initialise each momentum vector, p 0 at 0. We do 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 6 layers, but we include a 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 in every SMoE, MomentumSMoE and AdamSMoE layer and selects 2 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 https://github.com/ofirpress/sandwich_transformer and https://github.com/giangdip2410/CompeteSMoE/tree/main. AdamSMoE/AdamGLaM: It is observed that Adam has certain divergent behavior during largescale training leading to unstable loss curves [8, 44] . In line with this observation, during the initial implementations of AdamSMoE, we 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 use AdamSMoE in 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.Datasets: We use the ImageNet-1K 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 to the ImageNet-1K validation set with 5 levels of severity. We provide a breakdown of our results on each corruption type and the mean Corruption Error (mCE) across which contains various artistic renditions of images. This evaluates the model's generalization ability to abstract visual renditions.Metrics: On ImageNet-1K, ImageNet-A and ImageNet-R, we report the top-1 accuracies for all experiments. On ImageNet-C, the standard metric for evaluation is the mCE. To calculate this, we average the top-1 error rate for each corruption 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 or decreasing values of these metrics signifying greater robustness will be indicated in the table with an arrow.We use a small Vision Mixture of Experts (V-MoE) [62] model as the SMoE baseline for ImageNet-1K object recognition task as well as the standard robustness benchmarks. This variant of V-MoE consists of 8 Vision Transformer (ViT) blocks [15] with every odd block's MLP being replaced by a SMoE layer. In turn, in MomentumV-MoE, we replace every other MLP layer with a MomentumSMoE layer and similarly for Robust MomentumV-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 is available here https://github.com/google-research/vmoe/.For our MoE baseline on clean ImageNet-1K data, we use Soft Mixture of Experts (Soft MoE) [56] .A Soft MoE model is designed to side step the challenging discrete optimization problem of 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 a certain number of slots and each slot is allocated a weighted average of all tokens. These weights are dependant on both the tokens and the experts. In this case, Soft MoE is 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 6 layers replace the MLP in those blocks with a Soft MoE layer. In Momentum-Soft MoE, we implement the momentum parameter into each Soft-MoE layer as in a SMoE layer.As there were no training details provided for Soft MoE on ImageNet-1K, we follow the training procedure in [69] for 120 epochs, using their published code base https://github.com/facebookresearch/deit. We train Momentum-Soft MoE and the baseline model using the PyTorch implementation of Soft MoE at https://github.com/bwconrad/soft-moe.In this section, we provide the pseudocode as written in python for MomentumSMoE, AdamSMoE, and Robust MomentumSMoE for clarification on our implementation. These are found in Figures 11, 12 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 γ fixed at 1.0, and the MomentumSMoE time-varying (TV) layer is as followsp t = -f (x t ) + t -1 t + 2 p t-1 ; x t+1 = x t + γp tZero-order hold µ and γ: Recall Eqn. 10, the proposed accelerated MomentumSMoE layerp t = -f (x t ) + µp t-1 ; x t+1 = x t + γp t .When we replace MLP layers with MomentumSMoE layers, as we do in our language model Mo-mentumSMoE, each expert function is a linear network. Explicitly expressing the MomentumSMoE layer results in the following expressionp t = (- E i=1 g(x t ) i W ⊤ i )x t + µp t-1 = Bx t + µp t-1 ; x t+1 = x t + γp tand can be interpreted as a state space representation with constant state and output matrices, µ and γ.From this perspective, we experiment with learning a discretized µ t and γ t by applying the Zero-order hold (ZOH) such that they become adaptive parameters as is a common practise in optimization algorithms. µ and γ are then parameterized as follows µ t = e ∆µ ; γ t = (∆µ) -1 (e ∆µ -1)∆γ where ∆ is the Softplus function of a learned scalar parameter, µ a learned scalar parameter and γ a linear network with scalar outputs.Results: We report the PPL for MomentumSMoE (TV) and MomentumSMoE (ZOH) on 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 γ as the sigmoid of a linear network. In this setting, every input has a different learning rate. While these models do improve over the baseline, there is no advantage over the standard MomentumSMoE and the other learnable µ and γ settings. Hence, we do not recommend implementations in this section over those and keep the results here for reference.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 as well. This aligns with the analogy of each layer being a GD step and with a higher number of iterations, the momentum term becomes more effective. While beneficial for large models, as implementing MomentumSMoE is computationally cost efficient, and hence, minimally affected by model depth, we note that there is an adverse effect in small models such as SMoE-small with only 3 layers. We hypothesize that the primary reason for this are the insufficient layers for the momentum term to make a positive impact and is a limitation of our method.We provide the full results of the top-1 accuracy and mCE of all 15 corruption types in ImageNet-C for V-MoE, MomentumV-MoE, Robust MomentumV-MoE and Sharpness Aware MomentumV-MoE (SAM-V-MoE), developed in Appendix E.4, in Table 8 . Included in Figure 14 is a 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 MomentumV-MoE outperforms the baseline V-MoE, with as high as a 6.5% increase in top-1 accuracy and 8 mCE decrease on fog corruption.The advantage of an optimization perspective for SMoEs are the countably many descent algorithms available that can be used to improve the model. In this section, we elaborate on six more extensions and Robust MomentumV-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 momentum method a 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 in an undesirable direction too quickly, preventing large oscillations especially when the learning rate is high. We implement NAG in our SMoE gradient framework asp t = -γf (x t -µp t-1 ) + µp t-1 ; x t+1 = x t + p twhere f (x t ) is the SMoE output and µ ∈ (0, 1) and γ > 0 are two hyperparameters corresponding to the momentum coefficient and step size respectively. We refer to this implementation as a NAG-SMoE.Time-varying momentum with scheduled restart: An enhancement of the time-varying momentum covered in Section E.1 is to include a scheduled restart for the momentum parameter, t -1/t + 2. Such a modification can help to recover an optimal convergence rate as accelerated methods do not necessarily maintain a monotonic decrease in objective value [64, 47] . As our model has 6 layers, we choose to restart after layer 3 and the RS-SMoE update isp t = -f (x t ) + t mod 3 t mod 3 + 3 p t-1 ; x t+1 = x t + γp twhere f (x t ) is the SMoE output and γ > 0 is the step size.RMSprop: RMSprop is an adaptive step size algorithm that scales the gradient by an exponentially decaying average of the squared gradients [27] . The algorithm replaces a global learning rate, which is 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 a running 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, we have the following update in our rms-SMoE layer,p t = µp t-1 + (1 -µ)f (x t ) 2 ; x t+1 = x t - γ √ p t + ϵ f (x t )where f (x t ) is the SMoE output, µ ∈ (0, 1) is the moving average parameter and, γ > 0 the step size. Similar to AdamSMoE, 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, a flatter minima should yield greater generalization ability [17, 29, 31] . Further, this would increase the model'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 l 2 -norm in the SMoE optimization framework yields the SAM-SMoE layerε(x t ) = ρf (x t )/∥f (x t )∥ 2 ; p t = f (x t + ε(x t )); x t+1 = x t -γp twhere f (x t ) is the SMoE output, ρ > 0 is a hyperparameter controlling the size of the neighbourhood and, γ > 0 the step size.Complex and Negative Momentum: Classic momentum works well in a gradient-based optimization when the Jacobian of our update step, as a fixed-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 its oscillations at fixed frequencies between adding and subtracting the momentum term. We oscillate the sign of the momentum 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-MomentumSMoE layer:EQUATIONwhere µ ∈ C and γ > 0 are hyperparameters corresponding to the momentum coefficient and step size, respectively and ℜ extracts the real component of the momentum term.Results: On the language modeling task, though all models, except NAG-SMoE on 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 with the exception of Complex-MomentumSMoE, which has an even greater improvement over the baseline than MomentumSMoE. Complex-MomentumSMoE's enhanced results further verify the power and promise of our momentum-based framework in designing SMoE. On the computer vision task, we find that SAM-V-MoE does perform relatively well on corruption benchmarks, exceeding the baseline by more than 1% on ImageNet-A and ImageNet-C.A 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 MomentumSMoE and AdamSMoE as compared to 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 metrics at both training and test time with negligible computational cost.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 AI chatbot assistants. We show significant advancements in resisting data perturbation, aiming to protect critical AI systems from malicious actors. Furthermore, we achieve competitive performance in language modeling with contaminated data, reflecting real-world scenarios where data is often imperfect. While the potential for AI misuse exists, our work provides substantial improvements in fundamental architectures and theory, which we hope will lead to further socially beneficial outcomes. 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: The paper poses no such risks.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a 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 to usage guidelines or restrictions to 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 is challenging, and many papers do not require this, but we encourage authors to take this into account and make a 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 use and the baselines we compare with in our manuscript.Guidelines:• The answer NA means that the paper does not use 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 is used and, if 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.• If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
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/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was 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. 
crowdsourcing 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 as details about compensation (if any)?Answer: [NA]Justification: The paper does not involve crowdsourcing or research with human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution 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. 
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 an equivalent approval/review based on the requirements of your country or 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 nor research with human subjects. • Depending on 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.• We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to 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. 
 
 Appendix:  


Decision: 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.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 1. This method has little effect on models with few layers.

2. The largest models for evaluation only have  388M parameters, which are much smaller than mainstream MoE LLMs.

3. From a theoretical standpoint, developing a framework to explain the enhanced robustness of MomentumSMoE would be interesting.
questions: questions: please refer to weaknesses
limitations: limitations: The authors have adequately discussed the limitations.
paper_summary: paper_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.
strength_weakness: strength_weakness: (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. 

(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.

(3) There are some grammar issues. Please use `` instead of " in the paper (line 665).

(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.

I'd like to hear a more insightful discussion regarding all the points above from the authors.
questions: questions: (1) Please explain more of line 140 ("Thus, it is expected that these two terms learn to reduce ...).
limitations: limitations: N/A
paper_summary: paper_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.
strength_weakness: strength_weakness: The pseudocode may be provided to better illustrate the implementation of the proposal.
questions: 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?
2. 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?
3. As stated in the weaknesses, the authors could provide pseudocode to better clarify their proposal.
limitations: limitations: The authors adequately addressed the limitations.
paper_summary: paper_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.
The 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.
strength_weakness: strength_weakness: 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.

But, 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

The 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.

Azizian, Waïss, et al. "Accelerating smooth games by manipulating spectral shapes." International Conference on Artificial Intelligence and Statistics. PMLR, 2020.
Lorraine, Jonathan P., et al. "Complex momentum for optimization in games." International Conference on Artificial Intelligence and Statistics. PMLR, 2022.
Gidel, Gauthier, et al. "Negative momentum for improved game dynamics." The 22nd International Conference on Artificial Intelligence and Statistics. PMLR, 2019.
questions: questions: How can the assumptions about the fixed point operator's spectrum and the Jacobian's conservativeness be validated or relaxed in practical scenarios?

Are there more general acceleration tools than momentum you might want to use for this problem?
limitations: limitations: The limitations are discussed briefly, but a delineated section elaborating on all the limitations would be valuable.

Comments: title: Thanks for your endorsement!
Thanks for your endorsement!
After Rebuttal
Any Questions from Reviewer puL2 on Our Rebuttal?
Any Questions from Reviewer qEGD on Our Rebuttal?
Any Questions from Reviewer oJG3 on Our Rebuttal?
Code Snippet Part 4/4
Code Snippet Part 3/4
Code Snippet Part 2/4
Code Snippet Part 1/4
Reply to Reviewer GFpG
Clarification on Figure 1 in Our Attached PDF
Thanks for your endorsement!
Response to Rebuttal
comment: Thanks the reviewer for your constructive responses and interesting suggestions on the spectrum analysis of the fixed point operator. We are grateful for the time taken to provide good insights into further improvements to our paper. We will continue to explore these new perspectives in our future work. We appreciate your endorsement and the fruitful discussions.
I have run the code you provided, and those visualizations look good.  Perhaps a heatmap would look reasonable in the plots, but that is a personal preference. This is, of course, not required for this paper, but for future analysis, might I suggest looking at the eigenvalues of fixed point operator (as in Figure 4, https://arxiv.org/pdf/1807.04740), which combined information about your optimization scheme, with the spectrum of $\nabla_x f$.  For example, for gradient descent which uses an update of $x = x - f$, for a fixed point operator $g(x) = x - f$, then the you would want to look at $eigs(\nabla_x g) = eigs(I - \nabla_x f)$.  The eigenvalues of the fixed point operator directly tell you convergence rates (in idealized scenarios), and looking at how your optimizer warps the spectrum may provide insight into the best acceleration schemes. You might want to look at the union of the eigenvalues over multiple different "layers" and optimization setups / different problems for the experts.  It could also be worth looking at what eigenspaces your updates are actually accumulating in, as this is non-convex (so most theoretical results will say it doesn't converge).  Overall, I am happy with your response, and I think this a fun topic to ponder.

```fig, ax = plt.subplots(1, 3, figsize=(16, 4), sharey=True)
fontsize = 15
bins = 51  # Number of bins for the heatmap. Better visualization if this number is odd imo

# Heatmap for the first set of data
h0 = ax[0].hist2d(abs0, phase0, bins=bins, cmap='Blues')
ax[0].set_ylim([-3.2, 3.2])
ax[0].set_yticks(ticks=[-3.14, -1.57, 0, 1.57, 3.14])
ax[0].set_yticklabels(labels=[-3.14, -1.57, 0, 1.57, 3.14], fontsize=fontsize-2)
ax[0].grid()
ax[0].set_title("Layer 1", fontsize=fontsize)
ax[0].tick_params(axis='both', which='major', labelsize=fontsize-2)
fig.colorbar(h0[3], ax=ax[0])

# Heatmap for the second set of data
h2 = ax[1].hist2d(abs2, phase2, bins=bins, cmap='Oranges')
ax[1].grid()
ax[1].set_title("Layer 3", fontsize=fontsize)
ax[1].tick_params(axis='both', which='major', labelsize=fontsize-2)
fig.colorbar(h2[3], ax=ax[1])

# Heatmap for the third set of data
h5 = ax[2].hist2d(abs5, phase5, bins=bins, cmap='Greens')
ax[2].grid()
ax[2].set_title("Layer 6", fontsize=fontsize)
ax[2].tick_params(axis='both', which='major', labelsize=fontsize-2)
fig.colorbar(h5[3], ax=ax[2])

plt.tight_layout()
plt.show()
Thanks for your response, and we appreciate your endorsement.
Thanks to the detailed response from authors. Most of my concerns are well addressed. I'd like to increase the score to 5.
We would like to thank the reviewer again for your thoughtful reviews and valuable feedback.

We would appreciate it if you could let us know if our responses have addressed your concerns and whether you still have any other questions about our rebuttal.

We would be happy to do any follow-up discussion or address any additional comments.
We would like to thank the reviewer again for your thoughtful reviews and valuable feedback.

We would appreciate it if you could let us know if our responses have addressed your concerns and whether you still have any other questions about our rebuttal.

We would be happy to do any follow-up discussion or address any additional comments.
We would like to thank the reviewer again for your thoughtful reviews and valuable feedback.

We would appreciate it if you could let us know if our responses have addressed your concerns and whether you still have any other questions about our rebuttal.

We would be happy to do any follow-up discussion or address any additional comments.
~~~#python
# layer 1
abs0 = torch.log(eig0_base.abs())
phase0 = eig0_base.angle()

# layer 3
abs2 = torch.log(eig2_base.abs())
phase2 = eig2_base.angle()

# layer 6
abs5 = torch.log(eig5_base.abs())
phase5 = eig5_base.angle()

fig, ax = plt.subplots(1, 3, figsize=(16,4), sharey=True)
fontsize = 15
linewidth = 4
ax[0].scatter(abs0, phase0, s = 30, c = "blue", alpha=0.7,edgecolors="blue")
ax[1].scatter(abs2, phase2, s = 30, c = "orange", alpha=0.7,edgecolors="orange")
ax[2].scatter(abs5, phase5, s = 30, c = "green", alpha=0.7,edgecolors="green")
ax[0].set_ylim([-3.2,3.2])
ax[0].set_yticks(ticks=[-3.14, -1.57, 0, 1.57, 3.14])
ax[0].set_yticklabels(labels=[-3.14, -1.57, 0, 1.57, 3.14], fontsize=fontsize-2)
ax[0].grid()
ax[1].grid()
ax[2].grid()
ax[0].set_title("Layer 1", fontsize=fontsize)
ax[0].tick_params(axis='both', which='major', labelsize=fontsize-2)
ax[1].set_title("Layer 3", fontsize=fontsize)
ax[1].tick_params(axis='both', which='major', labelsize=fontsize-2)
ax[2].set_title("Layer 6", fontsize=fontsize)
ax[2].tick_params(axis='both', which='major', labelsize=fontsize-2)
plt.tight_layout()
plt.show()
~~~
~~~#python 
# layer 6
eig5_base = torch.tensor([ 2.64+1.62j, 2.64-1.62j, -16.31+0.77j, -16.31-0.77j, -17.69+0.09j, -17.69-0.09j, -16.41+0.37j, -16.41-0.37j, -7.03+1.32j, -7.03-1.32j, -12.89+0.83j,
-12.89-0.83j, 11.58+0.26j, 11.58-0.26j, 11.76+0.04j, 11.76-0.04j, 9.81+0.68j, 9.81-0.68j, 1.99+1.23j, 1.99-1.23j, 0.71+1.20j, 0.71-1.20j,
-4.73+1.19j, -4.73-1.19j, -1.76+1.14j, -1.76-1.14j, -8.65+1.06j, -8.65-1.06j, -6.52+1.10j, -6.52-1.10j, -3.00+1.12j, -3.00-1.12j, -12.92+0.64j,
-12.92-0.64j, -14.16+0.48j, -14.16-0.48j, -10.68+0.80j, -10.68-0.80j, -15.07+0.00j, -7.74+0.91j, -7.74-0.91j, -11.19+0.62j, -11.19-0.62j, -13.21+0.32j,
-13.21-0.32j, -13.05+0.00j, -12.84+0.08j, -12.84-0.08j, -8.55+0.79j, -8.55-0.79j, -11.98+0.28j, -11.98-0.28j, 2.43+1.05j, 2.43-1.05j, 0.87+1.02j,
0.87-1.02j, 5.79+0.84j, 5.79-0.84j, 7.31+0.63j, 7.31-0.63j, 3.47+0.94j, 3.47-0.94j, 2.38+0.95j, 2.38-0.95j, 5.07+0.81j, 5.07-0.81j,
8.30+0.45j, 8.30-0.45j, 9.42+0.00j, 8.86+0.13j, 8.86-0.13j, 7.75+0.41j, 7.75-0.41j, -5.44+0.93j, -5.44-0.93j, -2.96+0.92j, -2.96-0.92j,
-10.01+0.47j, -10.01-0.47j, -10.96+0.00j, -10.29+0.25j, -10.29-0.25j, -5.31+0.73j, -5.31-0.73j, -6.92+0.61j, -6.92-0.61j, 0.95+0.88j, 0.95-0.88j,
5.88+0.58j, 5.88-0.58j, 4.95+0.67j, 4.95-0.67j, 7.76+0.07j, 7.76-0.07j, 3.51+0.73j, 3.51-0.73j, -9.62+0.00j, -9.44+0.00j, 6.15+0.31j,
6.15-0.31j, 0.37+0.84j, 0.37-0.84j, 1.72+0.74j, 1.72-0.74j, -2.77+0.74j, -2.77-0.74j, -4.83+0.73j, -4.83-0.73j, -8.99+0.00j, -7.44+0.44j,
-7.44-0.44j, -5.92+0.59j, -5.92-0.59j, 7.24+0.00j, 6.34+0.25j, 6.34-0.25j, -2.43+0.74j, -2.43-0.74j, -7.86+0.28j, -7.86-0.28j, -8.53+0.08j,
-8.53-0.08j, -7.30+0.29j, -7.30-0.29j, 5.34+0.41j, 5.34-0.41j, 6.51+0.16j, 6.51-0.16j, 6.51+0.00j, -4.92+0.49j, -4.92-0.49j, -2.04+0.63j,
-2.04-0.63j, 0.72+0.63j, 0.72-0.63j, 3.77+0.51j, 3.77-0.51j, 2.83+0.54j, 2.83-0.54j, 5.87+0.00j, 5.03+0.23j, 5.03-0.23j, 4.88+0.10j,
4.88-0.10j, 3.61+0.38j, 3.61-0.38j, -6.84+0.04j, -6.84-0.04j, -6.94+0.00j, -5.80+0.31j, -5.80-0.31j, -6.03+0.19j, -6.03-0.19j, 0.66+0.57j,
0.66-0.57j, 0.01+0.58j, 0.01-0.58j, -3.12+0.51j, -3.12-0.51j, -2.70+0.53j, -2.70-0.53j, -0.77+0.55j, -0.77-0.55j, -4.14+0.40j, -4.14-0.40j,
4.56+0.09j, 4.56-0.09j, 2.87+0.40j, 2.87-0.40j, 1.61+0.50j, 1.61-0.50j, -5.38+0.12j, -5.38-0.12j, -4.35+0.34j, -4.35-0.34j, -3.38+0.41j,
-3.38-0.41j, -1.32+0.49j, -1.32-0.49j, 0.34+0.47j, 0.34-0.47j, -4.58+0.25j, -4.58-0.25j, -5.38+0.00j, 3.55+0.26j, 3.55-0.26j, 3.23+0.29j,
3.23-0.29j, -4.87+0.00j, -1.79+0.46j, -1.79-0.46j, 1.35+0.42j, 1.35-0.42j, 0.38+0.45j, 0.38-0.45j, 3.69+0.08j, 3.69-0.08j, 3.33+0.17j,
3.33-0.17j, 0.72+0.42j, 0.72-0.42j, -2.44+0.39j, -2.44-0.39j, -3.58+0.27j, -3.58-0.27j, -4.32+0.18j, -4.32-0.18j, -4.83+0.09j, -4.83-0.09j,
-4.77+0.00j, -2.68+0.33j, -2.68-0.33j, -3.98+0.08j, -3.98-0.08j, -1.17+0.36j, -1.17-0.36j, 2.79+0.20j, 2.79-0.20j, 0.92+0.33j, 0.92-0.33j,
3.30+0.00j, 3.13+0.04j, 3.13-0.04j, -0.09+0.35j, -0.09-0.35j, 1.96+0.25j, 1.96-0.25j, 2.61+0.14j, 2.61-0.14j, -1.34+0.31j, -1.34-0.31j,
-2.04+0.28j, -2.04-0.28j, -2.48+0.25j, -2.48-0.25j, -3.01+0.15j, -3.01-0.15j, -3.68+0.00j, -0.08+0.29j, -0.08-0.29j, 1.42+0.23j, 1.42-0.23j,
1.89+0.18j, 1.89-0.18j, -3.32+0.00j, -3.03+0.07j, -3.03-0.07j, -2.68+0.00j, -1.82+0.23j, -1.82-0.23j, -1.23+0.26j, -1.23-0.26j, -2.22+0.17j,
-2.22-0.17j, -2.34+0.14j, -2.34-0.14j, 1.01+0.24j, 1.01-0.24j, 2.39+0.03j, 2.39-0.03j, -0.34+0.25j, -0.34-0.25j, -1.49+0.18j, -1.49-0.18j,
1.89+0.12j, 1.89-0.12j, 1.97+0.02j, 1.97-0.02j, 1.40+0.16j, 1.40-0.16j, 1.00+0.20j, 1.00-0.20j, 0.01+0.21j, 0.01-0.21j, 0.67+0.21j,
0.67-0.21j, -0.83+0.18j, -0.83-0.18j, -1.95+0.10j, -1.95-0.10j, -2.07+0.03j, -2.07-0.03j, 1.13+0.12j, 1.13-0.12j, 1.35+0.08j, 1.35-0.08j,
0.87+0.11j, 0.87-0.11j, -0.51+0.15j, -0.51-0.15j, -2.04+0.00j, -1.30+0.11j, -1.30-0.11j, -1.57+0.00j, -1.23+0.08j, -1.23-0.08j, -0.24+0.14j,
-0.24-0.14j, 0.31+0.14j, 0.31-0.14j, 1.19+0.05j, 1.19-0.05j, 1.30+0.00j, 1.27+0.00j, -0.65+0.12j, -0.65-0.12j, -1.41+0.02j, -1.41-0.02j,
-1.33+0.00j, 0.50+0.10j, 0.50-0.10j, 0.03+0.11j, 0.03-0.11j, -0.95+0.06j, -0.95-0.06j, 0.79+0.06j, 0.79-0.06j, 0.88+0.04j, 0.88-0.04j,
-0.63+0.06j, -0.63-0.06j, -0.45+0.09j, -0.45-0.09j, -0.92+0.02j, -0.92-0.02j, -0.83+0.00j, -0.12+0.07j, -0.12-0.07j, 0.63+0.02j, 0.63-0.02j,
0.35+0.05j, 0.35-0.05j, -0.69+0.03j, -0.69-0.03j, -0.27+0.06j, -0.27-0.06j, -0.44+0.05j, -0.44-0.05j, 0.45+0.02j, 0.45-0.02j, 0.40+0.02j,
0.40-0.02j, 0.03+0.04j, 0.03-0.04j, -0.47+0.00j, -0.01+0.02j, -0.01-0.02j, -0.09+0.02j, -0.09-0.02j, 0.03+0.01j, 0.03-0.01j, -0.32+0.00j])
~~~
~~~#python
# layer 3
eig2_base = torch.tensor([-46.60+0.00j, -17.27+0.00j, -11.27+1.00j, -11.27-1.00j, 15.41+0.00j, 14.28+0.31j, 14.28-0.31j, -13.24+0.34j, -13.24-0.34j, -13.74+0.00j, -12.13+0.48j, -12.13-0.48j, -9.66+0.79j,
-9.66-0.79j, -6.73+1.06j, -6.73-1.06j, 6.03+1.16j, 6.03-1.16j, -11.59+0.00j, 11.38+0.67j, 11.38-0.67j, -8.53+0.72j, -8.53-0.72j, -2.86+1.17j, -2.86-1.17j, 10.74+0.54j,
10.74-0.54j, 9.84+0.64j, 9.84-0.64j, 8.81+0.74j, 8.81-0.74j, -3.41+1.02j, -3.41-1.02j, -2.00+1.04j, -2.00-1.04j, 4.80+0.95j, 4.80-0.95j, 0.95+1.05j, 0.95-1.05j,
2.01+1.03j, 2.01-1.03j, 10.79+0.33j, 10.79-0.33j, 10.51+0.02j, 10.51-0.02j, 9.58+0.31j, 9.58-0.31j, 8.68+0.53j, 8.68-0.53j, 6.59+0.78j, 6.59-0.78j, 0.32+0.98j,
0.32-0.98j, -5.63+0.88j, -5.63-0.88j, -10.85+0.33j, -10.85-0.33j, -9.14+0.52j, -9.14-0.52j, -9.88+0.38j, -9.88-0.38j, -10.63+0.07j, -10.63-0.07j, -4.31+0.86j, -4.31-0.86j,
8.75+0.22j, 8.75-0.22j, 6.86+0.57j, 6.86-0.57j, -6.09+0.76j, -6.09-0.76j, -8.47+0.52j, -8.47-0.52j, -6.74+0.64j, -6.74-0.64j, 4.72+0.76j, 4.72-0.76j, 3.99+0.78j,
3.99-0.78j, 7.07+0.46j, 7.07-0.46j, 8.62+0.04j, 8.62-0.04j, -8.81+0.31j, -8.81-0.31j, -9.04+0.13j, -9.04-0.13j, 0.41+0.84j, 0.41-0.84j, -0.70+0.84j, -0.70-0.84j,
-2.68+0.77j, -2.68-0.77j, -6.54+0.54j, -6.54-0.54j, -7.23+0.43j, -7.23-0.43j, -2.63+0.71j, -2.63-0.71j, -4.57+0.55j, -4.57-0.55j, 7.40+0.25j, 7.40-0.25j, 2.30+0.71j,
2.30-0.71j, 0.30+0.72j, 0.30-0.72j, 4.17+0.61j, 4.17-0.61j, 5.69+0.48j, 5.69-0.48j, 6.00+0.38j, 6.00-0.38j, 4.52+0.52j, 4.52-0.52j, 7.35+0.00j, -1.39+0.70j,
-1.39-0.70j, 6.76+0.12j, 6.76-0.12j, 6.03+0.23j, 6.03-0.23j, 2.19+0.63j, 2.19-0.63j, 4.58+0.46j, 4.58-0.46j, 3.79+0.52j, 3.79-0.52j, -2.84+0.61j, -2.84-0.61j,
-7.68+0.00j, -7.45+0.15j, -7.45-0.15j, -6.98+0.27j, -6.98-0.27j, -3.79+0.54j, -3.79-0.54j, -7.11+0.05j, -7.11-0.05j, -5.99+0.35j, -5.99-0.35j, -6.31+0.28j, -6.31-0.28j,
-4.61+0.47j, -4.61-0.47j, 5.44+0.20j, 5.44-0.20j, -1.53+0.58j, -1.53-0.58j, 1.56+0.58j, 1.56-0.58j, 0.67+0.59j, 0.67-0.59j, -0.20+0.59j, -0.20-0.59j, -6.01+0.00j,
-5.20+0.29j, -5.20-0.29j, -4.55+0.36j, -4.55-0.36j, -2.16+0.54j, -2.16-0.54j, 3.15+0.47j, 3.15-0.47j, 5.63+0.06j, 5.63-0.06j, 5.05+0.26j, 5.05-0.26j, 4.02+0.35j,
4.02-0.35j, -5.68+0.07j, -5.68-0.07j, -3.48+0.41j, -3.48-0.41j, -1.57+0.50j, -1.57-0.50j, -5.02+0.23j, -5.02-0.23j, 4.92+0.00j, 2.04+0.49j, 2.04-0.49j, 1.56+0.48j,
1.56-0.48j, 3.29+0.36j, 3.29-0.36j, 4.04+0.24j, 4.04-0.24j, -4.53+0.25j, -4.53-0.25j, -2.23+0.40j, -2.23-0.40j, -4.83+0.02j, -4.83-0.02j, 4.18+0.18j, 4.18-0.18j,
4.11+0.14j, 4.11-0.14j, 2.30+0.36j, 2.30-0.36j, 0.68+0.43j, 0.68-0.43j, -0.38+0.43j, -0.38-0.43j, -0.70+0.43j, -0.70-0.43j, -2.98+0.35j, -2.98-0.35j, -3.33+0.33j,
-3.33-0.33j, -4.04+0.22j, -4.04-0.22j, -4.40+0.08j, -4.40-0.08j, 3.82+0.08j, 3.82-0.08j, 2.92+0.25j, 2.92-0.25j, 0.94+0.38j, 0.94-0.38j, -0.12+0.36j, -0.12-0.36j,
-1.38+0.36j, -1.38-0.36j, -2.21+0.33j, -2.21-0.33j, -3.19+0.26j, -3.19-0.26j, -4.05+0.00j, -3.84+0.13j, -3.84-0.13j, -3.50+0.16j, -3.50-0.16j, -1.58+0.34j, -1.58-0.34j,
2.05+0.30j, 2.05-0.30j, 3.73+0.00j, 3.29+0.10j, 3.29-0.10j, -3.47+0.08j, -3.47-0.08j, -2.78+0.20j, -2.78-0.20j, -0.64+0.33j, -0.64-0.33j, -0.21+0.34j, -0.21-0.34j,
0.40+0.31j, 0.40-0.31j, 1.79+0.27j, 1.79-0.27j, 2.17+0.22j, 2.17-0.22j, 2.48+0.17j, 2.48-0.17j, 2.75+0.04j, 2.75-0.04j, 2.42+0.14j, 2.42-0.14j, 0.67+0.29j,
0.67-0.29j, 1.02+0.26j, 1.02-0.26j, -2.32+0.19j, -2.32-0.19j, -1.39+0.24j, -1.39-0.24j, -0.33+0.27j, -0.33-0.27j, -2.95+0.07j, -2.95-0.07j, -2.45+0.13j, -2.45-0.13j,
-2.82+0.03j, -2.82-0.03j, -2.79+0.00j, -0.57+0.24j, -0.57-0.24j, -0.90+0.21j, -0.90-0.21j, -1.10+0.18j, -1.10-0.18j, -1.72+0.13j, -1.72-0.13j, -1.79+0.10j, -1.79-0.10j,
0.88+0.21j, 0.88-0.21j, 0.13+0.21j, 0.13-0.21j, 2.16+0.10j, 2.16-0.10j, 2.35+0.00j, 2.17+0.03j, 2.17-0.03j, 1.93+0.10j, 1.93-0.10j, 1.64+0.14j, 1.64-0.14j,
0.71+0.17j, 0.71-0.17j, 1.39+0.11j, 1.39-0.11j, 1.53+0.07j, 1.53-0.07j, 1.76+0.01j, 1.76-0.01j, -2.01+0.00j, -1.83+0.04j, -1.83-0.04j, -1.44+0.08j, -1.44-0.08j,
-0.07+0.17j, -0.07-0.17j, -0.75+0.14j, -0.75-0.14j, -0.51+0.16j, -0.51-0.16j, 0.72+0.12j, 0.72-0.12j, -1.58+0.00j, -1.06+0.10j, -1.06-0.10j, -1.35+0.00j, -1.32+0.03j,
-1.32-0.03j, -0.13+0.13j, -0.13-0.13j, 0.92+0.06j, 0.92-0.06j, 1.18+0.02j, 1.18-0.02j, 0.34+0.11j, 0.34-0.11j, 0.84+0.01j, 0.84-0.01j, -0.62+0.08j, -0.62-0.08j,
0.16+0.08j, 0.16-0.08j, -0.23+0.09j, -0.23-0.09j, 0.62+0.03j, 0.62-0.03j, 0.20+0.05j, 0.20-0.05j, -0.90+0.00j, -0.53+0.02j, -0.53-0.02j, -0.26+0.05j, -0.26-0.05j, -0.24+0.00j])
~~~
~~~ #python
import torch
import matplotlib.pyplot as plt

# layer 1
eig0_base = torch.tensor([ -28.21+0.00j, -18.35+0.21j, -18.35-0.21j, -14.52+0.68j, -14.52-0.68j, -7.76+1.04j, -7.76-1.04j, -15.45+0.00j, -14.99+0.12j, -14.99-0.12j, -14.32+0.00j,
-12.24+0.64j, -12.24-0.64j, -4.99+1.07j, -4.99-1.07j, -12.49+0.36j, -12.49-0.36j, -9.83+0.76j, -9.83-0.76j, 1.91+1.01j, 1.91-1.01j, -10.79+0.61j,
-10.79-0.61j, -7.73+0.85j, -7.73-0.85j, -11.27+0.42j, -11.27-0.42j, -11.91+0.10j, -11.91-0.10j, -10.23+0.53j, -10.23-0.53j, -6.54+0.85j, -6.54-0.85j,
-4.04+0.96j, -4.04-0.96j, 7.57+0.00j, 6.55+0.47j, 6.55-0.47j, 0.15+0.98j, 0.15-0.98j, -7.79+0.71j, -7.79-0.71j, 4.91+0.65j, 4.91-0.65j,
-9.64+0.14j, -9.64-0.14j, -3.78+0.85j, -3.78-0.85j, 2.30+0.79j, 2.30-0.79j, -1.28+0.88j, -1.28-0.88j, 0.89+0.82j, 0.89-0.82j, -10.34+0.16j,
-10.34-0.16j, -7.92+0.60j, -7.92-0.60j, -6.47+0.67j, -6.47-0.67j, -10.16+0.00j, -7.87+0.54j, -7.87-0.54j, -8.77+0.29j, -8.77-0.29j, -2.87+0.80j,
-2.87-0.80j, -0.06+0.79j, -0.06-0.79j, 3.06+0.70j, 3.06-0.70j, 5.00+0.55j, 5.00-0.55j, 5.16+0.40j, 5.16-0.40j, 6.78+0.13j, 6.78-0.13j,
6.83+0.00j, 5.82+0.31j, 5.82-0.31j, 6.09+0.19j, 6.09-0.19j, -8.94+0.00j, 2.90+0.67j, 2.90-0.67j, -4.00+0.74j, -4.00-0.74j, -2.03+0.74j,
-2.03-0.74j, 0.21+0.73j, 0.21-0.73j, -5.03+0.65j, -5.03-0.65j, -8.39+0.00j, 1.27+0.66j, 1.27-0.66j, 2.54+0.62j, 2.54-0.62j, 4.75+0.43j,
4.75-0.43j, 3.28+0.54j, 3.28-0.54j, 5.78+0.16j, 5.78-0.16j, 5.42+0.24j, 5.42-0.24j, -0.75+0.69j, -0.75-0.69j, -4.10+0.61j, -4.10-0.61j,
-6.44+0.47j, -6.44-0.47j, -7.08+0.35j, -7.08-0.35j, -7.75+0.18j, -7.75-0.18j, -7.46+0.27j, -7.46-0.27j, 5.36+0.00j, 4.62+0.25j, 4.62-0.25j,
-2.09+0.61j, -2.09-0.61j, -5.08+0.48j, -5.08-0.48j, -5.27+0.45j, -5.27-0.45j, 4.99+0.00j, -3.61+0.55j, -3.61-0.55j, -6.72+0.24j, -6.72-0.24j,
-6.60+0.22j, -6.60-0.22j, -7.18+0.01j, -7.18-0.01j, -7.04+0.06j, -7.04-0.06j, -1.82+0.59j, -1.82-0.59j, -0.18+0.59j, -0.18-0.59j, 4.10+0.28j,
4.10-0.28j, 3.61+0.35j, 3.61-0.35j, 1.72+0.50j, 1.72-0.50j, 4.05+0.13j, 4.05-0.13j, 2.38+0.45j, 2.38-0.45j, 2.71+0.39j, 2.71-0.39j,
-3.14+0.51j, -3.14-0.51j, -0.91+0.56j, -0.91-0.56j, -1.11+0.54j, -1.11-0.54j, 0.36+0.51j, 0.36-0.51j, -2.56+0.51j, -2.56-0.51j, -6.10+0.00j,
-4.94+0.29j, -4.94-0.29j, -4.73+0.30j, -4.73-0.30j, 2.97+0.33j, 2.97-0.33j, 4.08+0.07j, 4.08-0.07j, 3.95+0.13j, 3.95-0.13j, 0.57+0.47j,
0.57-0.47j, 1.69+0.41j, 1.69-0.41j, 3.44+0.21j, 3.44-0.21j, 2.96+0.20j, 2.96-0.20j, -0.15+0.45j, -0.15-0.45j, -0.59+0.46j, -0.59-0.46j,
-2.92+0.42j, -2.92-0.42j, -5.39+0.23j, -5.39-0.23j, -4.43+0.33j, -4.43-0.33j, -3.61+0.35j, -3.61-0.35j, -5.59+0.09j, -5.59-0.09j, -3.97+0.29j,
-3.97-0.29j, -4.96+0.17j, -4.96-0.17j, -5.13+0.12j, -5.13-0.12j, -5.21+0.03j, -5.21-0.03j, 0.88+0.39j, 0.88-0.39j, 2.95+0.14j, 2.95-0.14j,
2.02+0.29j, 2.02-0.29j, -2.25+0.38j, -2.25-0.38j, -1.58+0.39j, -1.58-0.39j, -2.39+0.33j, -2.39-0.33j, -0.05+0.38j, -0.05-0.38j, 0.32+0.37j,
0.32-0.37j, 1.51+0.31j, 1.51-0.31j, 1.26+0.29j, 1.26-0.29j, 3.21+0.01j, 3.21-0.01j, 2.26+0.19j, 2.26-0.19j, 2.75+0.06j, 2.75-0.06j,
-3.49+0.21j, -3.49-0.21j, -3.91+0.14j, -3.91-0.14j, -4.12+0.00j, -4.04+0.04j, -4.04-0.04j, -1.50+0.33j, -1.50-0.33j, -0.80+0.35j, -0.80-0.35j,
2.50+0.09j, 2.50-0.09j, 1.70+0.21j, 1.70-0.21j, 0.99+0.26j, 0.99-0.26j, -3.46+0.13j, -3.46-0.13j, 0.17+0.26j, 0.17-0.26j, -0.02+0.27j,
-0.02-0.27j, -2.45+0.25j, -2.45-0.25j, -0.99+0.27j, -0.99-0.27j, -1.47+0.28j, -1.47-0.28j, -2.07+0.23j, -2.07-0.23j, -2.83+0.16j, -2.83-0.16j,
-3.11+0.11j, -3.11-0.11j, 1.85+0.15j, 1.85-0.15j, 2.20+0.00j, 2.16+0.03j, 2.16-0.03j, -3.17+0.00j, -2.25+0.18j, -2.25-0.18j, 1.36+0.16j,
1.36-0.16j, 1.79+0.07j, 1.79-0.07j, -2.97+0.00j, -2.66+0.06j, -2.66-0.06j, -1.69+0.19j, -1.69-0.19j, -0.91+0.23j, -0.91-0.23j, -1.31+0.20j,
-1.31-0.20j, -2.42+0.03j, -2.42-0.03j, -2.05+0.12j, -2.05-0.12j, -0.37+0.22j, -0.37-0.22j, 0.68+0.20j, 0.68-0.20j, -0.68+0.20j, -0.68-0.20j,
1.81+0.00j, 0.05+0.19j, 0.05-0.19j, 0.55+0.16j, 0.55-0.16j, 1.37+0.00j, 1.20+0.08j, 1.20-0.08j, 0.90+0.12j, 0.90-0.12j, -1.09+0.15j,
-1.09-0.15j, -1.81+0.10j, -1.81-0.10j, -1.81+0.05j, -1.81-0.05j, -0.28+0.16j, -0.28-0.16j, 1.07+0.04j, 1.07-0.04j, 0.99+0.05j, 0.99-0.05j,
0.80+0.09j, 0.80-0.09j, -1.22+0.10j, -1.22-0.10j, -0.53+0.13j, -0.53-0.13j, -1.53+0.05j, -1.53-0.05j, -1.45+0.03j, -1.45-0.03j, -0.62+0.11j,
-0.62-0.11j, -0.11+0.11j, -0.11-0.11j, 0.16+0.09j, 0.16-0.09j, 0.52+0.03j, 0.52-0.03j, 0.04+0.07j, 0.04-0.07j, 0.52+0.00j, 0.22+0.04j,
0.22-0.04j, -0.93+0.05j, -0.93-0.05j, -0.42+0.05j, -0.42-0.05j, 0.19+0.00j, -0.02+0.00j, -1.12+0.00j, -0.94+0.00j, -0.65+0.00j, -0.68+0.00j])
~~~
Thanks for your detailed explanation and informative reply. We follow the reviewer's suggestion and derive the Jacobian $\nabla_x f$ as follows.

In Eqn. (9) in our manuscript, the SMoE is formulated as

$$
x_{t+1} = x_t - \gamma\sum_{i=1}^K \text{softmax}(\text{TopK}(g_{i}(x_t))) [-u_i(x_t)] = x_t - \gamma f(x_t),
$$

where, again, $f(x) = \nabla_{x}F = \sum_{i=1}^K \text{softmax}(\text{TopK}(g_{i}(x))) [-u_i(x)]$ plays the role of a fixed point operator and $F$ is the objective function defined in Eqn. (3) in our manuscript. The Jacobian $\nabla_x f$ is then given by:

$$
\nabla_x f = \sum_{i=1}^K \nabla_x  \, \{\text{softmax}(\text{TopK}(g_{i}(x))) [-u_i(x)]\}.
$$

For each data point, we differentiate $f$ using automatic differentiation in PyTorch. We then compute the eigenvalues of the Jacobian at layers 1, 3 and 6 in the SMoE model, corresponding to their respective steps in our optimization algorithm. These eigenvalues are listed in the code snippet below. You can run our code to visualize those eigenvalues in a scatterplot. As can be seen from the resultant figures, the eigenvalues cluster around 0, $\pi$ and $-\pi$, indicating that while they are not exactly real, the eigenvalues are close to real. It is also interesting that in the last step (layer 6), there seem to be slightly less eigenvalues with almost purely imaginary parts.
The goal is to visualize the eigenvalues of $\nabla_x f$ (or, similarly, the eigenvalues of the Jacobian of the fixed point operator) at each optimization step, so you should visualize exactly that.  Note that you are abusing notation with $\nabla_x f$ as it seems to assume a potential function, so, really, you are looking at the eigenvalues of the Jacobian of the expert's aggregated update.  This is uniquely defined, and there should not be ambiguity in how you establish these eigenvalues.

As such, you should use the exact weighting of the expert's matrices, such that it equals $\nabla_x f$ (at some specific optimization step).  I believe this is weighting them according to the gating scores. Still, the key point is figuring out how to visualize exactly the spectrum of $\nabla_x f$.  You need to derive exactly what form the Jacobian $\nabla_x f$ -- which will involve gating values and expert matrices -- then visualize exactly that, and this will tell you exactly how to weigh the different matrices uniquely.  You should not just choose arbitrary averages of the matrices or arbitrary averages of eigenvalues of different terms.

Note that the eigenvalues will likely be different at each step of your fixed-point operator. It is more common to visualize the union of them or them separately at varying points in optimization or only at the end of optimization, which dictates asymptotic convergence rates.

This specific spectrum is useful, as it dictates the spectrum of the Jacobian of the fixed point operator, which bounds on convergence rates to the fixed point and, relatedly, the optimal acceleration schemes (like momentum and its variants).  See Theorem 1 from Negative momentum (https://arxiv.org/pdf/1807.04740), which just re-states the classic result of Prop 4.4.1 from Bertsekas 1999.

D. P. Bertsekas. Nonlinear programming. Athena
scientific Belmont, 1999.
On the point of averaging the matrices of each expert, we agree that this would be a better way to visualize the eigenvalues of $\nabla_{x}f$ than averaging the eigenvalues. However, as only 2 experts are chosen at each layer for each data point (i.e., our model implements top-2 SMoE), and further, a convex combination of their matrices are applied where the coefficients of this convex combination are the corresponding affinity (gate) scores, there are many possible combinations. How would you like to best visualize them? Would it make sense to you to take the 2 experts chosen the most frequently during training/inference time and plot the spectrum of their convex combination? 
    
For example, at layer $i$, $i=1,...,6$, we have 16 expert matrices $A_{i_1}, A_{i_2}, ..., A_{i_{16}}$, and the top 2 most frequently chosen matrices are $A_{i_j}$ and $A_{i_k}$. Then, for $\lambda \in [0,1]$ decided by the affinity (gate) scores, we plan to compute the convex combination $A_i = \lambda A_{i_j} + (1-\lambda) A_{i_k}$ and plot the spectrum of $A_i$ in each layer. 
    
We would be willing to provide you with the eigenvalues and a snippet of python code that can be easily run in a jupyter notebook for visualization.
Thanks for your prompt response. We appreciate your detailed and valuable suggestions to better visualize the eigenvalues of $\nabla_{x}f$. If you allow us to include an anonymous link that includes additional figures in our reply, we will be very happy to make the figures you suggested and present them in this discussion. We will also include those figures in our revision.
The goal is to visualize the eigenvalues of $\nabla_x f$, and perhaps view how it changes as "optimization progresses".  I think this corresponds to taking the eigenvalues of the matrix, which is an average of the matrices for each expert, and the "optimization step" is like the layer.  So the "averaging" would be done for the matrices for the experts themselves and not the eigenvalues if I understand correctly.  Further, plotting the eigenvalues' union would be a more typical visualization strategy than averaging if there are multiple distinct matrices, as getting rid of the conjugate structure is confusing to me. I think a union over the different layers may make sense, or as separate plots.

You might also find it easier to visualize as a heatmap than a scatterplot. If you wanted to see a more granular structure, you could play with different ways to plot it—for example, a different figure/color for each layer/expert. For example, in [Lorraine, 2022] Figure 9 shows the spectrum at the start and the end of training.
We are sorry for this misunderstanding. Please let us clarify Figure 1 here. By taking the average over the layers, we mean that in a medium SMoE model, there are 6 layers and 16 experts. Each of these will correspond to 1 real nxn matrix. We find the eigenvalues of each of these matrices, which are either real numbers or conjugate pairs. We find the modulus and phase of each eigenvalue and average over the 6 layers and 16 experts. Positive real eigenvalues in the spectrum result in phases of $0$ while negative real eigenvalues in the spectrum result in phases of $\pi$. Consequently, averaging these $0$ and $\pi$ phases with the phases of complex eigenvalues causes the plot to be skewed upwards. 
    
We have tried our best to reproduce Figure 7 in [Lorraine, 2022] in the context of SMoE. We would appreciate your suggestions to better make this figure. 
 
**Reference**
    
Lorraine, Jonathan P., David Acuna, Paul Vicol, and David Duvenaud. "Complex momentum for optimization in games." International Conference on Artificial Intelligence and Statistics. PMLR, 2022.
"For Figure 1 in our 1-page attached PDF, the spectrum is not symmetric due to the presence of purely real eigenvalues. We also take the average over the layers in the SMoE, causing the spectrum to skew upwards in our plot."

In your plot, there are complex eigenvalues -- i.e., eigenvalues with a non-zero imaginary part -- which do not have their conjugate plotted.  Why is this occurring?  It would not be caused by real eigenvalues, as every complex eigenvalue will still have its conjugate.  I also don't understand what you mean by taking the average over layers.  You are computing the eigenvalues of a real nxn matrix, which always satisfies the property of having complex eigenvalues occur in conjugate pairs.
We would like to thank the reviewer for your further feedback, and we appreciate your endorsement. 

We will include Figure 1 in the Section "Empirical Evidences for the Gradient Descent Analogy of (S)MoE" in our revision. Following your suggestion, we will revise our manuscript to verify that MoE is “close to doing gradient descent on an objective function in a spectral sense” and update our stability analysis of MomentumSMoE (Section 3 in our manuscript) accordingly. We will also discuss the extension from our momentum-based framework for SMoE to accelerating fixed-point iterations and include the experiments with the ComplexMomentumSMoE and NegativeMomentumSMoE in our revision.  
    
Furthermore, we agree with the reviewer that introducing a hyperparameter into the model usually improves its performance and that a thorough comparison will help further validate our experimental results. As you mentioned, robust momentum is not our main contribution, so we leave this for future work. However, we would like to point out that when tuning the hyperparameter in RobustMomentumSMoE, we are tuning the model on clean ImageNet-1K data, where there is a slight improvement. While on perturbed datasets, such as ImageNet-A/R/C, there is a significant increase in accuracy. These results lead us to believe that the model's enhanced performance is not entirely due to the introduction of an extra parameter, but rather the design of robust momentum itself. 
    
In addition, further investigation of acceleration methods that consider different spectrums is an interesting direction that we will continue to explore. We believe that our momentum-based framework can be extended to these methods, as well as their theories, and will benefit from them. 
    
For Figure 1 in our 1-page attached PDF, the spectrum is not symmetric due to the presence of purely real eigenvalues. We also take the average over the layers in the SMoE, causing the spectrum to skew upwards in our plot.
    
Once again, we appreciate your suggestions and feedback which help enhance the quality of our paper and introduce exciting new directions for further development of our MomentumSMoE.
I appreciate the author's response and results, which ameliorate my main concerns.  As such, I am revising the rating to be higher.  Here are some elaborations on the response:

Figure 1 in the response really strengthens the paper, and it should be referenced in the Section “Empirical Evidence for the Gradient Descent Analogy of (S)MoE.”  However, I think it's quite important that the writing clearly emphasizes that the MoE is not doing gradient descent on a single objective function since you observed complex eigenvalues.  Instead, the MoE is “close to doing gradient descent on an objective function”, in a spectral sense.  Further, all subsequent theories should be noted as a simplified heuristic justification for your method, as the theory requires assuming a conservative, real-valued vector field.  I believe the theory could be generalized to work with complex eigenvalues, and I do not think it is required from this work to generalize the theory. Still, the gap between theory and practice should be easily documented for readers.

I believe the comparison with other fixed-point acceleration schemes also enhances the paper.

Also, regarding robust momentum, introducing an extra hyperparameter -- even randomly -- often leads to some improvement from tuning the new parameter.  What we need to do to show this is useful is some exploration showing that doing your search in this new 3-parameter space is better than optimizing the 2-parameter space – ex., finding better values when doing a grid search of N queries in 3 dimensions instead of 2.  Or, you could compare to other ways to introduce a third hyperparameter, and show you are better.  However, I don’t think this is non-essential for publication, as the robust momentum is not the key contribution.  More, you note that various momentum flavors can be projected to this setting.

It seems infeasible for NeurIPS, but it would improve the paper to rephrase the momentum theory to work with spectrums that are bound to be “near” the real axis in some sense. I think it's likely you could reuse theory from papers on fixed-point acceleration.

For Figure 1, shouldn’t the spectrum be symmetric, because if I have an eigenvalue at “x + iy” then I must have an eigenvalue at “x - iy” because the eigenvalues come in conjugate pairs.  Since your plot is not symmetric, what is being shown? Are you only plotting the positive roots?

Also, this does not affect my review, but I think is an interesting avenue for future exploration.  For your Figure 1 with the spectrum exploration, it is interesting that the spectrum is almost all positive and closely clustered to the real axis. Why is there one eigenvalue with a negative real part?  A better understanding of the shape of spectrums you encounter allows you to select optimal acceleration schemes.  Does this spectral shape hold for various MoE models?  Does the spectral shape change during training?
Also, there is a single eigenspace with a negative real part—this seems notable. Does this eigenspace correspond to anything interpretable? In what eigenspaces do the parameters accumulate as the iterations progress?  Answering these may lead to insights, allowing you to design better fixed-point acceleration schemes.