MomentumSMoE: Integrating Momentum into Sparse Mixture of Experts

Q1. Unfounded connection between the SMoE and gradient descent. Connection to accelerating fixed-point iterations. The authors should work with complex eigenvalues [Azizian et. al.] using negative (Gidel et. al.) or complex momentum (Lorraine et. al.). Empirical (or theoretical) investigation and an analysis of the eigenvalues of $\nabla_x f$ is needed.

Answer: We respectfully disagree with the reviewer’s comment that the connection between SMoE and gradient descent (GD) is largely unfounded. In Section 2.3, from line 128 to line 144, we discuss the empirical evidences (Fig. 2, 3) for connection between SMoE and GD.

GD updates can be considered as fixed-point iteration: at optimality where the gradient is 0, GD finds a fixed-point solution. Thus, techniques in accelerating fixed-point iterations can be used to find the equilibrium point of SMoE. Following the reviewer's suggestion, we have conducted additional experiments that incorporate negative momentum (Gidel, 2019) and complex momentum (Lorraine, 2022) into SMoE using our momentum-based design framework. Table 2 in the attached PDF compares the PPL of NegativeMomentumSMoE and ComplexMomentumSMoE with MomentumSMoE, AdamSMoE, and SMoE on the clean/attacked WikiText-103. Our AdamSMoE achieves the best PPL. ComplexMomentumSMoE obtains slightly better PPL than MomentumSMoE (positive momentum coefficient), while both of these methods significantly outperform NegativeMomentumSMoE and the SMoE baseline.

Furthermore, when analyzing the eigenvalues of $\nabla_x f$, where $f$ is the output of an SMoE layer in the trained SMoE above, we observe that the mean absolute real and imaginary parts of the eigenvalues are 4.06 and 0.34, averaged over all SMoE layers in the model, respectively. (Gidel, 2019) and (Lorraine, 2022) suggest that positive momentum converges when the eigenvalues are real, while negative momentum has better convergence when the eigenvalues have large imaginary parts. Also, the convergence of the complex momentum method is robust to the wider range of complex eigenvalues, including purely real and imaginary ones. Since the imaginary part of the eigenvalues of $\nabla_x f$ is small compared to their real part (0.34 vs. 4.06), NegativeMomentumSMoE should not work well, and MomentumSMoE with positive momentum should be a better design choice. Since ComplexMomentumSMoE can handle a wider range of complex eigenvalues, it helps improve the performance of MomentumSMoE. However, due to the small imaginary part of eigenvalues of $\nabla_x f$, the improvement is small. These explain our results in Table 2 in the attached PDF and we also provide a plot of the spectrum in Figure 1.

We would like to thank the reviewer for your suggestion on negative and complex momentum. The significant improvement of MomemtunSMoE, AdamSMoE, and ComplexMomentumSMoE over the SMoE baseline further verifies the power and promise of the momentum-based framework for designing SMoE. Moreover, given our framework, analyzing the eigenvalues of $\nabla_x f$ offers a principled way to select the right momentum method for SMoE design.

Q2. Methods like RobustSMoE seem to rely on knowing the spectrum to set various parameters, which we won’t have access in real settings.

Validating/Relaxing assumptions about the fixed point operator's spectrum and the Jacobian's conservativeness

Answer: The motivation behind the design of complex momentum is to work over a large range of eigenvalues, as it is common that, in practice, we do not know the spectrum [Lorraine, 2022]. Similarly, for MomentumSMoE and AdamSMoE, we need to tune the hyperparameters, i.e., $\mu$, $\gamma$, as well as $\beta$ in AdamSMoE, in order to achieve convergence.

In contrast, the Robust Momentum Method in [Cyrus, 2018] that we use to develop Robust MomentumSMoE does not require the spectrum information, and the convergence of the algorithm does not depend on the spectrum values. This Robust Momentum Method can be considered as a Lur’e feedback control system, and the hyperparameters $\gamma$, $\mu$, and $\alpha$ (see Eqn. (18), our manuscript) are designed to push the stability boundary into the negative real axis. This allows an additional margin for the stability conditions of the system to hold.

Q3. The novelty of theoretical results?

Answer: As mentioned in line 156 in our manuscript, we are inspired by the techniques in [Qian, 1999] and adapt these techniques to prove the convergence and stability results of Momentum SMoE (Proposition 1 and Corollary 1). The techniques we use are not new, and other techniques, such as those in [Azizian, 2020] or [Wilson, 2018], can be leveraged to prove similar results. Our objective is to provide theoretical guarantees on the convergence and stability of MomentumSMoE to prove that MomentumSMoE is more stable than the SMoE baseline, which justifies the empirical advantages of MomentumSMoE over SMoE shown in our manuscript.

Q4. More general acceleration tools than momentum?

Answer: Momentum methods, such as Nesterov’s accelerated gradient method, are simple approximations of the proximal point method [Ahn, 2022]. For future work, we can employ the proximal point method to develop better SMoE models. In our manuscript, we also explore the sharpness-aware minimization (SAM) [Roulet, 2017] approach to enhance the robustness of SMoE (lines 805 - 821 in Appendix E.4).

References

O’donoghue, B, et al. Adaptive restart for accelerated gradient schemes. 2015.

Lucas, J., et al. Aggregated momentum. 2018.

Cyrus, S., et al. A robust accelerated optimization algorithm for strongly convex functions. 2018

Qian, N. On the momentum term in gradient descent learning algorithms. 1999.

Wilson, A. C., et al. A Lyapunov analysis of accelerated methods in optimization. 2021.

Ahn, K., et al.Understanding nesterov's acceleration via proximal point method. 2022.

Roulet, V., et al. Sharpness, restart and acceleration. 2017.

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.

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.

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.

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.

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.

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

Q1. Why the MomentumV-MoE and Robust MomentumV-MoE have marginal gains on clean IN-1K data, is there any in-depth analysis available on this?

Answer: V-MoE's result reported in Table 2 in our manuscript is among the state-of-the-art results on clean IN-1k data for the models that have around 45M effective parameters without pretraining. To further improve this result is challenging, and the almost 0.43% improvement of our MomentumV-MoE vs. the V-MoE baseline is already nontrivial. Notice that both our MomentumV-MoE and Robust MomentumV-MoE significantly improve over the V-MoE baseline on robustness benchmarks, such as IN-R, IN-A, and IN-C, as shown in Table 2.

Q2. In the ImageNet-1K Object Recognition experiment, why was the popular top-5 accuracy metric not used, as it was in the Soft Mixture of Experts experiment?

Answer: Thanks for your comment. We provide the top-5 accuracy for ImageNet-1K object recognition task in Table 1 in the attached PDF.

Q3. As stated in the weaknesses, the authors could provide pseudocode to better clarify their proposal.

Answer: Thanks for your suggestion. We include the pseudocode for our MomentumSMoE, AdamSMoE, and Robust MomentumSMoE below.

Hyperparameters: mu

def MomentumSMoE(x, momentum):
	momentum = - SMoE(x) + mu * momentum 
	x = x + gamma * momentum
	return x

Hyperparameters: mu, beta, eps = 1e-8

def AdamSMoE(x, gradient, squared_gradient):
	gradient = - (1 - mu) * SMoE(x) + mu * momentum 
	squared_gradient = beta * squared_gradient + (1 - beta) * SMoE(x) ** 2
	x = x + gamma / (torch.sqrt(squared_gradient) + eps) * gradient - k * x
	return x

Hyperparameters: p, L, m

def RobustMomentumSMoE(x, momentum):
	k = L / m
	gamma = k * ((1 - p) ** 2) * (1 + p) / L
	mu = k * p ** 3 / (k - 1)
	alpha = p ** 3 / ((k - 1) * ((1 - p) ** 2) * (1 + p))
	y = x + alpha * gamma * momentum
	momentum = - SMoE(y) + mu * momentum 
	x = x + gamma * momentum 
	return x

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.

Q1. Formulating SMoE as a multi-objective optimization problem is doubtful to me.

Answer: We believe there is a misunderstanding of our formulation of SMoE as a multi-objective optimization problem. Please allow us to clear this misunderstanding by clarifying the role of expert networks in our multi-objective optimization (MOO) framework for Sparse Mixture of Experts (SMoE). We rewrite the minimization problem in Eqn. (3) as follows:

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

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

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

Step 2 - fix $x$ and minimize $F$ with respect to $\theta$

where $\theta^{i}$ lies in the parameter space $\Theta^{i}$.

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

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

Q2. Why $||f(x)||$ as the key metrics to measure the efficacy of (S)MoE?

Please explain more of line 140.

Answer: In Fig. 2 and 3 in our manuscript, $||f(x)||$ is not used as a key metric to measure the efficacy of SMoE/MoE. Instead, the results in Fig. 2 and 3 are to justify the connection between SMoE/MoE and gradient descent. In particular, from Eqn. (8) and (9), the output of (S)MoE corresponds to $-f(x_t)$, where $f(x_t) = \nabla_xF_(x_t)$ and is the gradient of the objective function $F$ with respect to $x$ at $x_t$. If (S)MoE indeed performs (stochastic) gradient descent in its forward pass, then the norm of the (S)MoE output must decrease when t (i.e., layer index) increases since the gradient norm $\||f(x_t)\||$ decreases when gradient descent updates are applied. Fig. 3 confirms this by showing that the norm of the (S)MoE output decreases over layers in a 6-layer (S)MoE model trained on the WikiText-103 language modeling task. At the last layer, the norm increases might be due to overshooting, a common phenomenon that can occur when using gradient descent.

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

Q3. Grammar issues: We have addressed the grammar issues in our revision.

Q4. Insufficient discussion of computation overhead.

Answer: The reviewer might have overlooked Appendix E.5 of our manuscript. Table 11 in E.5 reports the training/inference runtime, training/inference memory usage, and the number of parameters of our MomentumSMoE/AdamSMoE vs. the SMoE baseline. Our MomentumSMoE/AdamSMoE have the same training/inference memory usage and number of parameters as SMoE. Furthermore, our momentum models are on par with the SMoE baseline in terms of training/inference runtime. Compared to the SMoE baseline, our MomentumSMoE (see Fig. 1) only needs an additional momentum state $p$, which is accumulated from layer to layer. Thus, MomentumSMoE only uses an additional tensor shared between layers to store this momentum state $p$. Similarly, AdamSMoE only needs two additional tensors shared between layers to store $p$ and $m$. The memory overhead for introducing these additional tensors is minimal, thus explaining why memory increase and throughput decrease in our MomentumSMoE/AdamSMoE compared to the SMoE baseline are negligible.

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.

Q1. This method has little effect on models with few layers.

Answer: A momentum-based approach like MomentumSMoE needs more layers to show its advantages, just like a heavy-ball momentum, or Adam needs a couple of iterations to start showing its faster convergence compared to gradient descent. Also, models with few layers are not common in practice due to their worse performance compared to models with more layers. For example, as can be seen in Table 7 in our manuscript, the SMoE-small baseline has much lower validation and test perplexity (PPL) compared to SMoE-medium and large baselines, 84.26/84.81 PPL vs. 33.76/35.55 and 29.31/30.33 PPL, respectively. The poor PPL of the SMoE-small baseline renders its performance on the WikiText-103 language modeling tasks negligible.

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

Answer: Thanks for your comment. During the rebuttal, we have tried our best to conduct an additional experiment on a Sparse Mixture of Experts (SMoE) backbone with more than 1B parameters. However, due to the shortage of computing resources and the short time of the rebuttal, our trainings have not finished yet. We will include these results of our momentum-based SMoE on a larger SMoE backbone in the revised manuscript and during the discussion period if possible.

Q3. From a theoretical standpoint, developing a framework to explain the enhanced robustness of MomentumSMoE would be interesting.

Answer: Thanks for your suggestion. We agree with the reviewer that developing a framework to explain the enhanced robustness of MomentumSMoE would be interesting. The approach to theoretically proving the robustness of MomentumSMoE is to consider a minimizer $x^*$ of the objective function $F(x)$ in Eqn. (3) in our manuscript, which is found by heavy-ball updates when starting at the initial point $x_0$. Then, we prove that by tuning the momentum $\mu$ and step size $\gamma$, the heavy-ball iterations will find a point $x_k$ such that $\||x_k - x^*\|| \le c\rho^{k}$ for a constant c, $\rho \in [0, 1)$, and $k \ge 1$. This result implies that the output of MomentumSMoE is close to $x^{*}$ as long as the input data is close to $x_0$, thus verifying the robustness of MomentumSMoE. This proof can be extended from the proof of Proposition 1 and Corollary 1 in our manuscript. We will include the detailed proposition and its proof for the robustness of MomentumSMoE in our revision.

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.