Oracle-MoE: Locality-preserving Routing in the Oracle Space for Memory-constrained Large Language Model Inference

Thank you for your professional review comments and suggestions. We will add them in the updated version. Figures mentioned can be accessed at: https://anonymous.4open.science/r/ICML2025REBUTTAL-E158/README.md

Q1 Semantic locality: While the paper claims that tokens with higher mutual attention scores share similar high-level semantics, I am not convinced there is strong evidence of this.

A1 Previous studies [A][B] on representation space analysis have shown that semantically similar samples exhibit higher similarity in their embeddings compared to semantically dissimilar ones, which is also widely validated in experiments with general-purpose large models. We corroborate this observation and further identify a more fine-grained similarity pattern: token representations encapsulate both high-level semantics and token identity semantics. Among tokens with the same identity, the embeddings of those that share the same high-level semantic meaning tend to be more similar. This pattern is consistently observed in various models, including widely used large models like DeepSeek-16B-2.8B and Qwen1.5-MoE-A2.7B, which are illustrated in Figure 2 in our paper and Figure a in the link. Theoretical insights into how attention mechanisms compute correlations between tokens using the inner product of query (Q) and key (K) vectors are also supported by existing studies[C][D][E][F]. The computation of attention scores involves first assessing token correlations through inner products of query (Q) and key (K) vectors, followed by normalization of these correlations via softmax, and finally allocating contextual information through value (V) vectors weighted by the normalized scores. Among which, the Q-K inner product effectively captures token similarity and reflects high-level semantic alignment, as visualized in Figure b,d in the link.

[A] A Survey on Word Embeddings: From Shallow Models to Deep Learning (Goldberg, 2017)

[B] Deep Learning for NLP and Speech Recognition" (Hinton et al., 2012)

[C] Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer (Raffel et al., 2020)

[D] Effective Approaches to Attention-based Neural Machine Translation(Luong et al., 2015)

[E]Analyzing the Structure of Attention in Transformers (Kobayashi et al., 2020)

[F]On the Expressive Power of Self-Attention Matrices （Valerii Likhosherstov et al.,2021）

Q2 Effectiveness of oracle space: The paper claims that the oracle space efficiently describes various high-level semantics and that routing in this space preserves semantic locality; again, there is now strong evidence of this.

A2 As mentioned above, Self-Attention, especially attention score, can reflect and amplify the high-level semantics across tokens in the same context. Therefore, oracle space, which is constructed by grouping tokens with high attention scores and extracting high-level semantics from them, can capture semantic locality well. To further demonstrate the generalizability of this conclusion, we constructed oracle space with DeepSeekMoE-16B\Qwen1.5-MoE-A2.7B on several datasets, as illustrated in Figure b-e in the link. We found that semantic group embeddings in these oracle spaces all preserve semantic locality well, with semantic group embeddings varying slowly and smoothly, showing the potential for scaling up and generalization.

Q3 Narrow scope of application If you're in a memory-constrained LLM inference setup, maybe you shouldn't be using MoEs.

A3 Pursuing the trade-off between performance and latency is a key topic in the field of LLM inference. The MoE structure, which scales up model performance without increasing the number of activated parameters, is the most edge-friendly model architecture. Recently, sparse MoE models like DeepSeek and QwQ have exhibited impressive ability, demonstrating the value of MoE edge deployment. Companies like Qualcomm have also started research concerning MOE edge deployment, which was presented at the NeurIPS 2024 enterprise section, illustrating the huge market value of this problem. It is believed that this technology will further increase the accessibility to LLMs at edge devices, better highlighting the advantages of MoE on edge devices.

Thank you for your professional review comments and suggestions. We will add them in the updated version. Figures mentioned can be accessed at: https://anonymous.4open.science/r/ICML2025REBUTTAL-E158/README.md

Q1 The main concern I have with this paper is the issue of scalability. In the paper, the largest MoE model is 2B with 32 experts. However, the existing models like Mistral-8x7B [1] and DeepSeekMoE-16B [2] are much larger. It would be better if the authors could provide results on these models to ensure the scale-up and scale-out ability of the proposed method.

A1 Considering limited training resources and time, we train a model following the setting of DeepSeekMoE-16B but with fewer parameters(3B): 12 MoE layers with 64 routed experts each, where hidden size is set to 1536 and expert intermediate size is set to 1024. Top 6 experts are selected for each token. Our method still achieved a 75% latency reduction at 2.5GB memory. Meanwhile, our model maintains the performance of downstream tasks. On Trivia QA, our model achieves an F1 Score of 50.20, compared to the baseline of 50.75. On XSum, our model attains a ROUGE-1 score of 21.74, while the baseline score is 21.22.

Q2 How would the number of sampled data in oracle space initialization impact the accuracy performance? Would this cause a potentially high overhead for larger models?

A2 As for the sampling overhead, the largest batch size we sampled in our experiments was 16382, and it only took 20 minutes for our 8*GTX3090 platform to process the sampling. Compared with 32.94 hours of training time, the sampling and oracle space construction contribute to only 1% of the overall wall-clock time, which is negligible. At the inference stage, our method introduces no inference overhead, thus there is no need to worry about larger models.

Q3 In terms of figure location, I suggest putting all figures on the top of the page instead of in between the texts (see Figure 2/3/4).

A3 Thank you for your suggestions. We will make such modifications to the layout of these pictures in the updated versions.

Thank you for your professional review comments and suggestions. We will add them to the updated version. The figures mentioned can be accessed at: https://anonymous.4open.science/r/ICML2025REBUTTAL-E158/README.md

Q1 The paper could benefit from further discussions on practical deployment considerations and real-world constraints, such as more diverse hardware scenarios beyond the NVIDIA Jetson platform, like A100s and H100s.

A1 Results on A100s are listed below. Our method still speeds up inference by 50%~350%.

Q2 The concept of predicting expert activations in deeper layers based on shallow layers is promising, but the rationale and robustness of the 85%-95% prediction accuracy are left for future exploration. More details on this approach would strengthen the work significantly.

A2 We predict by training a linear classifier with representations in shallow layers as input and routing results in deep layers as target labels. It turns out that such a linear classifier can reach an accuracy of 85%~95% in our structure and 40%~60% in existing MoE structures. More fine-grained observations show that layers that are closer predict each other well, e.g., representations in layer 10 predict routing results in layer 12 better than that of layer 5, and layers that predict each other well show a grouped pattern. We have been studying this phenomenon and believe this can be attributed to residual connections maintaining semantics across layers to some extent.

Q3 While semantic locality is effectively leveraged, the paper does not deeply investigate scenarios where semantic locality is minimal (highly diverse or abrupt topic changes), potentially limiting generalizability.

A3 We tested scenarios where the topic changes frequently. We randomly sample sentences from different datasets and combine them into a whole sequence. We observed that our proposed oracle space can still distinguish semantic groups efficiently, both in our models and public large MOE models (DeepSeek-16B-2.8B\Qwen1.5-MoE-A2.7B), as shown in Figure g,h,i,j in the link. We also tested the expert activation variation of such highly diverse data with Oracle-MoE and switch-transformer. On average, in every 100 consecutive token generations, Oracle-MoE only changes 12.20 times while the switch transformer changes 90.54 times. This is because in human natural language, it takes at least dozens of tokens to express a complete meaning so our method still benefits from such "abrupt" semantic locality.

Q4 The paper did not discuss cases of using fine-grained experts ( e.g., number of experts > 128, like the ones in DeepSeek's model).

A4 Considering limited training resources and time, we train a model following the setting of DeepSeekMoE-16B but with fewer parameters(3B): 12 MoE layers with 64 routed experts each, where hidden size is set to 1536 and expert intermediate size is set to 1024. The top 6 experts are selected for each token. Our method still achieved a 75% latency reduction at 2.5GB memory. Meanwhile, our model maintains the performance of downstream tasks. On Trivia QA, our model achieves an F1 Score of 50.20, compared to the baseline of 50.75. On XSum, our model attains a ROUGE-1 score of 21.74, while the baseline score is 21.22.

Q5 typo "Lanuge".

A5 We will correct the typos in the updated version.

Thank you for your professional review comments and suggestions. We will add them in the updated version. Figures mentioned can be accessed at: https://anonymous.4open.science/r/ICML2025REBUTTAL-E158/README.md

Q1 It would be better to measure the temporal inconsistencies for the whole dataset and different layers. It would be more convincing to provide a metric to quantitatively measure the temporal inconsistency and show how Oracle-MOE reduces this inconsistency.

A1 We propose temporal activation inconsistency, defined as the average number of inconsistent expert activations per 100 consecutive tokens per expert. Results over the entire dataset and across different models and layers are listed below. Existing MoEs show strong temporal activation inconsistency within all layers, while Oracle-MoE reduces this.

Q2 It would be better to provide some preliminary experiments to show the evidence of semantic locality in real datasets, .. across different models/layers/samples.

Experiments with DeepSeekMoE-16B and Qwen1.5-MoE-A2.7B on real chat datasets(Wizard-of-Wikipedia and Synthetic-Persona-Chat) are shown in Figure a-e in the link. Semantic locality appears across different models/layers/samples. Semantic groups can still be distinguished based on attention score and obtained by our method. It indicates the potential of Oracle-MoE being a general-purpose solution.

Q3 Can the authors explain more about why the mapping of Q/K/attention score will group consecutive tokens with similar semantics? Why does this happen for different layers/samples?

A3 Previous studies [A][B] have shown that semantically similar samples exhibit higher embedding similarity than semantically dissimilar ones, widely validated in experiments with large models. We corroborate this observation and identify a more fine-grained pattern: token representations encapsulate high-level and token identity semantics. Among tokens with the same identity, embeddings of those sharing the same high-level semantics tend to be more similar. This pattern is consistently observed in models like DeepSeek-16B-2.8B and Qwen1.5-MoE-A2.7B, as illustrated in Figure 2 in the paper and Figure a in the link. Theoretical insights are also supported by [C][D]. Computing attention scores involves first assessing token correlations through inner products of Q and K vectors, normalization via softmax, and allocating contextual information through V weighted by the normalized scores. Among these, the Q-K inner product effectively captures token similarity and reflects high-level semantic alignment, as visualized in Figure b,d in the link.

[A] A Survey on Word Embeddings: From Shallow Models to Deep Learning (Goldberg, 2017)

[B] Deep Learning for NLP and Speech Recognition (Hinton et al., 2012)

[C] Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer (Raffel et al., 2020)

[D] Analyzing the Structure of Attention in Transformers (Kobayashi et al., 2020)

Q4 There are some approximations in the derivation of oracle-MOE. It would be better to validate such an approximation in the experimental design.

A4 Experiments validating this approximation are listed below: for token pairs with the distance between [x,y), we count on average how many inconsistent expert activations are triggered. There is a statistically positive correlation that the higher the distance between tokens, the more likely they activate different experts. More results are in Figure f in the link.

Q5 Ablation study on CSD and hyperparameter study of γ are missing. How does γ influence the proposed algorithm?

A5 CSD and γ are key metrics we defined for the optimization problem of minimizing the model's consecutive inconsistent activation while maintaining its task performance above γ. In our paper, we design oracle space routing to reduce CSD while maintaining γ. We found that the number of MoE layers and experts per layer can influence CSD and γ as is listed in the table below. CSD is given as the average number of inconsistent expert activations per 100 tokens in all layers and γ is given as the average of all tasks tested. Our model has a larger γ and less CSD, performing consistently well at large expert numbers, demonstrating our method's robustness to hyperparameters.