$\textit{Read-ME}$: Refactorizing LLMs as Router-Decoupled Mixture of Experts with System Co-Design

Thanks for your time and effort for the review. We answered your questions as follows.

[W1 - Effectiveness of Read-ME] Thank you for your question.

• We would like to first clarify that our method significantly outperforms Sheared-Llama. For example, Sheared-Llama nearly performs at random guess levels on the MMLU task (achieving only 26.4% accuracy on a 4-class classification), whereas Read-ME achieves a much higher accuracy of 38.9%.
• Second, our method is also more training-efficient, using only 1 billion tokens compared to Sheared-Llama's 50 billion tokens (x50 times higher than ours). To further back up our claim, we perform the Sheared-llama continued tuning using the same amount of tokens as ours. We show the downstream task performance in Table A. By using 1B tokens for tuning, the performance of Sheared-llama reduces to 48.2%, which is significantly lower than our method (55.5%). Note that we use the publicly released Sheared-llama pruning checkpoint [1] as the initialization.
• Additionally, our method aims to convert a dense pre-trained model into a MoE for efficient inference, and such a compression method is typically not lossless. However, our method achieves a superior accuracy-cost trade-off compared to all baseline compression methods. To validate our approach, we plot the MMLU performance of our method alongside a large number of baseline methods in Figure 4. This data is also presented in tabular form in Table B. Our method outperforms all baseline models/methods of similar sizes.

Table A: Our method is more training-efficient, using only 1 billion tokens compared to Sheared-Llama's 50 billion tokens (x50 times higher than ours). If running sheared-llama with only 1 billion tokens, the performance drops significantly.

Table B: Evaluation of Read-ME on MMLU benchmark, compared with other representative open-source models and compression techniques.

[1] https://huggingface.co/princeton-nlp/Sheared-LLaMA-2.7B-Pruned

[W2 - Perplexity Comparison] Thank you for your suggestion. We tested the perplexity on Wikipedia, with the results presented in Table C. We set the sequence length to 1024, which is within the sequence limit of all models. Table C demonstrates that our method shows negligible performance degradation (from 3.91 to 3.94) compared to Llama-2, the pre-trained dense model. Furthermore, our method achieves better performance compared to most of the baseline methods of similar sizes, including Sheared-Llama, and Pythia. We will add the column to Table 1 in the future version, and we have also provided a tentative Table 1 in the uploaded PDF.

Table C: Comparison on Wikipedia perplexity.

[W1 - Additional experimental results] Thanks for the suggestion! To validate that our method remains effective in other scenarios, we use the Mistral model as the pre-trained dense model, and convert it to the MoE structure, with the proposed method. The task is challenging because we do not have prior knowledge on the Mistral original training data, and our experiment in Table A shows that our method remains effective without the prior knowledge of the original training domain.

Table A: Ablation study on Mistral[1] pre-trained model.

[1] Mistral 7b.

[W2 - Computational cost of Auto-regressive Routers]

Thanks for the suggestion!

For a detailed analysis, we added: (1) FLOPs comparison, (2) latency, and (3) latency breakdown with a larger batch size (high-throughput scenarios) of a Traditional Router (TR) and an Autoregressive Router (AR). To focus solely on the router’s impact on latency, we controlled other variables (e.g., the number of activated parameters) to be the same.

(1) flops comparison

Table B. Flops comparison between Traditional router and Auto-regressive router

(2) latency [ms]

Table C. Latency comparison between Traditional router and Auto-regressive router

(3) latency breakdown [%]

Table D. Latency breakdown comparison between Traditional router and Auto-regressive router

Note that the computational cost of both the traditional router and the autoregressive router is theoretically linear to batch size. Therefore, when the batch size is high (in high-throughput scenarios), the cost increases linearly. In both cases, the computation can be parallelized, so this remains negligible in end-to-end latency even in high-throughput scenarios. In fact, we would like to clarify that the bottleneck in high-throughput scenarios is actually the expert layers, as seen in table D – Expert/MLP row. This issue can be addressed by the methods discussed in Section 4. Traditional layerwise routers do not allow for efficient system design, which underscores the need for a careful co-design of routers.

[W3 - Reproducibility and Pseudocode]

In Appendix A, we provide the pseudocode for the batching algorithm and pre-gating. In summary, at each scheduling step, we find the expert with the most requests and select that expert for the current step. We then check whether the scheduled tokens exceed the maximum token length or the maximum number of requests that can fit. This process is repeated until no more requests can be scheduled. We will release the code publicly and refine the text to improve understanding as well.

Thanks for all the interesting questions. Please see below.

[W1 - MoEs Motivation] We validated that MoE achieves better cost-accuracy trade-off than small dense models acquired by pruning, and provided the results in Appendix C.1.

We mentioned that prior compression efforts focus on converting large dense pre-trained models into smaller dense models. However, we argue that a smaller MoE model, with fewer activation parameters, is a better target architecture. To ensure a fair comparison, we (1) create a 4.7B parameter-dense model matching a single expert network size, and (2) fine-tune it for the same number of steps. Table 5 shows refactorizing the pre-trained model into an MoE structure, rather than a smaller dense variant, leads to significant performance improvement.

Table 5: By adopting an MoE as the target structure instead of a dense model, our model achieves better overall performance.

[W2 - Figure 2 Details] For Figure 2, we used the Arxiv dataset, a subset of Red-pajama. The observation holds with Wikipedia and Github subsets as well. Please see the uploaded PDF for visualization results; we will add these in a future version.

The visualized high mutual information between two adjacent layers’ expert selection is sufficient to support our claim. Since $H(S_1, …, S_L) \le \sum_l H(S_l) - \sum_l I(S_{l+1}; S_l)$, high layer-wise mutual information implies low-entropy (deterministic) path selection.

Regarding visualizing end-to-end routing paths, firstly we would like to mention that since the model has 32 layers and each layer performs Top-2 selection out of 8 experts, there will be ${8 \choose 2}^{32}$ (approximately $2 \times 10^{46}$) possible paths.

Instead, we calculated routing statistics for 180k tokens, finding only 8.4k paths out of $2\times 10^{46}$ possible paths activated at least once. The top 20 paths are selected by 2805 tokens, and the top 40 paths by 4947 tokens. The observation validates that only a small fraction of routing paths are frequently selected. Please see the uploaded PDF for more visualization.

[W3 - Comparison with SiDA] SiDA is an alternative approach that separates the router from the inference path. However, Read-ME offers better opportunities to improve inference efficiency because it is designed to enable expert-aware batching. Additionally, we would like to emphasize that, by design, the Read-ME router provides exact expert selection, while SiDA and prior works offer approximate selection [1,2]. In detail,

(1) Batching and latency: With SiDA, tokens cannot be batched together because they do not share routing decisions across all layers, unlike Read-ME (Since SiDA is distilled from the Switch Transformer, it is not possible to change decisions of routers). This vastly magnifies the expert space when batching and makes it difficult to compose an efficient batch that is aware of expert selection. In effect, SiDA activates more experts for each batch at each layer, leading to increased latency. Table A compares the inference latency between SiDA and Read-ME. (SiDA did not release a checkpoint, so we used SwitchTransformer-8, the teacher model from which SiDA was distilled, as a proxy) Table A. Inference latency of SiDA and Read-ME.

(2) Router accuracy: SiDA’s offline routing function is an approximation method that is distilled from original layerwise routers. Thus, the accuracy of routing is not 100%. Table B shows the “failure rate” of the SiDA router (defined as prediction miss rate on the expert activation of the trained router), and Table C shows the resulting degradation in task performance. With an incorrect guess, the entire expert selection can fail, leading to a collapse in inference, especially as the model scales. In contrast, our method is exact, ensuring no performance drop during inference, regardless of model size.

Table B. SiDA Router failure rate

Table C. Accuracy drop due to SiDA’s Router Failure

(3) Training cost: Note that SiDA is based on an MoE model and only distills the router, whereas our model is based on a dense model and builds both the expert network and the router, training them together. This means that SiDA's training time only accounts for the router training cost, while our training time includes both the router training cost and the expert specialization cost. Additionally, SiDA did not report the training cost in terms of tokens, making a meaningful comparison of training costs impossible.

In addition, our work introduces novel contributions to caching and prefetching, which SiDA lacks. SiDA relies on on-demand expert loading and FIFO-based expert eviction, which can negatively impact performance. Please refer to Section 4 for a detailed discussion of our contributions in these aspects. We will add discussion of SiDA in the revised version.

[1] Fast inference of mixture-of-experts language models with offloading

[2] SiDA: Sparsity-Inspired Data-Aware Serving for Efficient and Scalable Large Mixture-of-Experts Model

[W4 - Theoretical Analysis] In this work, we empirically examine the redundancy of layer-wise routers and propose a system-oriented method to convert a pre-trained dense model to an MoE with minimal additional training costs. While our focus is on empirical validation, providing a theoretical justification is beyond the scope of our paper and something we hope to explore in future study.

Thanks for catching the missing references and indices. We will definitely correct it in our revised version.

[W1 - Comparison with layer-by-layer loaded Llama ] Thanks for the question. We compared the latency of the layer-by-layer loaded llama and Read-ME model in the following table.

Table A. Latency comparison

The difference arises from the size of the parameters to be loaded. The size of Read-ME's expert layer (including both the top-1 expert and residual expert) is approximately 25.0% that of the MLP counterpart in the LLaMA baseline. Regarding peak memory usage, Read-ME exhibits a 10% lower profile, though both models consume an insignificant amount of memory overall.

Our method converts a dense pre-trained model into a MoE for efficient inference. Although this compression is generally not lossless, please note that our approach offers a better accuracy-cost trade-off than all baseline compression methods.

[Q1 - Understanding of Equation 3] Sorry for the confusion. We found this is a typo in our original submission., As we are selecting the mask $\boldsymbol{M}$ to maximize the magnitude of activated channels, the operator should indeed be $\arg\max$ in Equation 3.

[Q2 - Ablation on Routing distillation loss] Thank you for the suggestion. We have ablated the routing distillation loss by removing the router training step and using a random routing mechanism instead, tuning only the expert weights using the language modeling loss. To ensure a fair comparison, we maintained the same number of training tokens and the same training schedule. The resultant performance is provided in Table B. The results show that with the routing distillation loss, the average downstream task accuracy increased from 51.5% to 55.5%, validating the necessity of the routing distillation loss.

Table B: Ablation study on routing distillation loss. We compare the performance with and without the routing distillation loss, while keeping the number of training tokens, to validate the necessity of routing distillation loss.

[Q3 - Read-ME compared to Figure 3] This is a great suggestion! Please check Fig 3 of the PDF uploaded.
In Fig 3-left, Read-ME batches tokens directed to the same expert, so all arrows point to a single expert at each layer. Also, there is only one router instead of three, resulting in all arrows following the same path from layer 1 to layer 3, leading to a unique activated expert count of 1. In Fig 3-middle, Read-ME shifts the distribution to the lower left, where most points stay within the range of x < 4.5 and y < 55. Fig 3-right already compares the traditional approach with Read-ME.

[Q4 - Baseline numbers in Figure 4] Thanks for asking.

• MMLU is the common evaluation benchmark for LLMs, so most of the numbers in Figure 4 are collected from their original paper, including OpenMoE, Llama-2, LaCo, and Compresso; For Sheared-Llama, Open-Llama and Pythia, as they do not report MMLU in their original paper, we test their publicly released checkpoints with the lm-eval-harness library[1]; LLM-Pruner and SliceGPT didn’t report MMLU nor release checkpoints, so we use the numbers reported in LaCo [3].
• Yes, all of the methods reported in Figure 4 use the same base model - Llama-2.

[Q5 - Training Cost Calculation] Thanks for the suggestion. First, our training cost calculation follows Sheared-Llama[2] (see Table 1 of the paper), a representative post-training method to generate a small model out of a large one by pruning. Second, the training cost here measures the computational resources needed by the LLM deployer to obtain a new model of an auxiliary architecture (a smaller model in the Sheared-Llama case, and a MoE in our case), given all the publicly available resources. Using the publicly released checkpoints (e.g. Llama-2) will not incur additional training costs. We appreciate the suggestion and will mention this in the caption. Please see the PDF for our tentative Table 1.

[1] https://github.com/EleutherAI/lm-evaluation-harness

[2] Sheared LLaMA: Accelerating Language Model Pre-training via Structured Pruning

[3] LaCo: Large Language Model Pruning via Layer Collapse