Layerwise Recurrent Router for Mixture-of-Experts

Responses to reviewer oYLa

Thanks for your review. To further validate the effectiveness of our method, we supplement experiments with larger scales, i.e., larger models and more training tokens. We **train models up to 43B total parameters with 6.9B activated parameters and 400B tokens**. We evaluate our trained model on MMLU, GSM8K, Hellaswag, and Humaneval. We observe **consistent performance gains** on these tasks compared with the standard MoE (MMLU, 65.8 v.s. 64.4, GSM8K 67.4 v.s. 65.1, human eval 37.8 v.s. 35.4). Although we do not scale tokens to 1T level, the model trained on 400B high-quality data has already achieved comparatively high performance. We humbly believe that experiments with this scale could validate the effectiveness of our method. For more detailed results, please refer to the general comment, common concern 1. Thanks for the mentioned related work. We have cited it in our revised paper. Please let us know if there is anything we could do to let you recommend our paper further!

Responses to reviewer Cw7i

> Current performance improvement appears marginal; val and test losses alone do not demonstrate statistically significant gains over baselines. Benchmarking on standard language modeling tasks would provide more practical context for the model's effectiveness. To further validate the effectiveness of our method, we supplement experiments with larger scales, i.e., larger models and more training tokens. We train models up to **43.1B total parameters and 400B tokens**. We evaluate our trained model on MMLU, GSM8K, Hellaswag, and Humaneval. We observe **significant performance gains on these tasks** compared with the baseline (MMLU, 65.8 v.s. 64.4, GSM8K 67.4 v.s. 65.1, human eval 37.8 v.s. 35.4). We humbly believe that experiments with this scale could validate the effectiveness of our method. For more detailed results, please refer to the general comment, common concern 1. > Analysis of gradient flow could be more rigorous; while recurrent gradient flow is highlighted as crucial in Section 5, more sophisticated gradient tracking techniques would strengthen this claim. Thanks for your suggestion! We track the gradient norm of different routers at different training steps. We empirically found that the introduced recurrent gradient flow could partially prevent the router from being dominated by the load balancing loss through training, thus bringing improvements. We think this extra analysis could strengthen our claim. Please refer to the general comment, Common Concern 2, for more detailed results. > The additional computation stage introduced by the GRU increases complexity, yet the impact on training efficiency is not thoroughly analyzed across different model sizes. We reported the training time and memory usage in Table 1 and Table 2 in our submission. The training speed and the memory usage are approximately the same as the SMoE baseline, because the bottleneck remains the calculation of other parameters in the model. > Clarify the architecture of the layer-specific projectors used in Eq. 4; are there structural variations across layers, or have alternative projector architectures been tested for potential performance improvements? Thanks for your question! We have improved the clarity of this part in our revised paper. To be clear, each projector is just one single linear layer. By mentioning layer-specific, we are emphasising that the parameters of projectors are not shared across layers. We empirically validate that employing layer-specific/unshared projectors is beneficial. This may be because the distributions and norms of hidden states in different layers shift. We have not introduced more sophisticated projectors, which we believe is promising but is not the main focus of our paper. > In Table 1, improvement over baselines is minor, yet gains are more noticeable in Table 2 (billion-level settings). Could the authors explain the differing results, specifically regarding training sizes and conditions? The difference between Table 1 and Table 2 is mainly because of metrics. Table 1 uses BPC and PPL, whereas Table 2 uses accuracy. A similar trend is also observed in our supplemented experiments in the large-scale setting, presented in the general comment, common concern 1. The difference in the ppl is usually small, but the improvement for downstream tasks is usually more significant. > Column names in Table 8 appear inconsistent with descriptions in the text; could these be corrected or clarified? Thanks for your question! From your description, we guess that the inconsistency may be: in the text description, we formally define the **Inner Balance** with the ratio top1 / top-k; and in the caption of Fig. 8, we mentioned the top1 / top2. This is because we set k=2 in all our experiments. We have revised this part to make it more clear!

Responses to Reviewer bumc

> Larger MoE models Thanks for your suggestion. In our submission, we showed the results not only on the 36M small model, but also the results with approx. 1B model trained with 20B / 40B tokens. To further validate the effectiveness of our method, we supplement experiments with larger scales, i.e., larger models and more training tokens. We train models up to **43.1B total parameters and 400B tokens**. We evaluate our trained model on MMLU, GSM8K, Hellaswag, and Humaneval. We observe **significant performance gain on these tasks compared with the baseline** (MMLU, 65.8 v.s. 64.4, GSM8K 67.4 v.s. 65.1, human eval 37.8 v.s. 35.4). We humbly believe that experiments with this scale could validate the effectiveness of our method. For more detailed results, please refer to the general comment, common concern 1.

Reminder for discussion

Dear reviewer bumc, We apologize for reaching out with another reminder of the discussion. We want to highlight that we have provided additional experimental results following your suggestions. Specifically, we pre-trained models with up to **43.1B total parameters** for up to **400B tokens** and evaluated the models on MMLU, gsm8k, HumanEval, etc. We observe significant improvements in these tasks. We would appreciate your thoughts on whether these updates satisfactorily address your concerns. Best, Authors

Responses to Reviewer rPzv

> More complete evaluation. We have provided the results of training larger models on more tokens (15B models on 120B tokens, 15B models on 400B tokens, and 43B models on 400B tokens). Our proposed RMoE network demonstrates substantial improvements regarding training perplexity and downstream task performance. When training the **43B MoE with 400B tokens, RMoE achieves 1.35 performance gains on MMLU, 2.27 on GSM8K, and 2.43 on Humaneval compared with the standard SMoE**. The experimental details are explained as presented in the General Comment Common Concern 1. > A more direct analysis Please refer to our general comment, common concern 2. We analyse the grad norm and the drop_ratio of different routers through the training. We validate that the extra recurrent gradient effectively reduces the influence of the load balancing loss on router training, therefore bringing improvements. > The paper could benefit from a deeper comparative analysis with recent MoE models that also use recurrence or other advanced routing strategies (such as XMoE, HyperMoE etc). This would help clarify how RRMoE’s benefits compare to similar innovations in this rapidly evolving field. Thanks for your suggestion. To the best of our knowledge, we are the first to introduce cross-layer recurrence into the MoE architecture and provide an analysis of gradient flow. We respectfully believe that the ablations in our submission and additional experimental results in our rebuttal message could well explain how RMoE benefits: an extra gradient pathway! > The paper emphasizes the importance of gradient flow through the GRU-based recurrence in RMoE but does not mention whether Truncated Backpropagation Through Time (BPTT) or any specific method was used to manage gradients across layers. Did the authors consider using truncated BPTT for controlling gradient flow through the recurrent router? Including details on gradient management would help clarify if additional methods, like truncation, could further optimize training. Thanks for your question. We have not used *truncated* BPTT or any other specific method to manage the gradients across layers. The GRU is introduced to capture the dependency of routing decisions between different layers, and the largest model we train (43.1 B parameters) has 32 layers, which is still a relatively short sequence for a standard GRU model. > Would it be possible to conduct a more direct analysis, such as by visualizing or measuring gradient norms through layers, to further validate the specific impact of recurrent gradients? Thanks for your suggestion. We have provided a more direct analysis of the recurrent gradient's impact. We investigate how routers' gradient norms and drop ratio vary during training. We validated that the extra recurrent gradient path could help prevent the router from overfitting the load balancing loss during the early training stage. Please refer to our general comment, common concern 2, for more detailed results. > Could the authors provide insights or quantitative analysis on how the recurrence depth impacts the RMoE’s routing decisions? We humbly believe that the recurrence depth does not have much impact on GRU to capture dependencies of routing decisions across layers. This is because the models are usually not that deep. Even our trained 43.1B models only have 32 layers, which is still a relatively shorter sequence for the GRU model. Moreover, we think the recurrence depth mainly helps the gradient backward propagation. In section 5, we find that only using recurrency in the forward process will not improve the performance. > The current evaluation omits challenging benchmarks like MMLU and GSM8K, mentioning that the models are still under-trained. Can the authors comment on whether they plan to scale the training or incorporate task-specific fine-tuning to gauge RMoE’s performance on these harder tasks? Additional results on these benchmarks could help better gauge RMoE’s robustness and adaptability in real-world scenarios. Thanks for your questions. To provide further empirical results, we train models up to **43.1B total parameters and 400B tokens**. We evaluate our trained model on MMLU, GSM8K, Hellaswag, and Humaneval. We observe significant performance gains on these tasks compared with the baseline (MMLU, 65.8 v.s. 64.4, GSM8K 67.4 v.s. 65.1, human eval 37.8 v.s. 35.4). We humbly believe that experiments with this scale could validate the effectiveness of our method. For more detailed results, please refer to the general comment, common concern 1.

Reminder for discussion

Dear reviewer rPzv, We apologize for reaching out with another reminder of the discussion. We want to highlight that we have made significant revisions to our paper following your suggestions, including more complete, large-scale evaluations, a more direct analysis of the recurrent gradient, and other clarification responses to your questions. We wonder if it is possible for you to let us know if our responses could address your concerns. We would appreciate your thoughts on whether these updates satisfactorily address the theme. Best, Authors

General Comment Part 2 / 4

| Setting: **15B activate 2.54B param. Training 400B tokens** | Hellaswag | MMLU | GSM8K | training ppl | | ----------------------------------------------------------- | --------- | --------- | --------- | ------------ | | 50k steps (200B tokens) | | | | | | linear router | 69.48 | 49.96 | 33.21 | 7.718 | | MLP router | 69.76 | 50.27 | 31.77 | 7.736 | | RNN router (Ours) | **70** | **52.21** | **32.98** | **7.608** | | 70k steps (280B tokens) | | | | | | linear router | 72.4 | 54.66 | 42.61 | 6.477 | | MLP router | 72.62 | 55.33 | 38.51 | 6.502 | | RNN router (Ours) | **73.18** | **56.06** | **44.35** | **6.4** | | 100k steps (400B tokens) | | | | | | linear router | 76.39 | 59.54 | 52.16 | 5.685 | | MLP router | 76.09 | 59.96 | 51.71 | 5.709 | | RNN router (Ours) | **76.72** | **60.6** | **52.99** | **5.62** | | Setting: **43.1B activate 6.9B param. Training 400B tokens** | Hellaswag | MMLU | GSM8K | Humaneval | training ppl | | ------------------------------------------------------------ | --------- | --------- | --------- | --------- | ------------ | | 50k steps (200B tokens) | | | | | | | linear router | 74.47 | 56.1 | 47.01 | 28.05 | 6.184 | | RNN router (Ours) | **74.96** | **57.21** | **48.6** | **30.49** | **6.104** | | 70k steps (280B tokens) | | | | | | | linear router | 77.13 | 61.37 | 56.1 | 28.66 | 5.585 | | RNN router (Ours) | **77.95** | **62.42** | **56.86** | **31.71** | **5.521** | | 90k steps (360B tokens) | | | | | | | linear router | 79.2 | 64.34 | 64.97 | 34.15 | 5.101 | | RNN router (Ours) | **79.71** | **65.45** | **65.96** | **37.2** | **5.047** | | 100k steps (400B tokens) | | | | | | | linear router | 79.62 | 64.41 | 65.13 | 35.37 | 4.995 | | RNN router (Ours) | **80.29** | **65.76** | **67.4** | **37.80** | **4.936** | The tables validate that **our proposed method is still effective when training larger MoE language models with more tokens**. The empirical findings are consistent with those of our small-scale settings. 1. Our proposed RNN router outperforms the linear and MLP routers in all settings, training steps, and tasks, demonstrating significant improvements in practical scenarios. 2. Although the MLP router adds more parameters to the original linear router, it still underperforms the simple linear router on most tasks. This indicates that simply involving more computation in the router stage cannot boost the MoE’s performance. Our results in Table 1 of the paper support this observation.