Sparse Mixture of Experts (SMoE) has become the key to unlocking unparalleled scalability in deep learning. SMoE has the potential to exponentially increase 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 demonstrate 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.
Paper
Similar papers
Peer review
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.
Strengths
To the best of my knowledge, attempting to accelerate the fixed point iteration in SMoE is an original idea. It seems like there is comprehensive empirical evidence for the method, but I am not an expert on metrics for the SMoE, and will have to rely on other reviews to be confident in this strength. The paper is fairly clear, with well-organized sections and figures.
Weaknesses
My largest negative for this paper is the largely unfounded connection between the SMoE and gradient descent. If the authors had made a connection to accelerating fixed-point iterations in general, I would want to accept this paper. Essentially, the authors are assuming that $\nabla_x f$ has strictly real eigenvalues when they should just work with truly, potentially complex, eigenvalues, ex., using tools as in Azizian et. al. For example, when performing this analysis, various other acceleration schemes are often better, like negative momentum (Gidel et. al.) or complex momentum (Lorraine et. al.). I would be curious to see some empirical investigation (or theoretical) or what the eigenvalues of $\nabla_x f$ are – ex., as in Figure 7 of https://arxiv.org/pdf/2102.08431 -- to validate any theoretical claims about what acceleration schemes should be used. 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
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?
Rating
7
Confidence
4
Soundness
3
Presentation
3
Contribution
4
Limitations
The limitations are discussed briefly, but a delineated section elaborating on all the limitations would be valuable.
Thanks for your endorsement!
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.
"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.
Clarification on Figure 1 in Our Attached PDF
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.
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.
Reply to Reviewer GFpG
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.
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$ (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.
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.
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 endorsement!
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.
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.
Strengths
1. The idea of integrating momentum into sparse mixture of experts is interesting. 2. Both the theoretical proof and extensive empirical results are provided to demonstrate that the proposed MomentumSMoE is more stable and robust than SmoE; the experimental results are appealing. 3. The code is provided.
Weaknesses
The pseudocode may be provided to better illustrate the implementation of the proposal.
Questions
1. Why the MomentumV-MoE and Robust MomentumV-MoE have marginal gains on clean IN-1K data, is there any in-depth analysis available on this? 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.
Rating
7
Confidence
3
Soundness
3
Presentation
3
Contribution
3
Limitations
The authors adequately addressed the limitations.
Any Questions from Reviewer puL2 on Our Rebuttal?
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.
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.
Strengths
(1) The integration of momentum into SMoE is a non-trivial innovation that addresses instability and inefficiency issues in existing models. (2) The paper provides convincing empirical evidence showing the effectiveness of MomentumSMoE across multiple benchmarks. (3) The proposed method's compatibility with other momentum-based optimizers, like Adam, suggests it can be broadly applied to various SMoE architectures.
Weaknesses
(1) Formulating SMoE as a multi-objective optimization problem is doubtful to me. Every expert network is continually changing during the model training, which makes each objective nonstatic, which violates the basic assumption of multi-objective optimization, whose objectives should be very clear and stable. (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
(1) Please explain more of line 140 ("Thus, it is expected that these two terms learn to reduce ...).
Rating
5
Confidence
4
Soundness
2
Presentation
3
Contribution
3
Limitations
N/A
Any Questions from Reviewer oJG3 on Our Rebuttal?
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.
After Rebuttal
Thanks to the detailed response from authors. Most of my concerns are well addressed. I'd like to increase the score to 5.
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.
Strengths
1. The paper tackles a critical issue in the training of SMoE models. 2. The proposed method is generalizable and can be applied to various SMoE models such as V-MoE and GLaM. 3. Experimental results demonstrate that this method is more stable than SMoE during the training process.
Weaknesses
1. This method has little effect on models with few layers. 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
please refer to weaknesses
Rating
6
Confidence
4
Soundness
3
Presentation
3
Contribution
3
Limitations
The authors have adequately discussed the limitations.
Any Questions from Reviewer qEGD on Our Rebuttal?
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.
Response to Rebuttal
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.
Decision
Accept (poster)