Mixture of Lookup Experts

We thank the reviewer for their thoughtful feedback. We address specific concerns and questions below.

Q1. As the model scales up, the usability of MoLoE may come into question.

Firstly, the number of LUT parameters that need to be offloaded by our method is $dN|V|$, whereas MoE needs to offload the expert size of $2dND_r$. Since $|V|$ is typically on the order of tens of thousands, the difference between the two is not greater than an order of magnitude. Additionally, the method is designed for edge deployment of LLMs, where the parameter size is constrained by the available computational resources, typically resulting in smaller models. Therefore, the size of the LUTs stored on lower-level storage devices is generally acceptable.

Q2. When the prompt length exceeds 10k, MoLoE-4E would need to load more parameters than MoE-10E

In the extreme case of very long texts, the amount of parameters our method loads during the pre-filling phase may indeed be greater than that of MoE. However, during LLM inference, the time required for multi-step decoding remains the dominant factor, meaning the cost during the prefilling phase can be somewhat offset. Additionally, directly loading the expert parameters before reparameterization during the pre-filling phase and using computation to replace loading is another potential solution.

Q3. Why is MoE significantly slower than the dense model?

Our implementation is based on the Mixtral code from the Huggingface transformers package. Unlike cluster environments that can utilize expert parallelism, edge-side single-card deployment requires each activated expert to process tokens individually, resulting in lower parallelism. When the batch size is greater than 1, the parallelism is further reduced due to different tokens requiring different experts, leading to slower speeds. It is worth noting that the inefficiency of MoE computation could potentially be addressed in the future through more friendly designs of hardwares or better operator implementations. However, even if the single-card execution efficiency of MoE becomes comparable to that of dense models, our method still maintains a significant advantage in terms of the latency associated with loading parameters.

Q4. Is it possible to use modules other than FFN as experts?

In fact, as long as the module does not involve information exchange between tokens, i.e., the module’s input is context-independent, it can serve as an expert in MoLoE. This characteristic is similar to the FFN layer in transformers. Therefore, we adopt FFN as the expert, ensuring that the type of expert is consistent with the MoE model.

We thank the reviewer for their thoughtful feedback. We address specific concerns and questions below.

Q1. Pruning LUT entries corresponding to rare token IDs.

We prune half entries of the LUTs of MoLoE-16E with 160M activated parameters based on frequency. When the input token ID is pruned, the LUTs return a zero vector. The results are as follows:

From the experimental results, it is evident that post-training LUT pruning significantly degrades model performance. If smaller LUTs are required, using a more compact vocabulary during the training phase is a better choice.

Q2. Why did the MoE baselines not use shared experts?

The different settings for MoE and MoLoE are chosen in order to prioritize aligning both total parameter count in training and activated parameter count in inference for a fair comparison.

• MoLoE adopts shared experts because, unlike MoE, MoLoE reparameterizes all routed experts into LUTs during inference. If shared experts were not used, MoLoE would not include any FFN layers during inference, which would make it difficult to align its activated parameter count (VRAM usage) with that of dense and MoE models.
• The reason MoE does not adopt shared experts is that, in our experiments, MoE without shared experts performs better, which has also been thoroughly validated in [1]. We provide our experimental results below (No.1 vs No.2).
• In inference stage, MoLoE only has 1 shared expert, while MoE activates 2 experts. Therefore, to make a fair comparison, we set the size of each expert in MoE to half of that in MoLoE (to align activated parameter count in inference), but double the number of experts (to align total parameter count in training). As a result, the total number of experts cannot be aligned (No.2 vs No.3). But as a reference, we have still provided the results below for MoLoE and MoE both using 1 shared expert and 33 routed experts (No.2 vs No.4).

[1] OLMoE: Open Mixture-of-Experts Language Models. ICLR 2025

Q3. Add citations to the papers mentioned above.

Thank you for your suggestion, we will update it in the revised version.

We thank the reviewer for their thoughtful feedback. We address specific concerns and questions below.

Q1. Will MoLoE require more training time than MoE?

Since MoLoE activates all parameters during training, its training cost is approximately equal to that of a dense model with the same number of parameters, and higher than that of a sparse MoE model with the same parameter count. However, it should be noted that our method is designed for edge deployment of LLMs, where the parameter size is limited by the computational resources available on the edge, typically resulting in smaller models. Therefore, the cost of training these smaller-scale LLMs on a cluster remains manageable.

Q2. Why is MoLoE able to perform better than MoE?

• Although the experts' inputs do not contain context information, they can still indirectly influence the model's handling of context by altering the input to subsequent attention layers. Additionally, the retained shared expert operates similarly to MoE and dense models, where it takes the output from the previous layer as its input, and thus can also handle context-related information.
• By changing the expert inputs to embedding tokens, the routed experts can be reparameterized as LUTs, resulting in almost zero computation during inference. This allows us to activate all experts without incurring additional costs, which helps recover some of the performance loss caused by altering the expert inputs (see Table 7). Additionally, since activating all experts does not require auxiliary losses, the training and inference objectives are more aligned, leading to improved model performance (see Table 4).

We thank the reviewer for their thoughtful feedback. We address specific concerns and questions below.

Q1. The comparison between MoE and MoLoE is inconsistent.

The different settings for MoE and MoLoE are chosen in order to prioritize aligning both total parameter count in training and activated parameter count in inference for a fair comparison.

• MoLoE adopts shared experts because, unlike MoE, MoLoE reparameterizes all routed experts into LUTs during inference. If shared experts were not used, MoLoE would not include any FFN layers during inference, which would make it difficult to align its activated parameter count (VRAM usage) with that of dense and MoE models.
• The reason MoE does not adopt shared experts is that, in our experiments, MoE without shared experts performs better, which has also been thoroughly validated in [1]. We provide our experimental results below (No.1 vs No.2).
• In inference stage, MoLoE only has 1 shared expert, while MoE activates 2 experts. Therefore, to make a fair comparison, we set the size of each expert in MoE to half of that in MoLoE (to align activated parameter count in inference), but double the number of experts (to align total parameter count in training). As a result, the total number of experts cannot be aligned (No.2 vs No.3). But as a reference, we have still provided the results below for MoLoE and MoE both using 1 shared expert and 33 routed experts (No.2 vs No.4).

[1] OLMoE: Open Mixture-of-Experts Language Models. ICLR 2025

Q2. Results under complex language modeling tasks like WikiText or Winogrande.

We reported the results for a 1B activated parameter model on WikiText and Winogrande. MoLoE outperforms MoE on Winogrande, but performs slightly worse on WikiText.

Although the routed experts in MoLoE do not directly receive context information, they can still indirectly influence the model's handling of context by altering the input to subsequent attention layers. Additionally, the retained shared expert operates similarly to MoE and dense models, taking the output from the previous layer as its input, and thus can also handle context-related information.

Q3. Different tokenizer settings.

In addition to Pythia's tokenizer with a vocabulary size of 50k that we used, we also tried the 32k vocabulary tokenizer from Mixtral. The results are shown below.

The primary impact of the tokenizer on MoLoE is that its vocabulary size is proportional to the size of the LUTs. However, since different tokenizers also affect the performance of all kinds of models, it is difficult to isolate the tokenizer's specific effect on MoLoE's performance. Nonetheless, at least from the experimental results, we can observe that MoLoE is effective under different tokenizers.

Q4. Compared with Pre-gated MoE.

Pre-gated MoE allows for the parallelization of expert prefetching and computation. Compared to this approach, our method has the following advantages:

• Pre-gated MoE does not increase inference latency under the assumption that the time to load parameters is less than or equal to the computation time. However, during batch decoding, since different tokens require different experts, the number of experts to prefetch increases (in the worst case, all experts need to be loaded), which in turn increases the latency.
• Additionally, when deployed on edge devices, it may be necessary to offload to lower-tier storage devices such as SSDs, where the transfer bandwidth is much smaller, causing the time to load parameters to be significantly higher than computation time, leading to a substantial increase in overall latency.
• In contrast, our MoLoE fundamentally reduces the amount of data that needs to be transferred by more than three orders of magnitude, thus avoiding additional latency during decoding. The two methods actually offer two different solutions in distinct directions.