Oracle-MoE: Locality-preserving Routing in the Oracle Space for Memory-constrained Large Language Model Inference
Authors: Jixian Zhou
Fang Dong
Ruijun Huang
Hengjie Cao
Mengyi Chen
Yifeng Yang
Anrui Chen
Mingzhi Dong
Yujiang Wang
Dongsheng Li
David A. Clifton
Qin Lv
Rui Zhu
Chun Zhang
Fan Yang
Tun Lu
Ning Gu
Li Shang
Venue: ICML 2025 poster (2025)
Keywords: MoE; Edge Device; Inference Latency; Semantic Space
Primary area: deep_learning->large_language_models
TL;DR: Oracle-MoE optimizes expert activation consistency and reduces inference latency by leveraging semantic locality, enabling efficient deployment of large language models on memory-constrained edge devices.

Abstract: Mixture-of-Experts (MoE) is widely adopted to deploy Large Language Models (LLMs) on edge devices with limited memory budgets.
Although MoE is, in theory, an inborn memory-friendly architecture requiring only a few activated experts to reside in the memory for inference, current MoE architectures cannot effectively fulfill this advantage and will yield intolerable inference latencies of LLMs on memory-constrained devices. 
Our investigation pinpoints the essential cause as the remarkable temporal inconsistencies of inter-token expert activations, which generate overly frequent expert swapping demands dominating the latencies. 
To this end, we propose a novel MoE architecture, Oracle-MoE, to 
fulfill the real on-device potential of MoE-based LLMs. 
Oracle-MoE route tokens in a highly compact space suggested by attention scores, termed the *oracle space*, to effectively maintain the semantic locality across consecutive tokens to reduce expert activation variations, eliminating massive swapping demands. 
Theoretical analysis proves that Oracle-MoE is bound to provide routing decisions with better semantic locality and, therefore, better expert activation consistencies. 
Experiments on the pretrained GPT-2 architectures of different sizes (200M, 350M, 790M, and 2B) and downstream tasks demonstrate that without compromising task performance, our Oracle-MoE has achieved state-of-the-art inference speeds across varying memory budgets, revealing its substantial potential for LLM deployments in industry.

Oracle-MoE: -oocalliyyppeseevinn Routing in the Oracle Space for Memory-constrained Large Language Model Inference
Jixian Zhou" " Fang Dong" I
Ruijun Huang 12 Hengjie Cao Mengyi Chen Yifeng Yang' Anrui Chen' Rui Zhu? Chun Zhang*
Mingzhi Dong' Yujiang Wang* Dongsheng Li David A. Clifton
Qin Lv
Fan Yang
Tun Lu ! Ning Gu Li Shang'
Abstract
fore, better expert activation consistencies. Exper- iments on the pretrained GPT-2 architectures of different sizes (200M, 350M, 790M, and 2B) and downstream tasks demonstrate that without com- promising task performance, our Oracle-Mot has achieved state-of-the-art inference speeds across varying memory budgets, revealing its substantial potential for LLM deployments in industry.
of (MoE) is widely adopted to deploy Large Language Models (LLMs) on edge devices with limited memory budgets. Although MoE is, in theory, an inborn memory-friendly architecture requiring only a few activated ex- perts reside in the memory for inference, rent MoE architectures cannot effectively fulfill this advantage and will yield intolerable inference latencies of LLMs on memory-constraie vices. Our investigation pinpoints the essential cause as the remarkable temporal inconsistencies of inter-token expert activations, which generate overly frequent expert swapping demands dom- inating the latencies. To this end, we propose a novel MoE architecture, Oracle-MoF, to fulfill the real on-device potential of MoE-based I.LMs. Oracle-MoB route tokens in a highly compact space suggested by attention scores, termed the oracle space, effectively maintain the seman- tic locality across consecutive tokens reduce expert activation variations, eliminating massive swapping demands. Theoretical analysis proves that Oracle MoE is bound to provide routing de- cisions with better semantic locality and, there-
1. Introduction
de-
The unprecedented prevalence of Large Language Models (LLMs) calls for their sink from remote servers local edge devices for better accessibility, privacy protections, and personalizations. Mixture-of-Experte (MoE) (Fedus et al. 2022 Cai et al. 2024) is the most widely adopted technique for the on-device deployment of LLMs. A well- established advantage of MoE is that could effectively scale up LLMs' capacities without increasing the inference computational complexity, which is vital for their execu- tions edge devices of less powerful computing hardware. However, we would like to underline crucial property of MoE: it is, in theory, memory- -fiendlly architecture that could naturally satisfy the memory con- straint edge devices. Only a few activated experts are needed reside in the device memory for a feedforward MoE inference, and thus, the minimum memory occupa- tions of MoE-based LLMs can be significantly less than their dense counterparts, which indicates the edge-fiiendly nature of MoE architectures.
missmneni but
'qqual contribution 'College of Computer Science and Arti- ficial Intelligence, Fudan University, Shanghai, China *Greater Bay Area National Center of Technology Innovation. Research Institute of Tsinghua University in Shenzhen, Shenzhen, China University of Bath, Bath, UK 'Oxford Suzhou Centre for Ad- vanced Research, Suzhou, China 'Microsoft Research Asia, Shang- hai, China DDpaarmenn of Computer Science, University of Col- orado Boulder, Colorado, USA Baaes Business School, City St George's, ss University of London, London, UK "Research Institute of Tsinghua University in Shenzhen. Shenzhen, China "School of Microelectronies, Fudan University. Shanghai, China. Corre- spondence to: Shang <lishang @ fudan. aod.a.edd..ns. edu.en>, Yujiang Wang <yujiang wang @ oscar.ox.ac.ut
In practice, the fulfilling of MoE's memory-saving advan- tage, however, is still far from satisfying. To date, deploy- ing MoE-based LLMs on memory-constraiee commonly introduce intolerable inference latencies that sub- stantially impair the industrial values of LLMs. Those la- tencies are predominantly caused by the overheads from frequently loading activated experts into the memory and cleaning old ones during successive inferences of tokens. As shown Figure [(a) (Top), the constried in- ference of MoE-based LLMS can result latencies approxi-
devices will
Proceedings of the 42nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Oracle-MoE: Locality-preserving Routing in the Oracle Space for Memory-consttaineed Large Language Model Inference
is typically consistent, and successive queries from user usually share similar knowledge grounds. In other words, token sequence represents notable semantic locality across nearby tokens. Inspired by this observation, Oracle-MoE conducts routing in a highly compact space, termed the oracle space to assign tokens instead of the usual token embeddings.
Latency Broakdlown
Expert Autivesiion in Owen Layer? 7 &.
On decosed
Pull Model
10
MesseryCiio
and 20
Throgiage
1.0
20
Expert Comperitng
Export Swappcng
Jestalto Mertecy(ti)
Tolers the sequizee
150
Particularly, our analysis revealed that tokens with higher mutual attention scores share similar high-level semantics, which we call semantic groups. To efficiently extract high- level semantics, we use average token embeddings in each group as semantic group embeddings and denote the set of these embeddings as oracle space. Empirical evidence and theoretical analysis demonstrate that this concise calculation of oracle space captures semantic locality and routing in the oracle space could preserve semantic locality better when selecting experts, considerably decreasing the swapping fre- quencies. On GPT--2bbseed networks (Radford et al.. 2019) of different capacities (200M, 350M, 790M, and 2.06B) and the downstream tasks including classification, QA and sum- marization, our Oracle-MoF. has attained state-of-the-an inference speeds than those baseline strategies across vary- ing memory budges without comprising the performance.
(a)
(b)
Figure 1. a). Top: compared with full model residing in the memory, memory-comsan conttrreed inference of LLMs suffers from severe latency issues; Bottom: 50%-85% latencies are incurred by the I/O overheads from expert swapping. An edge device with 8GiB GPU memory is adopted. b). Visualization of 60 expert activations of layer 7 in Qwen Yang et al. 2024) across sequence of 300 tokens.
mately 15 to 30 times higher than the same model entirely residing in the memory. 50%-85% of those latencies, as portrayed in Figure ((a) (Bottom), were incurred by the I/O overheads from expert swappings. Despite several prior attempts (Kong et al. 2024; Yi et al. 2023] Huang et al. 2023b; /Zhang et al. 2024] Rajbhandari et al. 2022) to op- timize those latencies, they followed an engineering- perspective to develop comparatively naive and empirical strategies, such as quantization (Zhu et al. 2023), expert
2. Analysis
In this section, we begin with defining the latency opti- mization problem in semantic space with introducing the Consecutive Semantic Difference (CSD). Given that con- secutive tokens show semantic localities, we aim to find low-variance semantic embedding for each token. We first model the token embeddings the combination of high- level semantics and iokentidentity semantics. Then, we use attention scores to discover high-level semantics similar- ity and obtain semantic groups, and construct the Oracle Space with semantic group embeddings. Finally, we show that oracle-space- based routing yields significantly lower CSD compared to token-level routing, both theoretically and empirically.
pruning (Cheng et al., 2023 Rajbhandari et al./ 2022). re-
taining frequently-ssed experts (Kong et al.
2024: Yi et al.
2023 Huang et al. 2023b Artetxe et al. 2022). which lack systematic, in-depth investigations.
This work thoroughly examines the underlying mechanism behind the high latencies of nemory-constraing MoE- based LLM inference. Figure [(b)) the inter-token activations of 60 experts of one layer across a sequence of 300 tokens, which exhibit a highly spare and temporally inconsistent pattern. Such a pattern will inevitably create a plethora of expert-swapping demands to be resolved, and is challenging develop an effective empirical strategy due to the complexity, sparsity, and inconsistency of activa- tions. The essential cause for this pattern can arguably be attributed to that current MoE routings are predominantly influenced by token embeddings, which tend to be domi- nated by loken-specifi information with high inter-token variations.
2.1. Lateney Problem Formalization
Deploying Mixture of Expert (MoE) models edge devices faces the challenge of expert swapping latency. Meanwhile, we do not want latency optimization to degrade the model's ss performance. So, we introduced constraint that the la- tency optimization must ensure that the measurement of performance M that the user interested in is better than a threshold 7:
To this end, we devise Oracle-MoE, novel MoE architec- ture to generate routing decisions that naturally activate iden- tical experts across consecutive tokes to reduce inter-token variations, climinating massive expert-swapping needs and accelerating memory-cone contory-onsttrained inference. Oracle-MoE is bridged on a characteristic of edge scenarios such that the linguistic meaning between consecutively generated tokens
min Leotal s.t. M >
Our experiments in Section 5 will show that latency opti- mization does not necessarily lead to performance degra- dation. For simplicity, we omit the subject to item in the
2

---

Oraclc-MoE: Locality-pessevinng Routing in the Oracle Space for Memory-constrained Large Language Model Inference
following of this section. The total latency Looal be written as:
where C(Wgr kijis a constant dependent W k.
However, since there is large variance between token em- beddings, as is shown in Figure 2 left, the CSD is hard to optimize for existing token-level MoE. Xue et al. (0244 also mentioned that token-level MoEs tend to dispatch to- kens according to token ID. But the linguistic meaning of consecutive user interactions on edge devices notably similar, which we term as semantic locality. This inspires to minimize swapping latency by extracting high-level semantics and designing a new gating mechanism.
Liotal
compute (t) + Liswap(es 62-1),
where Tecopaaee (t.) denotes computation latency for token embedding tz, and et denotes the set of activated expert indices for token t.. Lewa et-1) captures swapping latency caused by loading experts that are activated by t, but not activated by 1- S.eciffiifiily
Les et-1) = lswap les \ et-1/:
2.2. Semantic Groups and Oracle Space
Token Embeddings In token-level MoEs, tokens in the same sequence do share some similar high-level semantics, with their high-level semantics' locality and tendency to cluster logically, as shown in Figure 2 But their embed- dings are still influenced primarily by the token-identity fea- tures, leading us to explore semantic spaces. These tokens simultaneously contain two types of semantic information.
where lex) \ 1t-1 is minus operation. In practice, MoB- based LLMs edge devices generate continuous semantics, the large variance in the token-level MoE routing mecha- nism is essentially due to the high inter-token semantic vari- ance, resulting in frequent changes in routing results. We define Consecutive Semantic Difference (CSD) to measure the variation expert selection over consecutive tokens:
Definition (Token Embedding). For each token embed- ding t..
T
t: S; U,
CSD
Aet
t=2
where represents shared high-level semantic information between consecutive tokens and u, represents unique token- identity information.
Given fixed hardware and swapping algorithms, Leompute and waap are fixed, thus CSD determines Taeual- The final optimization goal is:
Selected Token embeddings
15
Canada Hocky Night matchups
min Σ De
10
the
B
Remark: Setting Ae; 0 leads to dense model that would trivially minimize CSD, which is not considered in our work as it negates the fundamental advantage of MoE models multiple experts for specialized processing.
-5
sport
-10-
controverstal sentence case
net
-15-
-10
15
Figure 2. UMAP visualization of embedding space in existing token-level MoF models. Left: Tokens tend to cluster accord- ing to token-identity semantics. Right: Tokens from the same sequence are colored the same. They share similar semantics and stay closer to each other in each token cluster.
In existing token-level MoE models, experts are usually selected through linear gate Wy E Nxd (Cai et al. 2004).. which projects the token embedding t, E Rd into expert scores (t,)) W tt E NN The experts are then chosen as:
6 = Top-k kei1,... (softmaax Seaimaxxxx,,e ())
Semantic Groups Studies ( (Vaswani et al. 2017 Koval- eva et al. 2019 Clark et al., 2019)) indicate that attention captures high-level semantic correlations between tokens. Therefore, we intuitively claim that the mapping of the Q/K matrices and the calculation of the attention scores will group consecutive tokens with similar high-level semantics together through significant attention score distribution dif- ferences. Our analysis in Appendix demonsstats that under our definition of token embedding and analysis the Q/K matrix, high attention score between tokens indicates
where Top-k denotes the indices of top k largest elements. Inspired by previous work (uue et al., 2024), current token- level MoEs tend to dispatch on token-identity semantics. So. approximate the change in k expert selection between consecutive tokens by the difference in their trans- formed embeddings:
T
T c(W, g k)llt, t--1ll
CSD DSDokenn taken
Σ
De s3
2
2
3

---

Oraclc-MoE: Locality-presevving Routing in the Oracle Space for Memory-consttained Large Language Model Inference
that they share similar high-level semantics. So, we propose to model this through causal graph perspective, where semantic groups emerge from connectivity in the attention score matrix.
semantics. Among tokens with the same identity, the em- beddings of those that share the same high-level semantic meaning tend to be more similar. This pattern is consis- tently observed in various models, including widely used large models like DeepSeek-166: 168-228 (Dai et al. 2024) and Qwen Dwenl..5MofEs 1.5-MoE-A2..7 (Team] 2024), which are illus- trated in Figure Figure 10 and Figure I1 in Appendix [B.2 Theoretical insights into how attention mechanisms com- pute correlations between tokens using the inner product of query (Q) and key (K) vectors are also supported by ex- isting studies (Raffel et al. 2020 Vig & Belinkov 2019).
Consider a directed acyelic graph Gglobal = (Vglobat, Feghoba), ) where Vilabal contains all tokens and Egglobal consists of edges t. + t2 weighted by at from the lower-triangular attention score matrix Aglobal = [a.y] (i.e., diy exists only for > j). We define semantic groups as maximally connected components that only tokens with attention score larger than predefined threshold e are considered connected:
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 correla- tions via softmax, and finally allocating contextual informa- tion 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 Appendix B.2
Definition 2 (Semantic Group). A subset containing tokens S {tay.... ..... tx... J (indices kn < ... < kim) is called a semantic group if:
Vi,j E (kI,,..., No proper superset of S satisfies the condition above
(i>
Di >E
This can be regarded as reformulation of the Minimum Clique Cover Problem (Gavril 1972) for DAGs. The defini- tion leverages the block structure of attention score matrices as is shown in Figure 3, which is also observed in previous works((Lin et al. 2024)). So, although the Minimum Clique Cover is Na-hard, we claim that it can be computation- ally tractable on attention score matrix via polynomial-time greedy algorithms (Farjas 2018). We first initialize each token as a singleton group, then for the token x; from left to right, we find the maximal j < with a.j and merge on into zyys 's group if Vor in the group,
Oracle Space Following sentence meta meta-embeddingg tech- niques (Poemer et al., 2019 Takahashi & Bollegala 2022). we compute semantic group embeddings as the average token embeddings in it. Formally:
Definition 3 (Semantic Group Embedding). For a Semantic Group S its semantic group embedding zs, is defined as:
1
2s.
1/
S
6.
ES
As proven in (Xu et al. 2018 Soltanolkotabi et al. 2013) and demonstrated in Appendix A.2, this aggregation reduces token-identity noise while preserving essential high-level semantics. Thus, we can efficiently extract various high- level semantic information from the embedding space using semantic group embeddings. We collect semantic group em- beddings from different data and name the space consisting of these embeddings as Oracle Space, efficiently describing various high-level semantics.
Alleution Score
Attention Score
Tolaso Indox
10
10
12
25
30
30
Token Indox
Tokch Inde
In inference task, new tokens arrive sequentially over time. To model the evolution of semantic groups and derive token embeddings based on these groups, we propose the semantic embedding of each token as its semantic group embedding. Given a token t, at time step t,, let S(t) denote the semantic group corresponding to t, We use the embedding of S(t) as the token's semantic embedding:
(a)
(b)
Figure 3. Visualization of attention score matrix. There are two semantic groups where tokens in each group show high attention scores with each other.
Discussion: Previous studies (Kamath et al., 2019) on repre- sentation space analysis have shown that semantically sim- ilar 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 -rained similarity pattern: token representations encapsulate both high-level semanties and token identity
Zt. = ZS(s)
=
|s(t)
t. tES(C)
The S(t) includes all tokens from previous time steps t' < t such that:
Vtp E S(t), au > e,
4

---

Oraclc-MoE: Locality-preservigg Routing in the Oracle Space for Memory-constrained Large Language Model Inference
where our is the attention score between tokens b and ty- When the model generates consecutive tokens, it retains the KV cache of previous tokens, so that der can be obtained by adding a new row to Agiobal
Next, highlight the low-variation gains of oracle-space- based routing compared to token-level routing, which re- duces unnecessary routing changes for semantically similar tokens. This leads to the following theorem:
This provides a way to compute semantic groups with a streaming input, which the case for auto-regressive gen- eration. As shown in Figure 4 in an regrrssieee generation process, the token's semantic embedding varies smoothly and slowly in the oracle space, preserving high- level semantic information.
Theorem 1. With a high probability. vt
N. te, tt-1, aut-1 > with a sufficiently high dimen- sion of token identity information:
C(Wg; k)llt tz-1ll > ||z(() zs(t-D)
Through the CSD estimation framework established in Sec- tion 2. this induces the relationship:
Semantic Group Embedding
Semantic Group Embedding
20
20
CSDroken > CSD
This indicates that oracle-space-based routing provides lower variation by reducing routing changes for semanti- cally similar tokens. The detailed proof is provided in the Appendix A..4
0-
-10-
-10
-10
10
-10
lnken
Last tolon
loken
Last tolon
3. Oracle-MoE
(a)
(b)
Based on the above analysis, we propose Oracle-MoF, re- placing the token-level MoE routing mechanism with bet- ter semantic locality in the oracle space and expert activation consistencies for edge-side devices to generate text scenar- ios.
Figure 4. UMAP visualization of sampled semantic group embed- dings (token's semantic embeddings) in different model layers. Each color represents sequence(user interaction). As token gen- eration goes on, the embeddings based on semantic groups vary slowly and smoothly.
Average
Semantic Group Embeddings
2.3. Oracle-Space-Basee Routing Yields Lower CSD
Based on the above analysis, routing tokens with similar se- mantic embedding in the oracle space to the same expert is highly likely yield low expert variation. Since semantic group embeddings evolve locally in the oracle space (Fig- 4 employ clustering algorithms allocate tokens that belong to the same cluster in the oracle space to the same expert, or:
Attention
Semantic Groups
Score
000
Expert
Oracle Space
Actention
Expert 2
Route
e arg min [tt, cxll arg min [(s()) ck
Input Tokens
Expert
where denotes the cluster center in oracle space, which will be detailed in Section 3] So, the change in expert assignments can be approximated by the change of semantic group of 2t,- We provide a detailed analysis of this trend in the AppendixA.33 and the following form:
Figure 5. Overview of an Oracle-MoE layer. Residual connections omitted.
Oracle Space Initialization We obtain the initial oracle space after a short warm-up training phase of the token- level MoE model. After the warm-up stage, we randomly sample N data, and get the semantic group embeddings of each data as is mentioned in Section 2 Definition 2 and
Definition 4 (CSD for oracle-space-basee based routing). Lei S(t - and S(t) denote the semantic groups of token t;, jat consecutive time steps. Without loss of generality, based on semantic variations on semantic groups, we have CSD of oracle-space-bun based routing:
These semantic group embeddings form an initial ora- cle space. Routing in the oracle space does not necessitate complete information; only distinguishing high-level se- mantics is required. Therefore, to improve computational
T Det.omoie oomele &
T
CSDenacle
Σ
(3(1)) 2 25(11))
2
2
5

---

Oraclc-MoE: Locality-preserving Routing in the Oracle Space for Memory-constrained Large Language Model Inference
efficiency, we adopt SVD to reduce these embeddings to lower dimensions (Schmidt 2020).
2017 Dauphin et al. 2017 Vaswani et al. 2017) has been proposed and achieved great success in variety of tasks. MoE has been widely explored to improve the training ef- ficiency of Large Language Models(LLMMs), with various routing strategies like (i) letting tokens select the top-k ex-
Oracle-MoE Routing at Pretraining Since the semantic group embedding varies only locally within small region in the oracle space as tokens are generated, to ensure tokens in the same semantic group be dispatched to the same expert, we run K-means (Jain 2008) in the oracle space to get i oracle clusters. The parameter k here is equal to the expert number of the original token-level MoE, and Sec- tion 5J will show that our method yields persistently good results with ranging from 4 to 32. In the following training process, we replace the routing of the original token-level MoE with the routing mechanism shown in Figure [5] For each new incoming data, we first divide it into semantic groups according to the attention scores and get their seman- tic group embeddings in the oracle space with the same SVD transform matrix computed in the previous stage. Then we calculate which oracle cluster each semantic group belongs to as Equation [] and dispatch tokens in the semantic group to the corresponding expert.
perts (Lppikhnn et al. 2021; Fedus et al. 2022 Zuo et al.
2022, Chi et al. 2022 Dai et al. [022]
Chen et al.
2023
iii letting experts select the top-k tokens (Zhou et 2022
to (iii) globally decide expert assignment
(ewiis et al.
2021
Clark et al. 2022).
Memory-constrained Inference Solutions for Mixture of Experts Huang et al. propose swapping experts from GPU memory to CPU memory to reduce memory usage, though this approach introduces significant latency over- head (Huang et al., (023a). SE-MoE (Shen et al.] 2022) employs Ring Memory offloading to minimize GPU mem- ory consumption. EdgeMoE (Yi et al.) (023), the on-device inference engine tailored for MoE-based LLMs, reduces ex- pert I/O overhead through expert-wise bitwidth adaptation and in-memory expert management, achieving low-latency inference while maintaining task performance. While Swap- MoE (Kong et al. 2024), by dynamically managing a small set of Virtual Experts based on activation locality and hard- ware profiling, reduces memory consumption and latency, but sacrifices accuracy. These works do not take into ac- count the characteristics of semantic space, especially the cluster characteristics of high-level semantics and in-depth investigations.
Oracle-MoE Routing at Inference The inference stage can be divided into prefill and decode stages. The prefill stage is the same as the training stage routing. In the decode stage, Oracle-MoE first decides which semantic group the coming token belongs to, and updates the semantic group embedding with the coming token. Since KV cache is widely adopted strategy (Yuan et al., (024)), this does not introduce memory overhead. Then we dispatch it to the expert corresponding to the oracle cluster of its semantic group. In our experiments, there are often fewer than 5 semantic groups in an input session within a length of 1024, and semantic groups from the same session are likely to belong to the same oracle cluster. So, the oracle space routing preserves the semantic locality of input tokens and yields a low expert variation, contributing to low expert swapping latency.
5. Experiments
5.1. Settings
Platform Since mainstream mobile phones(like Android and Apple and NPU manufac- turers (c.g.. Apple, Hisilicon, Qualcomm, samsung) do not provide commercial APIs for low-level GPU memory operations, we adopt NVIDIA Jetson Xavier NX as our experimental platform. The NVIDIA Jetson Xavier NX is equipped with a 384-core NVIDIA Volta architecture GPU with 8 GiB of GPU memory and an estimated 21 TOPS AI computing power.
Hardware
4. Related Works
Mixture of Experts The Mixture of Experts (MoE) is fundamental model in machine learning (Jacobs et al.) [991 Jordan& & Jacobs| 1994) and an instance of the conditional computation framework where different experts are respon- sible for different inputs. To increase the model's capacity to deal with complex data, (Eigen et al. (013) extended the MoE structure to deep neural networks and proposed deep MoE model composed of multiple layers of routers and experts. (Shazeer et al. 2017) simplified the MoE layer by making the output of the gating function sparse for each example, which greatly improves the training stability and reduces the computational cost. Since then, the MoE layer with different base neural network structures (Shazeer et al.
Models We mainly compare Oracle-MoE with Switch Transformer (Fedus al. 2022). a representative token- level MoE architecture. Experts are loaded demand under our experiment settings. We use models containing m MoE layers with n experts each and with a total parameter of p. denoted as n * m(p). In our experiments, we use models of 2*4((9,M) 4*8(2MMM), 8* 16(729M) and 9*24(2.06B)) Detailed model configurations are in Appendix B
6

---

Oraclc-MoE: Locality-preserving Routing in the Oracle Space for Memory-consttrineed Large Language Model Inference
Oracio Xol Expert Activation Layer
Oracie-Stt Expert Activatiin Layer 11
Switch Expect Activation Lager
Switch Expert Activation Layer 10
Di 88
192
154
20
190
150
240
250
(a)
(b)
(c)
(d)
Figure 6. Expert activation results of (8+ 16)729M models. Switch Transformer, almost every consecutive tokens activate different experts, and nearly all experts are activated, demanding frequent expert swapping. Whereas Oracle-MoF. requires only few expert swappings as the token generation continues.
(192M)
8 (295M))
8 16 (729M)
9 21 (2.068)
14
1120
EIPG
JIPU LAU
LMU
LEU
Soop
Swag
SnoppMall
60
Saeghiitt
Oradio
Dracte-Mot
Gracie
Full Model
20-
Mustel
Pull Medal
1:2
2
Xemory TrolgeneG))
BuigetecBt
Memory Dudgee:C3)
Momory DudortiCD)
(a)
(b)
(c)
(d)
Figure 7. Memory-Latency curve of models of different sizes, where mpp) denotes a model consisting of MoE Layers with experts each, and in total parameters. Full model refers to a model with the same number of activated parameters.
Swapping Strategies As introduced in Section I swap- ping strategies assign a certain priority to each expert so that lower-priority experts can be swapped out first. We evaluate both our model and the Switch Transformer with First-In-First-Our (FIFO), Least Recently Used (LRU), and the strategy SwapMoE (Kong et al., 2024p. FIFO swaps out the expert that is loaded first, and LRU swaps out the experts that have not been used for the longest time. In SwapMoE(Kong et al. 2024). experts are weighted by their frequency, magnitude, and input tokens. However, since our model eliminates the requirement of expert swapping, different strategies don't make a big difference in our model. So, we report the average of 3 strategies on our model.
for a given memory size. A larger memory provides models with redundancy to store temporarily unused experts and mitigate the penalty of expert activation variation. c) First token Latency measures the time before the first token is generated after the users provide the input. This is also an important metric for user experience.
5.2. Results
Expert Activation Figure illustrates the expert activa- tion of two models. During consecutive auto-regressive generation passes, our method shows a lower expert activa- tion variation, where expert swapping is only triggered after hundreds of tokens are generated. In Switch Transformer, expert activation changes frequently.
Data Workload Models are pretrained on Openweb- Text (Komatsuzaki) 2019), which is one of the pretraining datasets of GPT2 (Radford et al. 2019). We primarily use downstream tasks of 3 types: question answering, clas- sification, and summarization. For QA tasks, we adopt
Memory-Latencr Figure the memory-latene curve for methods of different sizes. The result of our model is reported as the average of different strategies. As is shown in the figure, while with a small-sized model, the latency is acceptable, the case becomes worse rapidly as the model size gets larger. For the 8* ##6(729M) model, even though only expert is allowed for each layer (which is about only 25% of the full-size memory), our method introduces only 3s additional latency compared with the full-size inference. Whereas Switch Transformer with FIFO, LRU, or Swap- MoE load-on-dee -nd-mnnddeand strategy introduces inevitable latency, increasing latency by up to 2000% compared to full-size memory inference. When the memory budget increases by up to 50% of the full-size model, the latency of the Switch Transformer with different strategies is still unacceptably high, whereas our model does not introduce latency.
Trivia QA (Joshi et al., 2017). For classification, we adopt GLUE (Wang et al./ 2019). MAG (Sinha et al. 2015) and Sci-Cite (Beltagy et al. 2019). We also use XSum (Narayan let al. 2018) for summarization tasks. In experiments, we always kept the batch size equal to I, which is the real sit- uation when running on edge devices like mobile phones, processing one user request at time.
We adopt mainly 3 evaluation metrics. a) Expert Activation, to evaluate the variation in expert activation of different models. b) Memory-latency curve measures the average time the model takes to process a single data point
Metries
7

---

Oraclc-MoE: Locality-preserni Routing in the Oracle Space for Memory-constrained Large Language Model Inference
Size
Model
TrivialQA(F1)
GLUE(Acc.)
MAG(Accc)
Sci-Cite(Acc.)
Xsum(Rouge-1)
Avg.
Switch
27.00
62.25
20.00
30.83
13.55
30.73
195M
Ours
26.72
62.86
18.33
32.50
13.36
30.75
Switch
30.10
64.76
22.67
31,25
14.62
32.68
295M
Ours
30.31
64.66
22.17
33.75
15.53
33.28
Switch
35.08
68.33
25,29
35.00
15.62
35.86
729M
Ours 35.56
68.00
25.50
36.67
16.05
36.35
Switch 46.06
77.75
31.33
46.75
16.77
43.73
2.06B
Ours
46.96
78.00
30.67
47,50
17.35
44.09
Table 1. The zero-shot performance of models of different sizes on different tasks. Metrics reported depend on specific tasks. Our method does not pose drawback to model performance on downstream tasks.
Model(Strategy)
First token latency(s)
Oracle-MoE Lalency Breakdown
Switch((SwapMon Lalency Breakdown
Switch(FIFO)
22.395
0.5-
0.5
Switch(LRU)
23.428
0.0
0.
10
Switch(SwapMoo)) Oracle-MoE
12.767
FIFO Latency Breakdown
Switch(FIPO) Latency Breakdown
4.910
1.0
0.5
0.5
Table 2. First token latency of 765M models on different archi- (ectures()trategiess) under 50% full-size memory. The memory budget only influences SwapMoE since it uses offline statistics to determine which expert to load first.
0.0
0.0
2
10
Avuiibble (snney(G))
Asuiibble Menory((25)
Other Compating
Expert Computing
Expert Swapping
Figure Latency composition our proposed Oracle-Mor and Switch Transformer equipped with different swapping strategies. Our model reduces the percentage expert swapping latency and thus reduces the overall latency.
5.3. Overhead Analysis
Performance on Downstream Tasks Although designed for edge-oriented scenarios, our model does not sacrifice performance on downstream tasks for edge-deploy inference latency. As is shown in Table group gating method shows similar task performance, in some tasks even surpasses, the widely accepted token- level gating MoE models. We believe this attributable to our proposed semantic group-level routing strategy. This setup allows each expert focus on subset of high-level semantics rather than requiring every expert to learn all possible high-level semanties present in their target tokens, thereby reducing redundancy among experts.
proposed semantic
The latency composition breakdown in Figure gives detailed visualization of the above results. It can be observed that even with the most limited memory budget, the latency introduced by expert swapping in our model only contributes to 50% of the overall latency. In token-level routing, such as the switch transformer, expert swapping contributes to more than 99% of latency.
First token latency Our model activates fewer experts for single input, so that only 1 or experts are needed for the prefilling stage. So, our model requires only one-time expert loading during the prefilling stage. However, existing token- level MoE methods still need to swap experts during the prefilling stage, leading to worse first token latency. Among the three expert swapping strategies, FIFO does not help at the prefilling stage, and the LRU strategy needs warm- up stage to decide on frequently used tokens. SwapMoF, however, uses off-line statistical information to decide the loaded expert at the beginning of inference, thus resulting in lower first token latency than baselines, but still not as good as ours.
Training Stage Overhead Our approach differs from ex- isting token-level MoE in that it includes one-time cluster analysis after the warm-up phase and cluster routing in each forward pass. In our experiments, with a sample size of 8192, the wall clock time for clustering analysis per layer is approximately 4 min, which is negligible compared with tens of hours of pretraining. For routing in each pass, token- level MoF is equivalent to performing a matrix multiplica- tion. It requires lc-4 seconds, whereas our now-dimensional cluster Euclidean distance computation requires three matrix multiplications and a square root operation. Thanks to the low-dimensional semantic space, the final wall-clock time
8

---

Oraclc-MoE: Locality-preserving Routing in the Oracle Space for Memory-consttrineed Large Language Model Inference
of our routing is 2.5e-4 seconds, which is also negligible pass taking 3.5
References
compared the single seconds.
Artetxe, M., Bhosale, S., Goyal, N., Mihaylov, T., Ott, M., Shleifer, S., Lin, X. V., Du, J., Iyer, S., Pasunuru, R., Anantharaman, G., Li. X., Chen, S., Akin, H.,, Baines, M., Martin, L., Zhou, X., Koura, P. S., O'Horo, B., Wang, J., Zettlemoyer, L.,, Diab, M., Kozareva, Z., and Stoyanov, V. Efficient large scale language modeling with mixtures of experts. In Goldberg, Y., Kozareva, Z., and Zhang, Y. (eds.), Proceedings of the 2022 Conference Empirical Methods in Natural Language Processing, pp. 1 11699- 11732, Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Linguistics. doi: 10. 10.18653w1/2022.mmmlpp 804. URL httpss:// ac wwwpntholoyy..rroooooo
5.4. Expert Prediction- Optimization
The Oracle-MoE has significantly diminished the neces- sity for expert swapping, thereby fundamentally reducing latency. However, the load-on-demant strategy still inher- ently suffers from another limitation: it cannot decide which expert to load for layer before that layer is reached. Conse- quently, we propose to predict deep layers' expert activation at shallower layers, enabling inferring current tokens and loading experts synchronously. Specifically, we use the em- beddings of the first layer to predict the expert activation in the following layers. Experimental results show that Oracle- MoE reaches an expert prediction accuracy of 85% to 95%, whereas Switch-Transfornn Transformer- -liee token-level routing struc- ture only has an accuracy of 40% to 60%. Employing this can further reduce the expert loading latency of Oracle-MoE by 10% to 15%. The underlying reasons for this expert pre- dictability will be left for future investigation.
emn p-
Beltagy, I., Lo, K., and Cohan, A. Scibert: Pretrained language model for scientific text. In MMNLP,2019
Cai, W., Jiang, J., Wang. F., Tang. J., Kim, S., and Huang. J. A survey on mixture of experts. ArXiv, abs/2407.06200 2024. URL https: www.aaissssaannncccllssssss org
Chen, T., Zhang. Z., JAISWAL, A. K., Liu, S... and Wang, Z. Sparse moe as the new dropout: Scaling dense and self- slimmable transformers. In The Eleventh International Conference on Learning Representations, 2023.
6. Conclusion
This paper introduces Oracle-MoE, an innovative MoE ar- chitecture designed to optimize edge deployment by ad- dressing latency issues inherent in token-level routing. By incorporating Oracle-Space Saaee--pace-Based Based routing mechanism, Oracle-Mofs reduces expert activation variation and mini- mizes memory swapping overhead. Extensive experiments validate that Oracle-MoE achieves lower latency without in- troducing any drawback on downstream tasks or pretraining computation, which is vital for the on-device deployments of LLMs. Future work could extend this approach other conditional computation frameworks and explore its appli- cability across diverse hardware environments.
Cheng, H., Zhang, M., and Shi, J. Q. A survey on deep neural network pruning: Taxonomy, comparison, anal- ysis, and recommendations. IEEE Transactions on Par- lern Analysis and Machine Intelligence, 16:10558-10578, 2023. URL https: //api. wew.aiicchoaaamcce.e.a...a.. org /CC.poo1D:26008811
Chi, Z., Dong, L., Huang. S., Dai, D., Ma, S., Patra, B., Sing- hal, S.,, Bujaj, P., Song. X., Mao, X.-L., Huang. H., and Wei, F. On the representation collapse of sparse mixture of experts. Advances in Neural Information Processing Systems, 2022.
Impact Statement
This work aims to deliver insights to the Machine Learning community and will not lead to any direct societal conse- quences. While is associated with LLMs that may output misleading or harmful content, such issues are outside the scope of this work and will not be particularly specified here.
Clark, A., De Las Casas, D., Guy, A., Mensch, A., Paganini, M., Hoffmann, J., Damoc, B., Hechtman, B., Cai, T., Borgeaud, S., Van Den Driessche, G. B., Rutherford, E., Hennigan, T., Johnson, M. J., Cassirer, A., Jones, C., Buchatskaya, E., Budden, D., Sifre, L., Osindero. S., Vinyals, O., Ranzato, M., Rae, J., Elsen, E., Kavukeuoglu, and Simonyan, K. Unified scaling laws for routed language models. Proceedings of the 39th International Conference on Machine Learning, 2022.
Acknowledgements
Yujiang Wang was supported by a Basic Research Program of Jiangsu (BK20240414) 14) and Leadership Talent Pro- gram (Science and Education) of Suzhou Industrial Park (KJQ2024204).
Clark, K., Khandelwal, U., Levy, O., and Man- ning, C. D. What does bert look at? an anal- ysis of bert's attention. In BlackboxNLP@ACL 2019. URL htttps://api... http:tts///p..eemmtttttooooe wewanticschola.. prg//Corppus1::18448777
9

---

Oraclc-MoE: Lorelity-reseevingg Routing in the Oracle Space for Memory-consttained Large Language Model Inference
Dai, D., Dong, L.., Ma, S., Zheng, B., Sui, Z.,, Chang. B., and Wei, F. Stablemoe: Stable routing strategy for mixture of experts. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguisties (Volume I: Long Papers), 2022.
Joshi, M., Choi, E., Weld, D., and Zetrleooerr L. triv- iaqa: A Large Scale Distantly Supervised Challenge Dataset for Reading Comprehension. arXiv e-prints, art. MxXiv:1705.00 100500351,,0111 2017.
Kamath, U., Liu, J., and Whitaker, J. Deep learning for NLP and speech recognition, volume 84. Springer. 2019.
Dai, D., Deng, C., Zhao, C., Xu, R., Gao, H., Chen, D., Li, J., Zeng, W., Yu, X., Wu, Y., et al. Deepscekmoe: Towards lan-
Komatsuzaki, A. Openwebtext corpus, 2019. https
ultimate expert specialization in mixture-of-exxpeeee guage models. arXiv preprini arXiv:2401.06066 2024.
//github. com/ sko akoda oopedbb//pewwb..x.tt..
Kong, R., Li, Y., Feng, Q., Wang, W., Ye, X., Ouyang, Y., Kong. I.., and Liu, SwapMoF: Serving off-the- shelf MoE-based large language models with tunable memory budget. In Ku, L.-W., Martins, A., and Srikumar, V. (eds.), Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume I: Long Pupers), pp. 6710-6720, Bangkok, Thailand, August 2024. Association for Computational Linguistics. doi: 10. 10.18865//1/20244.cc aclantholooy..orr//
Dauphin, Y. N., Fan, A., Auli, M., and Grangier, D. Lan- guage modeling with gated convolutional networks. In International conference on machine learning. PMLR, 2017.
Eigen, D., Ranzato, M., and Sutskever, Learning fac- tored representations in deep mixture of experts. arXiv preprint arXiv: 1312,4314, 2013.
URL [tttpts:::/ long. 363
Farjas, J. Clique cover.
httpp:://iithu.. c
Far [arjaajj/LCiuuueoveer 2018.
Kovaleva, O., Romanov, A., Rogers, A., and Rumshisky, A. Revealing the dark secrets of bert. ArXiv, abs/1908 08593, 2019. URL https: http:/::/..ppp.p.aaaa..aaa. o org/ Peg//oorpuuIDD:2115551.
Fedus, W., Zoph, B., and Shazeer, N. Switch transformers: Scaling to trillion parameter models with simple and ef- ficient sparsity. Journal of Machine Learning Research, 2022.
semanticscholes
Kowsher, M., Prottasha, N. J., and Yu, C.-N. Does self- attention need separate weights in transformers?, 2024. URL https:: wrw.wvorrggooo abs 2412. 2412.00359]
Gavril, F. Algorithms for minimum coloring, maximum clique, minimum covering by cliques, and maximum independent set of a chordal graph. SIAM Journal on Computing, 1(2):180-18)) 1972.
Lepikhin, D., Lee, H., Xu, Y., Chen, D., Firat, O., Huang. Y., Krikun, M., Shazeer, N., and Chen, Z. Gshard: Scaling giant models with conditional computation and automatic sharding. In International Conference on Learning Rep- resentations, 2021.
Huang, H., Ardalani, N., Sun, A., Ke, L., Lee, H.-H. S., Sridhar, A., Bhosale, S., Wu, C.-J,, and Lee, B. Towards moe deployment: Mitigating inefficiencies in mixture- of-expert (moe) inference. ArXiv, abs/2303.06182, 2023a. URL httpp::/atp.......e.e wtwppp..ematttccccc....o htg//oopputD:225445000
Lewis, M., Bhosale, S., Dettmers, T., Goyal, N.,, and Zettle- moyer, L. Base layers: Simplifying training of large, sparse models. Proceedings of the 38th International Conference on Machine Learning, 2021.
Huang. H., Ardalani, N.. Sun, A., Ke, L., Lee, H.-H. S., Sridhar, A., Bhosale, S., Wu, C.-J, and Lee, Towards moe deployment: Mitigating inefficiencies in mixture-of- expert (moe) inference, 2023b. URL htt.::tttts/aattt
Liu, G., Li, C., Zhao, J., Zhang, C., and Guo, M. Clus- terkv: Manipulating Ilm kv cache in semantic space for recallable compression, 2024. URL | htttps://axxiitit 03213
org/ abs/ 2303. 06182
Jacobs, R. A., Jordan, M. I., Nowlan, S. J., and Hinton, G. Adaptive mixtures of local experts. Neural computation, 3, 1991.
Narayan, S., Cohen, S. B., and Lapata, M. Don't give me the details, just the summary! topic-aware convolu- tional neural networks for extreme summarization. ArXiv, abs/1 8188808..4444 2018
Jain, A. K. k-means. Pattern Recognit. Lett. 31:651-666, 2008. URL http:p/://pp..smmattisssee.e org/ htt///ooosD:::11122000
Data clustering: 50 years beyond
Poemner, N., Waltinger, U., and Schütze, H. Sentence meta-embeddings for unsupervised semantic textual similarity. In Annual Meeting of the Association for Computational Linguisties, 2019. URL https htt.s//mmniiicccoooo.ooooooooaoe 207852340
Jordan, M. I. and Jacobs, R. A. Hierarchical mixtures of experts and the em algorithm. Neural computation, 6, 1994.
10

---

Oraclc-MoE: Locality-peseeving Routing in the Oracle Space for Memory-consttainedd Large Language Model Inference
Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., and Sutskever, I. Language models are unsu- pervised multitask learners. 2019. URL https:
Team, Q- Qwenl.5-moe: Matching 7b model per- formance with 1/3 activated parameters", February 2024.
wew.ssiscnnnncccla..... Org/Corpuss
wwen-moe/
1 160025533
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L.,, and Polosukhin, I. Atten- tion is all you need. In Advances in neural information processing systems, 2017.
Raffel, C.,, Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., and Liu, P. J. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research, 21 (140):1-67, 2020.
Vig, J. and Belinkov, Y. Analyzing the structure of atten- tion in a transformer language model. arXiv preprint arXiv: 042404284, 2019.
Rajbhandari, S., Li, C., Yao, Z., Zhang, M., Aminabadi, R. Y., Awan, A. A., Rasley, J., and He, Y. Deepspeed-mon moo:: inference and training to power next-generation ai scale, 2022. URL https://
Wang, A., Singh, A., Michael, J., Hill, F., Levy, O., and Bowman, S. R. GLUE: A multi-task benchmark and anal- ysis platform for natural language understanding. 2019. In the Proceedings of ICLR.
Advancing mixture-of-expee
arxiv. org/
05596
Xu, K., Hu, W., Leskovee, J., and Jegelka, S. How power- ful are graph neural networks? ArXiv, abs/1810.00826 2018. URL https://api: www.nticschoollaa....ao org/ http/CoruuSD:::23300
Schmidt, D. A survey of singular value decomposi- tion methods for distributed tall/skinny data. 2020 JEEE/ACM 11th Workshop on Latest Advances in Scal- able Algorithms for Large-Scale Systems (SculA), pp. 27-
34. 2020. URL htA.attpps//pp...emaatcceeeeee htg/CorpusID:2214422 semanticscoolar
Xue, F., Zheng, Z., Fu, Y., Ni, J., Zheng, Z., Zhou, W., and You, Y. Openmoe: An early effort open mixture- of-experts language models, 2024. URL https://
Shazeer, N., Mirhoseini, A., Maziarz, K., Davis, A., Le, Q. Hinton, G., and Dean, Outrageously large neural networks: The sparsely- gated -ifturr-of-experts layer. arXiv preprint orXiv:1701 006538,, 2017.
8 arxiv. axxiv.org/a hbsp/bb//// 2202/240..011
Yang, A., Yang, B., Hui, B., Zheng, B., Yu, B., Zhou, C., Li, C., Li, C., Liu, D., Huang, F,, Dong, G., Wei, H., Lin, H., Tang, J., Wang, J., Yang, J., Tu, J., Zhang, J., Ma, J., Xu, J., Zhou, J., Bai. J., Hc, J., Lin, J., Dang, K., Lu, K., Chen, K., Yang. K., Li, M., Xue, M., Ni, N., Zhang, P., Wang, P., Peng. R., Men, R., Gao, R., Lin, R., Wang. S., Bai, S., Tan, S., Zhu, T., Li, T., Liu, T., Ge, W., Deng, X., Zhou, X., Ren, X., Zhang, X., Wei, X., Ren, X., Fan, Y., Yao, Y., Zhang, Y., Wan, Y., Chu, Y., Liu, Y., Cui, Z., Zhang, 7.,, and Fan, Z. Qwen2 technical report. arXiv preprint arXiv:2407.10671, 2024.
Shen, L., Wu, Z., Gong, W., Hao, H., Bai, Y., Wu, H., Wu, X., Xiong, H., Yu, D., and Ma, Y. Se- moe: A scalable and efficient ofxture-of-experts distributed training and inference system. CORR, abs/2205.10034 2022. URL 48550/arxi www.:rxww..........m 2205. 10034.
orr//00
Sinha, A., Shen, Z., Song, Y., Ma, H., Eide, D., Hsu, B.- P., and Wang, K. An overview of microsoft academic service (mas) and applications. In Proceedings of the 24th International Conference on World Wide Web (WWW '15 Companion), pp. 243-246, New York, NY, USA, 2015. ACM. doi: 10. 145/27449908.
Yi, R., Guo, I.,, Wei, S., Zhou, A., Wang, S., and Xu, M. Edgemoe: Fast on-device inference of moo-based large language models. arXiv preprint urXiv:2308. 14352, 2023.
Yuan, Z., Shang, Y., Zhou, Y., Dong, Z., Xue, C., Wu, B., Li, Z. Gu, Q., Lee, Y. J.. Yan. Y.. Chen, B., Sun, G., and Keutzer, K. Llm inference unveiled: Survey and roofline model insights. ArXiv, abs/2402.16363. 2024. URL https://aa/ www.nteiiiio..m semanntccccolar org/ SNtrpusI::2680222:
Soltanolkotabi, M., Elhamifar, E., and Candès, E. J. Robust subspace clustering. ArXiv, abs/1301.2603, 2013. URL https: wwtpppssssmmtticccsssssss arg/Corpin ID: ID:
Takahashi, K. and Bollegala, D. Unsupervised attention- based sentence-level meta-embeddings from contex- tualised language models. In International Con- ference on Language Resources and Evaluation, 2022. URL https::/app....... www.api.sssnnniscco........ orre/Corrpusin:224227999
Zhang, 7., Xia, Y., Wang, H., Yang, D., Hu, C., Zhou, X., and Cheng, D. Mpmoe: Memory efficient moe for pre- trained models with adaptive pipeline parallelism. JEEE Transactions on Parallel and Distributed Systems, 35(6): 998-1011, 2024. doi: 10. 000//PPD...0223335569

---

Oraclc-MoE: Locelity-reserving Routing in the Oracle Space for Memory-consttrined Large Language Model Inference
Zhou, Y., Lei, T., Liu, H., Du, N., Huang, Y., Zhao, V., Dai, A. M., Chen, 7., Le, Q. V., and Laudon, J. Mixture- ofexperts with expert choice routing. Advances in Neural Information Processing Systems, 2022.
Zhu, X., Li, J., Liu, Y., Ma. C., and Wang, W. A survey on model compression for large language models. Transactions of the Associa- tion for Computational Linguistics, 12:1556-1577, 2023. URL org/ org/ http//opuu:::0000011..
Zuo. S., Liu, X., Jiao, J., Kim, Y. J., Hassan, H., Zhang, R., Gao, J., and Zhao, T. Taming sparsely activated transformer with stochastic experts. In International Con- ference on Learning Represemations, 2022. 2022
12

---

Oracle-MoE: Locality-peeserving Routing in the Oracle Space for Memory-constrained Large Language Model Inference
A. Proof of Analysis
where [|MI]| denotes the operator norm (spectral norm) of matrix M3 MM This bound reflects how the transformations MQ and Mr affect the original inner product (X, Y))
A.1. Property of semantic element
In this section, we delve into the detailed analysis of the approximate attention mechanism, specifically focusing on the impact of the inner product between query (WQ) and key (Wx) matrices.
Furthermore, although the above discussion is based on an approximate attention mechanism and focuses on the properties of token components in the Hy space, Figure 2 shows that tokens with the same identity (i.e., semantically similar tokens) tend cluster together in low-dimensional representations that preserve relative distances. But this clustering behavior also suggests that tokens with the same high-level semantic meaning have small relative distances in the semantic component space. Therefore, for tokens within the same subgraph or semantic context, we make the following assumption:
OK Rank Subspoce Angle
200
30
175
501
150
70
125
60
34
40
75-
Eifectiin Bank
20
Pitletivw Bank
Subsgoce Augle
10
Assumption 2. For tokens t; in the same subgraph, or Semantic Group Sj. their components in semantic space Hss St, there exist r, s.f. S; E Buss(cr.)
0
2.5
2.5 10.0
12.5
Law
Figure 9. Effective rank of matrices We and Wx in different lay- ers' self-attention and their subspaces angle. The space determined by Wo and Wx is very similar.
Assumption 3 (Uniform Distribution in Semantic Space). Tokens in the same semantic group 8j are uniformly dis- tributed in the semantic space with an expected value the center:
The attention score computation can be viewed as mapping from two token representations to scalar value:
E[s;] = Cj: Vi E
A: 7l x H - R,
A.2. The advantage of mean embedding of subgraph
In Figure 9 We and Wx exhibit similar ranks and mini- mal angles between their subspaces. Combined with the empirical success of shared QK matrices in transformers (e.g.. Kowsher et al. (2024)), we propose that Wo and WK k can be decomposed into shared projection followed by subspace-specific transformations. We therefore propose:
Given the definition of semantic groups, we now make an assumption about the distribution of token identity infor- mation uj which refers to the unique characteristics within each subgraph, distinguishing one token from another be- yond their shared high-level semantics:
We= MaP, =
Assumption 4 (Normal Distribution in Token Identity Space). Token identical information follows normal dis- tribution within each subgraph Sy:
where P: His shared projection matrix mapping tokens common subspace Hs of dimension (with S di)( and Mq. Mk Hss Hs are full-rank linear transformations within Hs.
~ N Eng,
Vi E Sj.
Given this understanding:
Then we have The variance of the semantic representation of token t; is defined as:
Assumption (Semantic Subspace). The token represen- tation space H. can be decomposed into a low-dimensional semantic subspace H. and its orthogonal complement H.S:
Var(z,.) Var
Σ t)
|S(I)
H Hs b H
tES())
Let the neighborhood size be n |S())), then:
Followed by the analysis and the experiment results in Sec- tion| 2.2 although we recognize that Wq * WK. meaning Me F MR, under our analysis, where both Mq and Mr are different full-rank linear mappings from Hs to Hs. we can derive a relatively symmetric bound for the inner products involving these mappings. Specifically, for any non-zero, unit vectors X, Y e Hs. we have:
1
1
ZS(()
t;
(s: + +u)
(2)
n
(ES(1)
(ES(()
It can be decomposed into two parts:
KM MR < [|M[| (X,Y)) (MMYYMKXX' < [|M[| [XXX)]-
2.((0)
n
1
Σ S;
n iES())
Σn
(3)
(ES())
Mcan semantic part
Mean of identity part
13

---

Oraclc-MoE: Locality-pesseving Routing in the Oracle Space for Memory-constrained Large Language Model Inference
Variance decomposition:
Under uniform cluster assumptions, the density-to-volume ratio simplifies to a constant, yielding:
Var(zs(e)) Var 1
Es s:
+ Var ( Σuu
T
T
77
(4)
De
~
(s(1)) - zs(--1)
+ 2Cov
Siz 71
11
A.4. Proof of Theorem 1
Then we have the analysis of each component:
First, have this lemma:
1. Covariance term: Given the orthogonal subspace decom- position H, T. H.g the semantic part s; and the identity part u, are independent, hence the covariance is zero:
Lemma 5 (Norm Comparison with Additive Threshold). For n < m, [| + m < Z| holds with probability approaching I as - 0o, where m and are fixed constants.
Cov L Σ
1 Σ
u1
0
(5)
Proof. Let {2i]< be iid. d-dimensional Gaussian vectors with:
2. Variance of the semantic part: According to As- sumption 3 (uniform distribution), within the same sub- graph, s; are independently and identically distributed with E[s,] = C.. Let Var(s;) = Eg, then:
E|ay] 4 Rd
Cov(x;) o² where la the d d identity matrix.
Var - Σ
Var(s;) = Σg
Let (li,,..,, .,. knj be a uniformly random subset of [l....,m) m_] (without replacement) with n m. Define:
8:
(6)
n2
1 m
m
Variance of the identity part: According to Assump- tion (normal distribution), within the same subgraph, Nuu En) and they are independent, then:
Tky 2 x1 - 22.
j=1
u;
For Y:
Var
1 n
Σ
uu
Var(u.) Σ
B[y] = 0,
n
Cov(Y) 0²
la:
m
thus we have
1
E|jlY 12] do² n
Eg + Ei Var(tt)
Var(zs(e))
(8)
m
n
For Z:
A.3. Approximation Analysis of Oracle CSD
B[Z] = 0,
The expert assignment change de; is defined as the sym- metric difference between consecutive expert sets:
Cov(Z) (2) 20
B|||Z| [2] = tr(Cov(Z)) = 2do
Ae, lee \\e er-1l
For n < m:
For simplicity, assume each token activates a single expert, so Ae, T[(: # e--1) (0 or 1).
<2
m
- 1
When (5(()) and (-((-1) reside in the same cluster, Ae, = O. When they lie in different clusters, de;- = 1. Let B denote cluster boundaries in the oracle space. The probability of crossing B between 1s(--1)||-
B||l|| = do²
n1 m
< 2do² = 1E11 1
and increases with ||zs()
Define the modified gap Wx = k21l212 - || We analyze:
For small displacements, we approximate the discrete Noundary-crossing event by the continuous embedding dis- placement:
P P(IY + +m < K| = P ||||2 2m
Cluster density at B
+ m² < x221112)
Ae, Rs [Zs(0) - 2 llll
Cluster volume
14

---

Oracle-MoE: Locality-preserving Routing in the Oracle Space for Memory-constrained Large Language Model Inference
Table Hyperparameters Models
Hyperparameters
195M MoE
295M MoE
729M MoE
2.06B MoE
Attention heads
12
12
12
16
Transformer layers
12
12
12
24
MoE layers
2
4
Expert Number
4
8
16
32
Activated Expert Number
1
1
1
Hidden dimension size
768
768
768
1024
Dropout
0.1
0.1
0.1
0.1
Attention dropout
0.1
0.1
0.1
0.1
Sequence length
256
256
512
1024
Batch size
320
320
160
80
Learning rate decay
Cosine
Cosine
Cosine 2e-4
Cosine
Maximum Learning rate
4e-4
4e-4
le-4
Activation inconsistency
DeepSeek
Qwen
Switch
Oracle
1st 1/4 layers avg
80.84
81.56
69.20
6.03
2nd 1/4 layers avg
65.35
71.04
64.87
4.82
3rd 1/4 layers avg
70.68
75.37
53.36
4.20
4th 1/4 layers avg
76.61
77.16
75.44
5.11
Table Activation inconsistency comparison across layers.
Using Chebyshev's inequality for Wx:
Therefore:
B[WK] = K²B - B[ = do²
2K
-
Var(W
0
(7) +0
Var(Wk) = K'Var([12])) Var(I
(independence),
lim P ((I + < K] 1.
1 1
= do + 2do
Note:we also give the analysis of the impact of parameter constrains: when approaches m:
Set
R[WK] 2m /m B[|
m
Substituting
1
[[W] x 2K2 m
2
1
Var(W) "k 4m.
6x = do²
2K
2m do2
m
The probability bound becomes:
Applying Chebyshev's 's inequality:
Var(Wk)
-o + 4m2 P(W > 0) d2K - m2 2
PWW 2 (x) (x) 21
Thus:
For large dominant terms yield:
P P( m K1
8K+ 2 4K
do + 2do (1 1)²
m)
P(W > 0121-11-1
= d'
d (2K². 1 11 2m / do² (1 11 m²
So we make an extreme assumption about the left and right components in semantic space, that is
The dominant terms scale as:
Var(Wx) = Odd), "k (d))
(t, t..1,))ll 1) 0, (Zs(t) =s(t-1))HHs 2r
15

---

Oracle-MoE: Locality-presse preserving Routing in the Oracle Space for Memory-constrained Large Language Model Inference
(s(()) 22(1-1)11 - (s((1) - zs(1-1))
H2s
Tokes Level Cluster
+ (2s(1) - =2s(e_1))
Technologhe
A1.s
7.5
5.0
= (2s(1) - 25((-11)
+ 2r
2.5
114
0.-
16
Hoalth
[[ - - 1 || ll(t; - 1z-1)]l
Since (t. bz-1)| (zs(t) Zs(t-1)) 17t follows
Figure IО, UMAP visualization of embedding space in DeepSeekMoE- 16B from layer on Wizard-of- -izzad-of-Wikipediia datasets. Left: Tokens tend to cluster according to token-identity semanties. Right: Tokens from the same sequence are colored the same. They share similar semantics and stay closer cach other in each token cluster.
L let 2r, C(W,, g: k) to be K in Lemma with existing d mentioned in lemma
C(W,, g. (ol(tt-tx-1))
(is())) -s((-1)) Hs + 2r|
Thus we have
t(-))l (2(()) zs(1-1>)|
Telwe- Clucker
Chatc Fleanor
B. Experiments
B.1. Pretraining Hyperparameters
Table
3
lists the hyperparameters used in our experiments.
Chatcon Computer
B.2. Activation Inconsistency and Semantic Locality in Existing MoEs
Figure 11. UMAP visualization of embedding space in DeepSeekMoE- 16B from layer on Wizard-of-Wikipedii Wikipedia datasets. Left: Tokens tend to cluster according to token-identity semanties. Right: Tokens from the same sequence are colored the same. They share similar semantics and stay closer cach other in each token cluster.
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 in Table Existing MoEs show strong temporal activation inconsistency within all layers, while Oracle-MoF: reduces this.
90.54 times. This is because in human natural language, it takes at least dozens of tokens to express complete meaning, our method still benefits from such "abrupt" semantic locality.
Experiments with DeepSeekMooEIIg 16B and Qwenl.5-MM Mo-E- A2.7B on real chat Srinsets(Wizard-fff Wikipedaa and are shown in Figure 10 and Fig- ure II Semantic locality appears across different mod- els/layers/samples Semantic groups can still be distin- guished based attention score and obtained by our method, as shown in Figure 12 to Figure 15 It indicates the potential of Oracle-Mof being general purposs solution.
B.3. MoE with fine-grained experts
We train model following the setting of DeepSeekMoF. 16B but with fewer parameters(3B): 12 MoE layers with 64 routed experts each as baseline, where hidden size is set to 1536 and expert intermediate size is set to 1024. The top experts are selected for each token. Our method still achieved 75% latency reduction 2.5GB memory. Mean- while, 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 21.74, while the baseline score 21.22.
We tested scenarios where the topic changes frequently. randomly sample sentences from different datasets and combine them into whole sequence. We observed that our proposed oracle space can still distinguish semantic groups efficiently, both in models and public large MoE models, shown in Figure 16 to Figure 19 We also tested the expert activation variation of such highly diverse data with Oracle-Ma and switch- transformer. On average, in every 100 consecutive token generations, Oracle-MoE only changes 12.20 times while the switch transformer changes
16

---

Oracle-MoE: Localiy-presserinng Routing in the Oracle Space for Memory-constrained Large Language Model Inference
Adwestte
Altentiin Scoee
Artentioon Score
Attertion
7akm Tocke
Ateetias Soare
Alleelkus Sixen
ATweatinn Scar
Alteciis Sooce
22
Teles Baket
Thkum Ioiies
Tokee Jales
Trims Todes
Figure 12. Attention scores randomly sampled data of Wizard-of- Wikipeediauupperrandd Synthetic-Persona Chat(bottom) in DeepSeckMoB- 16B from layers 5,10,1,220.
Scnantic Group PecsoMiiy
Seneetic Uraup Estecking
Susmantic Greap Enttesiing
Secnurte Greap tushodiing
A
Willt oeeeeeda
Lucbot eel podo
Leer
Figure 13. Semantic groups obtained the Oracle-MoF method Wizard-of- -irkkidd--Wiiipeiiaa and Synthetic- Synthetic-Persona- across different DeepScekMoF-16R layers with semantic groups from the same sequence or user interaction are colored the same.
Alcctiie
Altenika Scoce
Attattio Scare
Atteztion
Antertter
Attestkon Stxen
Antentionn Scom
Attertion
2 C.
Trles Babes
Tale
Figure 14. Attention scores randomly sampled data of Wizard-of- wiineedd--iikkppcdapppppee) and MoB-A2.7B from layers 5, 10,11,2,0,
in Qwen1.5-
Sexweric Croup Becondiing
Sccantoo Unaep Entesking
Secsastle Crrep Rnteediing
Secwentte Creap Frshedilng
Syehuts- terhetse
Laye 13
Lape
-10
Figure 15. Semantic groups obtained by the Oracle-MoE method Srrard-offWikipedde and Synthetic- across different Qwen1.5 D.ven..-MME-- M-5-MoEEA.... A2.7B layers with semantic groups from the same sequence or user interaction are colored the same.
17

---

Oraclc-MoE: Locality-preserving Routing in the Oracle Space for Memory-consttrineed Large Language Model Inference
Altertaee
Allcetbon Soore
Attoiiao Susrr
Allectaon Soure
24
Iokes
Trum
24
Trams
Altention Scoce
Actertiinn Score
Atterdan
300
300
Tears
Figure 16. Attention scores on dierent diverse data(the top and bottom rows are different data) in 1eepSeekMoE-16B from layers 1,10,15,20.
Servcrtt Eirbooiing
Semwarte Inteeriter
Servanto Greep Brboodiing
Semurrtte Fnshediteg
76
04
Figure 17. Semantic groups obtained by the Oracle-MoE method on diverse data from DeepSeekMoEE166 layers 5,10,15,20 with semantic groups from the same sequence or user interaction are colored the same.
00
X0
30
30
Troo Scdee
Toons Indes
Trors hnbe
Trons Indes
Toles
Telars
Teles
Teiars Bakes
Figure 18. Attention scores different diverse data(the top and bottom rows different data) in (wenn.55Moo--22.75 from layers 5,10,15,20.
Ureap Erteoking
Groep Exboddiag
Gropp Eistcolnn
-35
75
300
125
13
Figure Semantic groups obtained by the Oracle-MoE method on diverse data from Owenl.5-MoB-AA... layers 5,10,15,20 with semantic groups from the same sequence or user interaction are colored the same.
18


Reviews:
[{"review_id": "21I7A63To6", "reviewer": "Reviewer_dPcM", "summary": "The paper introduces Oracle MoE, a novel Mixture-of-Experts (MoE) architecture designed specifically for memory-constrained inference on edge devices. The main idea is to replace conventional token level routing with an oracle-space routing mechanism that leverages semantic locality. By grouping tokens based on their high level semantic similarity (extracted via attention scores) and using compact semantic group embeddings (the “oracle space”), the method reduces the frequency of expert swapping and thus significantly lowers inference latency. Extensive experiments on GPT 2 based models across various sizes and downstream tasks show that Oracle MoE achieves state of the art latency improvements while maintaining competitive task performance.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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\n\n_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.\n\n_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.\n\n|Activation inconsistency|DeepSeek|Qwen|Switch|Oracle|\n|--|--|--|--|--|\n|1st 1/4 layers avg|80.84|81.56|69.20|6.03|\n|2nd 1/4 layers avg|65.35|71.04|64.87|4.82|\n|3rd 1/4 layers avg|70.68|75.37|53.36|4.20|\n|4th 1/4 layers avg|76.61|77.16|75.44|5.11|\n\n_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.\n\nExperiments 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.\n\n_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?\n\n_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.\nTheoretical 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.\n\n[A] A Survey on Word Embeddings: From Shallow Models to Deep Learning (Goldberg, 2017)\n\n[B] Deep Learning for NLP and Speech Recognition (Hinton et al., 2012)\n\n[C] Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer (Raffel et al., 2020)\n\n[D] Analyzing the Structure of Attention in Transformers (Kobayashi et al., 2020)\n\n_Q4_ There are some approximations in the derivation of oracle-MOE. It would be better to validate such an approximation in the experimental design.\n\n_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.\n\n|Embedding L2 Distance|[0.0,22)|[22,33)|[33,44)|[44,55)|[55,66)|[66,77)|[77,88)|[88,+∞)|\n|---|---|---|---|---|---|---|---|---|\n|Inconsistent activation per token|0.0|0.083|0.222|0.377|0.576|0.770|0.878|0.903|\n\n_Q5_ Ablation study on CSD and hyperparameter study of γ are missing. How does γ influence the proposed algorithm?\n\n_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.\n\n|Model Scale|2*4(192M)|4*8(295M)|8*16(729M)|9*24(2.06B)|\n|--|--|--|--|--|\n|our model's outperform baseline's γ by(larger better)|+0.02|+0.60|+0.49|+0.36|\n|our model reduces CSD by(smaller better)|-25.10|-53.94|-71.23|-85.41|", "claims_and_evidence": "Some claims are supported by clear and convincing evidence. However, there are some important claims are problematic. For example, \n\n(1) the temporal inconsistencies of inter-token expert activation are not clear. The authors only show visualization results for a specific sample. It would be better to measure the temporal inconsistencies for the whole dataset and different layers. \n\n(2) The key assumption for semantic locality, i.e., linguistic meaning between consecutively generated tokens is typically consistent, is not supported by any evidence. \n\n(3) 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?", "methods_and_evaluation_criteria": "The proposed method makes sense to reduce latency via eliminating massive swapping demands.", "theoretical_claims": "I don’t check the theoretical claims carefully.", "experimental_designs_or_analyses": "Yes. I check the experimental design and analysis. Although the authors provide memory-latency results and demonstrate the overall performance for QA, classification and summarization tasks, there are several important experiment are missing:\n\n(1)\tThe experimental result on the temporal inconsistency of inter-token expert activation is missing. It would be more convincing to provide a metric to quantitatively measure the temporal inconsistency and show how Oracle-MOE reduces this inconsistency.\n\n(2)\tThe semantic locality, as a key assumption, is not well explained. It would be better to provide some preliminary experiments to show the evidence in real datasets. It would be interesting to see the semantic locality across different models/layers/samples.\n\n(3)\tThere are some approximations in the derivation of oracle-MOE, such as CSD_token in Line 147, page 3, CSD_oracle in Line 231, page 5. It would be better to validate such approximation in the experiments design.\n\n(4)\tAblation study on CSD and hyperparameter study of $\\gamma$ are missing. How does $\\gamma$ influence the proposed algorithm?", "relation_to_broader_scientific_literature": "The key contributions can accelerate the MOE model inference in edge devices scenario.", "essential_references_not_discussed": "Yes.", "other_strengths_and_weaknesses": "Strengths:\n\n1.\tThe key idea of maintaining the semantic locality across consecutive tokens is very simple and efficient in terms of latency.\n\n2.\tThis paper is well-organized and easy to read. The paper first shows the latency bottleneck on expert swapping for memory-constrained scenario, and then formulates latency optimization problem with a simple yet effective solution. The overall organization is easy to follow and smooth.\n\n3.\tThe authors demonstrate the effectiveness of the proposed method across different tasks and provide some visualization results (such as attention, semantic group embedding) to support some claims.\n\nWeakness:\n\n(1)\tThe experimental result on the temporal inconsistency of inter-token expert activation is missing. It would be more convincing to provide a metric to quantitatively measure the temporal inconsistency and show how Oracle-MOE reduces this inconsistency.\n\n(2)\tThe semantic locality, as a key assumption, is not well explained. It would be better to provide some preliminary experiments to show the evidence in real datasets. It would be interesting to see the semantic locality across different models/layers/samples.\n\n(3)\tThere are some approximations in the derivation of oracle-MOE, such as CSD_token in Line 147, page 3, CSD_oracle in Line 231, page 5. It would be better to validate such approximation in the experiments design.\n\n(4)\tAblation study on CSD and hyperparameter study of $\\gamma$ is missing. How does $\\gamma$ influence the proposed algorithm?", "other_comments_or_suggestions": "N/A", "questions_for_authors": "N/A", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "3As7bgHkDW", "reviewer": "Reviewer_jf9n", "summary": "This paper presents Oracle-MoE, a novel Mixture-of-Experts (MoE) architecture aimed at efficiently deploying Large Language Models (LLMs) on memory-constrained edge devices. Current MoE models, despite theoretical advantages for memory efficiency, suffer from high latency during inference due to frequent swapping of experts in and out of limited device memory. The authors identify that these latencies stem primarily from temporal inconsistencies in expert activations between consecutive tokens. To address this, Oracle-MoE introduces a new routing mechanism that groups tokens based on semantic similarity, routed within a compact \"oracle space\" defined by semantic group embeddings. This method significantly reduces the variability of expert activations, thereby minimizing expert swapping overhead and improving inference speed on devices with limited memory.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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\n\n_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.\n\n_A1_ Results on A100s are listed below. Our method still speeds up inference by 50%~350%.\n|Model size|Memory budget|Switch(FIFO)|Switch(LRU)|Switch(SwapMoE)|Ours(FIFO)\n|---|---|---|---|---|---|\n|9*24(2.06B)|1GB|16.613s|16.066s|15.900s|4.012s|\n||2GB|13.762s|15.229s|12.751s|3.630s||\n||4GB|11.165s|11.973s|9.754s|3.529s||\n||7GB|6.688s|7.115s|6.135s|3.410s||\n||Full Memory|3.252s|3.252s|3.252s|3.252s||\n\n_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.\n\n_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.\n|Avg Prediction Acc from|Qwen|Switch(Ourtrained)|Oracle(Ourtrained)|\n|---|---|---|---|\n|Layer 0|48.23|50.98|89.66|\n|Half of the model layers|55.59|59.72|92.82|\n|3/4 of the model layers|63.27|68.68|95.51|\n\n_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.\n\n_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.\n\n_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).\n\n_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.\n|Inference Latency|2.5GB|4.5GB|7GB|10GB|12GB(full)|\n|---|---|---|---|---|---|\n|Our Oracle Routing(s)|24.937|19.701|17.134|16.726|15.601|\n|Normal MoE(s)|113.828|89.082|44.013|27.058|15.601|\n\n_Q5_ typo \"Lanuge\".\n\n_A5_ We will correct the typos in the updated version.", "claims_and_evidence": "The Claims And Evidence are clear and well supported.", "methods_and_evaluation_criteria": "The paper proposes routing based on semantic embeddings derived from attention score clustering. In general it's innovative and well justified. However Oracle-MoE initialization requires a warm-up stage and clustering analysis (e.g., via K-means and SVD), which introduces complexity. Evaluation is done primarily on the NVIDIA Jetson Xavier NX, which could limit the scope of edge devices. Broader evaluations could enhance confidence in the method’s universality.", "theoretical_claims": "The Theoretical Claims are well justified.", "experimental_designs_or_analyses": "The Experimental Designs Or Analyses are sound.", "relation_to_broader_scientific_literature": "This is going to have most impact to the budget-constrained inference domain.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Strengths\n- The paper clearly identifies a practical limitation—high latency from expert swapping—providing insightful analysis into its underlying cause (temporal inconsistency).\n- The proposal of routing tokens based on semantic locality rather than token-specific embeddings is innovative and well-motivated \n- The authors offer a robust theoretical framework around Consecutive Semantic Difference (CSD)\n- A detailed breakdown of latency composition provides clear empirical insights into exactly how Oracle-MoE improves performance over baseline strategies.\n- Oracle-MoE is novel in its introduction of an oracle-space routing approach, leveraging semantic locality in attention scores\n\nWeaknesses\n- 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.\n- 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.\n- 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.\n- The paper did not discuss cases of using fine-grained experts ( e.g., number of experts > 128, like the ones in DeepSeek's model).", "other_comments_or_suggestions": "\"Lanuge\" -> Language in title", "questions_for_authors": "N/A", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xpCPF2CVEt", "reviewer": "Reviewer_fF5P", "summary": "This paper proposes Oracle-MoE, which improves the MoE inference efficiency by exploiting semantic locality to reduce swapping demands.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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\n\n_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.\n\n_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.\n|Inference Latency|2.5GB|4.5GB|7GB|10GB|12GB(full)|\n|---|---|---|---|---|---|\n|Our Oracle Routing(s)|24.937|19.701|17.134|16.726|15.601|\n|Normal MoE(s)|113.828|89.082|44.013|27.058|15.601|\n\n_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?\n\n_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.\n\n_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).\n\n_A3_ Thank you for your suggestions. We will make such modifications to the layout of these pictures in the updated versions.", "claims_and_evidence": "Please see **Other Strengths And Weaknesses**", "methods_and_evaluation_criteria": "Please see **Other Strengths And Weaknesses**", "theoretical_claims": "Please see **Other Strengths And Weaknesses**", "experimental_designs_or_analyses": "Please see **Other Strengths And Weaknesses**", "relation_to_broader_scientific_literature": "Please see **Other Strengths And Weaknesses**", "essential_references_not_discussed": "Please see **Other Strengths And Weaknesses**", "other_strengths_and_weaknesses": "**Strengths**:\n\n1. The paper is easy to follow, with clear writing and presentation.\n2. The authors provide detailed theoretical analysis to justify their method.\n3. The evaluation results are good and comprehensive.\n\n**Weaknesses**:\n\n1. The main concern I have with this paper is the issues of scalability. In the paper, the largest MoE models are 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.\n\n2. 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?\n\n3. 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).\n\n[1] Mixtral of Experts, Arxiv 2024.\n\n[2] DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models, Arxiv 2024.", "other_comments_or_suggestions": "Please see **Other Strengths And Weaknesses**", "questions_for_authors": "Please see **Other Strengths And Weaknesses**", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Gt868L3Osn", "reviewer": "Reviewer_bLcx", "summary": "This paper proposes a new MoE architecture, Oracle-MoE, to address the latency issues associated with deploying large language models (LLMs) on edge devices with limited memory.  The key idea is to route tokens in a compact space, called the oracle space, which is derived from attention scores to maintain semantic locality across tokens.  This approach reduces the frequent swapping of experts in and out of memory, which is a major cause of latency.  The paper provides theoretical analysis and experimental results on various GPT-2 models and downstream tasks, demonstrating that Oracle-MoE achieves state-of-the-art inference speeds without compromising task performance.\nThis work primarily targets the low throughput setting where \"only a few activated experts <are required> to reside in memory for inference\".", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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\n\n_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.\n\n_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.\nTheoretical 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.\n\n[A] A Survey on Word Embeddings: From Shallow Models to Deep Learning (Goldberg, 2017)\n\n[B] Deep Learning for NLP and Speech Recognition\" (Hinton et al., 2012)\n\n[C] Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer (Raffel et al., 2020)\n\n[D] Effective Approaches to Attention-based Neural Machine Translation(Luong et al., 2015)\n\n[E]Analyzing the Structure of Attention in Transformers (Kobayashi et al., 2020)\n\n[F]On the Expressive Power of Self-Attention Matrices （Valerii Likhosherstov et al.,2021）\n\n_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.\n\n_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. \nTo 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.\n\n_Q3_ Narrow scope of application If you're in a memory-constrained LLM inference setup, maybe you shouldn't be using MoEs.\n\n_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.", "claims_and_evidence": "Claims with Clear and Convincing Evidence:\n\n- High latency due to expert swapping: The paper provides clear evidence of the high latency caused by expert swapping in memory-constrained MoE-based LLM inference. They show that 50-85% of the latency is due to I/O overheads from expert swapping.\n- Temporal inconsistencies in expert activations: The visualization of expert activations over a sequence of tokens clearly demonstrates the temporal inconsistencies, leading to frequent expert swapping.\n- Oracle-MoE reduces expert activation variations: The paper provides evidence that Oracle-MoE effectively reduces expert activation variations compared to the Switch Transformer.  \n- Performance on downstream tasks: The results on various downstream tasks show that Oracle-MoE achieves similar or better performance compared to the Switch Transformer.  \n\nClaims that Need Further Support:\n\n- 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.  \n- 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.\n\n\nOverall, the paper presents a promising approach to address the latency challenges in memory-constrained MoE-based LLM inference. However, providing more detailed evidence and explanation for some of their claims would further strengthen their contribution. I'll also note that, if you're in the extremely memory-constrained MoE-based LLM inference setup, maybe an MoE is NOT the right solution.", "methods_and_evaluation_criteria": "Yes, the proposed methods and evaluation criteria in the paper generally make sense for the problem of memory-constrained LLM inference on edge devices.\n\nThey define an Oracle Space motivated by Semantic Grouping of tokens. They cluster tokens in to groups in the oracle space; and they use Expert Prediction to see how much they can accelerate async loading of experts.\n\nThe chosen evals (Trivia QA, GLUE, MAG, Sci-Cite, XSum) are standard and relevant for evaluating the performance of LLMs on various downstream tasks; other evals could have been chosen, but these are ok. Some of the chosen evals are not the most standard in literature; hard to say what is / is not cherry picked.\nMemory-Latency Curve and First Token Latency are used to measure the trade-off between memory usage and inference latency; these are reasonable choices.\n\nThe proposed methods and evaluation are good", "theoretical_claims": "I read theorem 1 and the skimmed the appendix proof. I have no strong statements about it.", "experimental_designs_or_analyses": "The experimental designs and analyses in the paper are generally sound and valid. They provide convincing evidence for the effectiveness of Oracle-MoE in reducing latency without compromising performance. However, as mentioned earlier, providing more details on certain aspects, such as semantic locality and expert prediction, would further strengthen the paper's contributions.", "relation_to_broader_scientific_literature": "There are a lot of works (eg Mixtral) which look at expert assignment based on domains. These are not exactly the same as semantic meaning of tokens, but it should be mentioned and the difference explicitly noted from what you're doing.", "essential_references_not_discussed": "not that i noticed", "other_strengths_and_weaknesses": "Narrow scope of application\nIf you're in a memory-constrained LLM inference setup, maybe you shouldn't be using MoEs.", "other_comments_or_suggestions": "none", "questions_for_authors": "none", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper presents Oracle-MoE, a memory-efficient Mixture-of-Experts architecture that exploits semantic locality to reduce inference latency. All four reviewers expressed positive opinions, i.e., three with Weak Accept and one with Accept and the rebuttal addressed key concerns, including semantic locality, temporal inconsistency, and scalability, thus prompting one reviewer to raise his/her score. Given the overall consensus and the paper’s well-structured contributions, I would recommend acceptance.
Decision: Accept (poster)