On the Role of Discrete Representation in Sparse Mixture of Experts
Authors: Giang Do
Kha Pham
Hung Le
Truyen Tran
Venue: Accepted by TMLR ()

Abstract: Sparse Mixture of Experts (SMoE) is an effective solution for scaling up model capacity without increasing the computational costs. A crucial component of SMoE is the router, responsible for directing the input to relevant experts; however, it also presents a major weakness, leading to routing inconsistencies and representation collapse issues. Instead of fixing the router like previous works, we propose an alternative that assigns experts to input via \emph{indirection}, which employs the discrete representation of input that points to the expert. The discrete representations are learned via vector quantization, resulting in a new architecture dubbed Vector-Quantized Mixture of Experts (VQMoE). We provide theoretical support and empirical evidence demonstrating the VQMoE's ability to overcome the challenges present in traditional routers. Through extensive evaluations on both large language models and vision tasks for pre-training and fine-tuning, we show that VQMoE achieves a 28\% improvement in robustness compared to other SMoE routing methods while maintaining strong performance in fine-tuning tasks.

Published in Transactions on Machine Learning Research (08/2025)
On the Role of Discrete Representation in Sparse Mixture of Experts
Giang Do
truong. do &deakin edu.ou
Applied Artificial Intelligence Initiative (A212) Deakin University
Kha Pham
phli@deakin edu.au
Applied Artificial Intelligence Initiative (A212) Deakin University
Hung Le
thai, Le &deakin edu.
Applied Artificial Intelligence Initiative (A212)
Deakin University
Truyen Tran
trugen. tran Edeakin, edu..ox on
Applied Artificial Intelligence Initiative (A212)
Deakin University
Reviewed on OpenReview:
Abstract
Sparse Mixture of Experts (SMoE) is effective solution for scaling up model capacity without increasing the computational costs. A crucial component of SMoE is the router, responsible for directing the imput to relevant experts; however, it also presents a major weakness, leading to routing inconsistencies and representation collapse issues. Instead of fixing the router like previous works, we propose an alternative that assigns experts to input via indirection, which employs the discrete representation of input that points to the expert. The discrete representations are learned via. vector quantization, resulting in new architecture dubbed Vector- Rector-Quantized Mixture of Experts (VQMoE). We provide theoretical support and empirical evidence demonstrating the VQMoE's ability to overcome the challenges present in traditional routers. Through extensive evaluations on both large language models and vision tasks for pre-training and fine-tuning, we show that VQMoE achieves 28% improvement in robustness compared to other SMoE routing methods while maintaining strong performance in fine-tuning tasks. The code is publicly available at http://ttttbb.oom///////gg/////p
1
Introduction
Scaling Transformer models with increasing data and computational resources has led to remarkable advances across a wide range of domains, including natural language processing (NLP) (Du al. 2022: Fedus al. 2022 Zhou et al. 2024 and visual reprosentation learning (Riquelme et al. 2021a Shen et al.) 2023b
Despite these successes, training and deploying large-scale dense Translormer models often require substantial computational resources, frequently amounting to hundreds of thousands of GPU hours and incurring costs in the millions of dollars (Kaddour et al. 2023). To address this scalability bottleneck, Sparse Mixture of
Experts (SMoE) architectures have emerged as a promising alternative (Sazeer et al. 2017 Zoph et al.
*Corresponding author

---

Published in Transactions on Machine Learning Research (08/2025)
Xue et al. 2024 Jiang et al. 2024). Inspired by classical Mixture of Experts formulations (Jacobs et al. 1991a SMoE models consist of multiple expert subnetworks with shared architectures, where a routing mechanism dynamically selects a small subset of experts (often one or two) for each imput token. This sparsity significantly reduces inference costs compared to dense counterparts of similar model capacity
2022
Artelxx et al. 2022
Krajewski et al. 2024), making SMoEs attractive for efficient scaling.
Despite their efficiency benefits, SMoE's face critical training challenges, most notably, representation col- lapse. This phenomenon occurs when only a small subset of experts are frequently activated, when all experts converge to similar representations, thereby negating the diversity and specialization that the ar- chitecture is intended to promote. Prior works have sought to mitigate this issuc by improving the routing policy through regularization and auxiliary losses (Chi et al.) 2022; Chen et al. 2023a Do et al. 2023). However, these approaches focus on the routers improvement rather than questioning its necessity.
In this work, we explore a more fundamental question: Is an explicit router necessary at all? We argue that incorporating discrete represenlations offers a principled alternative. Discrete latent variables are inherently suited to capturing structured and interpretable patterns within data, aligning with the symbolic nature of human cognition, where concepts are often discretized as words, tokens, categories. In the SMoE context, discrete representations can improve input routing by naturally clustering similar inputs, thereby enhancing expert specia lizatiniiation and utilization without relying solely on a learned gating mechanism.
Employing vector quantization (VQ) techniques to learn discrete representation, this paper proposes a novel mixture of expert framework, named VQMoE, which overcomes the representation collapse and inconsistency in training sparse mixture of experts. More specifically, we prove that the existing router methods are inconsistent and VQMoE suggests optimal expert selection for training SMoE. Additionally, our method guarantees superior SMoE training strategies compared to the existing methods by solving the representation collapse by design.
We evaluate the proposed method by conducting pro-training of Large Language Models (LLMs) on several advanced SMoE architectures, such as SMoE (Jiang et al. 2024), StableMoE (Dai el al. 2022). or XMoE CCh et al. 2022). followed by fine-tuning on downstream tasks on both Language and Vision domains.
In summary, the primary contributions of this paper are follows:
We theorctically demonstrate that learning discrete representations provides effective mechanism for expert selection, and that VQMoE intrinsically mitigates the problem of representation collapse.
We propose the use of vector quantization (VQ) to learn structured and interpretable expert clusters.
We conduct extensive experiments on large langnage models as well as vision pre-training and fine- tuning tasks to validate the effectiveness of our method.
We provide comprehensive analysis of VQMoE's behavior, offering insights into its performance and robustnoss.
2 Related Work
Sparse Mixture of Experts (SMoE). Sparse Mixture of Experts (SMoE) builds on the Mixture of Experts (MoF) framework introduced by Jacobs et al. (191bb) Jordan & Jacobs (1994), with the core idea that only subsct of paramcters is utilized to process cach example. This approach was first popularized by Shazeer et al. (2017). SMoE's popularity surged when it was combined with large language models based
Transformers (Zhou et al. 2022 Li et al.) 2022 Shen et al. 2023a), and its success in natural language procossing led to its application across various fields, such as computer vision (Riquelme et al. 2021b Hwang
et al. 2023 Lin et al. 2024), speech recognition (Wang et al. 2023 Kwon & Chung 2023 and multi-task
learning (Ye & Xu 2023 Chen et al. 2023b).
However, SMoE faces a major problem training known as representation collapse, i.c., the experts converge to similar outputs. To address this, various methods have been introduced. XMoE (Chi et al. 2022) cal-
culates routing scores between tokens and experts a low-dimensional hypersphere. SMoF-dropout (Chen
2

---

Published in Transactions on Machine Learning Research (08/2025)
et al. 2023a) uses fixed, randomly initialized router network to activate experts and gradually increase the number of experts involved to miligate collapse. Similarly, HyperRouter (Do et al. 2023 utilizes HyperNet- works (Ha et al. 2016) to generate router weights, providing another pathway for training SMoE' effectively, StableMoE (Dai et al. 2022) introduces a balanced routing approach where a lightweight router, decoupled from the backbonn model, distilled to manage loken-lo-expert, assignments. The StableMoE strategy ensures stable routing by freezing the assignments during training, while SimSMoE) Do et al. (2024) forces experts to learn dissimilar representations. Despite these extensive efforts, the representation collapse issue persists, highlighted by Pham et al. (002)). While most solutions focus on improving routing algorithms, our approach takes different path by learning discrete representation of input that points to relevant experts.
Discrete Representation. Discrete representations align well with human thought processes; for example, language can understood as a series of distinct symbols. Nevertheless, the use of discrete variables deep learning has proven challenging, as evidenced by the widespread preference for continuous latent variables most current research. VQVAE (van den Oord et al. 2017) implements discrete representation in Variational AutoEncoder (VAE) (Kingma & Welling) 2022 using vector quantization (VQ). IMSAT (Hu et al. 2017) attains a discrete representation by maximizing the information-thoorcte dependency between data and their predicted discrete representations. Recent works follow up the vector quantization ideas and make some enhancements for VAF, for example: (Yu et al. 2022); (Mentzer et al. 2023); J; and (Yang et al.] 2023 Mao et al. (2022) utilize a discrete representation to strengthen Vision Transformer (ViT) (Dosovitskiy et al. 2021p. To the best of our knowledge, our paper the first learn discrete representation of Sparse Mixture of Experts.
3 Method
We propose a novel model, Vector-Quantized Mixture of Experts (VQMoE), which learns discrete represen- tations for expert selection. As illustrated in Fig. Ia] our approach selects experts directly based on the input representation, eliminating the need for a trained router. To prevent information loss, we integrate discrete and continuous representations within the model.
3.1 Preliminaries
Sparse Mixture of Experts. Sparse Mixture of Experts (SMoE) is a variant of the transformer archi- tecture in which the conventional leed-forward layers (MLPs) are replaced with Mixture of Experts (MoE) layers (Shazeer et al. 2017). Given an input x E Raxd which represents the output of the multi-head attention (MHA) module, the SMoE layer computes a sparse weighted combination over a set of N expert networks. Each expert typically feed-forward neural network (FN, and its contribution to the final output is determined by routing function S(x). The resulting output of the SMoE layer is given by:
N
(x) = Σ S(a), FFN((a)
(1)
Σ S(x), W & (WFrx,*),
where o(.) denotes a non-linear activation function (e.g., ReLU or GELU), and WFFN, W are the learnable weights of the i-th expert. The routing weights S(x) are computed using a Top-k selection over the softmax scores derived from the dot product of the input with a learned expert embedding matrix W.. defined below:
S(x) = (onK(ssoaaxxWWa))
fo, if vs E top & largest elements of v,
(2)
TopK(v,
-00 otherwise.
3

---

Published in Transactions on Machine Learning Research (08/2025)
This sparse selection mechanism ensures that only small subset of experts activated for each imput, which significantly reduces computational cost while retaining model capacity.
Discrete Representation Learning. van den Oord et al. (017)) propose VQVAE, which uses Vector Quantization (VQ) to learn a discrete representation. Given an input I E Raxd, VQVAE discretized the Rxxd where K is the codebook size and d is the dimension of the embedding. Tet denote (x) Rnnd the output of the VQVAE and 1( the indicator function. The discrete representation 2x(2)) Uk: where k = argmin, (x;) Uyll2 is achieved by vector quantizer qe that maps integer
input into a codebook
for each imput as:
qo (2 = x) = 1 k = arg min 2|x) - V1ll2)
j=1:
3.2 Vector-Quantized Mixture of Experts (VQMoE)
Pre-training VQMoE. Traditional Sparse Mixture of Experts (SMoE) models utilize continnous token representations and route them to experts based on learned token-expert affinity scores. We propose novel architecture, VQMoE, that Iearns both continuous and discrete representations jointly during pre- training (see Figure Tab The contimuous component captures fine-grained data patterns, while the discrete component, learned via vector quantization, encodes robust latent structure useful for downstream transfer.
Let x E R denote the input to the VQMoE layer (o.g., output from a multi-hoad attention block), and let f4 denote the vector quantization operator. The VQMoE output during pre-training defined as:
f ( = ge(x)) SMoE (x) + ga(x))
Σ
PFN (āt)
(4)
(Continuous representation)
(Discrete representation)
In this formulation, (x)) denotes the output from a standard Sparse Mixture of Experts (SMoE) layer, capturing the contimuous expert representations. The second term corresponds to the discrete representation, where each JPFN is the I-th feedforward expert network. The input to each discrete expert, denoted as Et. = vù if the input vector my is assigned to the 1-th
is determined by vector quantization: specifically, codebook vector U;; otherwise, it is set to the zero vector, i.e., क = 0. Here, K is the number of vector quantization codebooks, and UR is learned codebook vector assigned by f', The gating functions g=(x)) and gz( as Equation 5 modulate the contributions of the continuous and discrete pathways, respectively, and are typically computed based on the input æ through learnable mechanisms.
[gc(z) ga(z)] softmax(W,xx) W E g2xd
(5)
To address the mismatch between the number of codebook vectors and the number of expert networks, we introduce a flemible code strategy. This approach enables consistent routing from quantized representations to experts, even when the two quantities differ. Specifically, we define hash-based mapping using a modulo operation. Let deb denote the index of a codebook vector, and let texp denote the index of the corresponding expert. The mapping is given by:
eex ich mod N,
(6)
where N is the total number of experts. This ensures each codebook index is deterministically assigned to one of the available experts
Fine-tuning VQMoE. Based on insights from Geva et al. (2011), which note that forw-forward layers (FFNs) constitute a significant portion of a transformer's parameters, we adopt a light finc-tuning strategy that retains only the discrete path of the VQMoE. This allows efficient adaptation while leveraging pre-trained latent. representations (see Figure Tby. The fine-tuning output becomes:
4

---

Published in Transactions on Machine Learning Research (08/2025)
yVQMoE (x)
Σ f SFPN (it)
3.3 Training Procedure
Pretraining. The training objective is jointly minimizing the loss of the target task and losses of the Vector Quantization module (C and ['ommitmentt ) as in (van den Oord et al. 2017). Equation S specifies the overall loss function for training VQMoE with three components: (1) task loss; (2) l2 loss; (3) commitment loss. While _ helps to move the embedding e towards the outputs (x)). the commitment loss makes sure the output of the Vector Quantization module commits to the embedding and its output does not grow. The Vector Quantization algorithm does not vary with 8, we follow introduce a new parameter, a, regulate the contribution of the Vector Quantization loss the overall loss. A higher value of favors stronger adherence to the discrete representation, and vice versa.
0.25 as van den Oord al. (2017). We
L= L Caask + a [z)] - 0112 + B (zzz(=) sg|u]l1>)
(8)
where sg(.) is the stop gradient operator defined as follows:
sg(x) = {
x forward pass
(9)
backward pass
In Equation Clask denotes the task-specific loss, which depends on applications. For example, in language modeling tasks, Caask is typically defined the negative log-likelihood (NLL) of the target tokens (Dai et al. 2019c), promoting accurate next-token prediction. In image classification tasks, gtask is usually implemented as the cross-entropy loss between the predicted class distribution and the ground-truth label (He et al. 2015 encouraging correct class assignment.
Fine-tuning. For downstream tasks, fine-tune the pretraining model by utilizing the codebook learned from the Equation g by freezing all parameters at the Vector Quantization module. Thus, the training
objective simply becomes: L C
4
Theory Analysis
4.1
Optimal Experts Selection
Problem settings. We consider an MoE layer with each expert being MLP layer which is trained by gradient descent and input data (x)), generated from data distribution D. Same (Chen el. al. 2022]; Dikkala et al. 2023), we assume that the MoE input exhibits cluster properties, meaning the data is generated from N distinct clusters (C1, C2,..,,
Definition 4.1 (Consistent Router) A sequence of points 21, 22, .. and corresponding sequence of clusters C1, C2..... . ON are said to be consistent if, for every point xp E Ct, the condition
........
dist(Ip u.) < min dist(ap, uy)
jf1
is satisfied, where dist(a,)) denotes the distance between and b, and u, is the center of cluster C,.
Definition 4.2 (Inconsistent Router) A sequence of points T1 .., In and a corresponding sequence of clusters C₁, Cz,..., On are said to be inconsistent if there exists point xp E Cs such that
dist( u.) min dist(App uy).
jft
where dist(a, b) represents the distance between and b, and un is the center of cluster Ct-

---

Published in Transactions on Machine Learning Research (08/2025)
Output
Output
FFN
FFN
FFN
FFN
FFN 1
FFN
FFN 3
FFN 4
Embeding ad
Embeding Space
Vector
Quantizer
Vector Quantizer
Gating
Router
Input
Input
(a) VQMoE Pre-training
(b) VQMoE Fine-tuning
Figure 1: Illustration of the proposed VQMoE architecture for Pre-training and fine-tuning. (a) At the Pre-training stage, VQMoE architecture learns simultaneously continuous and discrete representation at the Pre-training phase. The contimouu reprosentation is learned by the conventional SMoE, while the Vector Quantization block facilitates the learning of a discrete representation. The final output is then combined by gate layer. (b) VQMoF. learns a discrete representation that is capable of operating efficiently and robustly on downstream tasks. VQMoE computes the discrete representation only during the fine-tuning stage achieve robustness and efficiency.
Inspired by (Dikkala et al. 2023), conceptualize the router in Sparse Mixture of Experts clustering
problem. This leads us to define consistent router in Definition 4.1 Furthermore, we introduce definition for an inconsistent router in SMoE as outlined in Definition 4.2 along with the concept of inconsistent expert
selection presented in Theorem 4.3 during the training of SMoF.
Theorem 4.3 (Inconsistent Experts Selection) Lel SMHA a mulli-heed atlention (MHA) function producing an output x E Raxd, and consider N experts with embeddings es for expert i where i E [1, N). Assume that fMHA converges step lm, whill the eapert embeddings converge at step ter with im te. For each output x, an expert P [1, is selected such that
P= arg min dist(I,ez)) JE(L.N]
Under these conditions, the expert embeddings form inconsistent routing mechanism.
4.3 The proof of Theorem 4.3 is given in Appendix A.1.2, and we have the following insights. Theorem implies that an expert selection process by router as the conventional SMoE leads to the inconsistent router. Indeed, the router layer is designed as a simple linear layer, I is the output of MHA function in practice; and an SMoE router siguificantly simpler than the MHA function. Consequently, this design leads the router functioning as an inconsistent router, contributing the representation collapse issue and instability during training.
Proposition 4.4 (Optimal Experts Selection) Given input data partitioned into N clusters (C1, C2...., ,.C) and mixture of experts (MoE) layer with N experts (F1, F2,..., EN), the assign- ment of each cluster O₁ to expert E, for i [1, constitutes an optimal expert selection solution.
Proposition 4.4 demonstrates that we are given a clustering strueture as input, assigning each part of the input to its corresponding expert results in an optimal expert seloction. This implics that learning a discrete representation and directing each component to the appropriate expert yields an optimal solution. The proof of Proposition 4.4 can be found in Appendix A.1.3

---

Published in Transactions on Machine Learning Research (08/2025)
4.2 Experts Representation Collapse
The representation collapse problem in Sparse Mixture of Experts (SMoE), where all experts converge to similar representations, was first highlighted by Chi et al. (2022). Following Chi et al. (2022) and Do et al. (2023) we analyze this issue using the Jacobian matrix of the model output with respect to the input x. E Raxd_ The Jacobian for SMoE is expressed as:
JSMoE -(()) + Σ
(10)
= =S(I) FFN + Σ c.e
where c; = (()) (a.,((o)((())) -5(x)) JFFN is the Jacobian of the selected expert's feedforward network, and Ej are the expert embedding vectors. Equation 10 consists of two components: S(I) the main signal path from the input the output through the selected expert: and Σ c,e the contribution from the gating function's gradient with respect to the expert embeddings.
Since the summation projection restricts the output space from Rd to RN, which elfectively causes representation collapse.
expert embeddings lies in subspace of dimension N, and typically N < d, this
Jacobian Analysis of VQMoE. To examine whether VQMoE mitigates this collapse, we derive the Jacobian the VQMoE output with respect to the imput x E Rxxd The detailed expression of the VQMoE Jacobian matrix is provided Section A.I.I. Specifically, we have:
VQMoE = -o(x)) JSMOE Oge(a)
M (æ)
+ Ox
@ya(w)
K
+ ga(a))
Σ J.FN
dx
(11)
1=1
1=1
N+K+2
= J +
, j=1
Same as the Jacobian matrix of SMoE, the Jacobian matrix of VQMoE consists two terms: (1) J1 depends on input token and exports to the final output; (2) b N2 o indicates to learn better gating function minimize the lask loss. We can see that + k + >> N, implying that VQMoE is better than SMoE in solving the representation collapse issue. In theory, we can choose the number of codebook to be approximately d - N - 2 with hashing index to experts to address the issue. However, this involves trade-off with the computational resources required to learn the codebook.
5 Experiment
We conduct experiments to investigate the following hypotheses: (i) VQMoE offers an effective training algorithm for Sparse Mixture-of Experts (SMoE) in large language models (LLMs); (ii) VQMoE enables efficient fine-tuning; and (iii) VQMoE outperforms other routing methods across multiple domains.
5.1 Experimental Settings
To evaluate the three hypotheses, we conduct experiments across both vision and language tasks. For pre- training language models, we assess two standard benchmarks: (i) character-level language modeling using enwik8 and text8(Mahoney 2011), and (ii) word-level language modcling using 1ikkText-103(MMrrigt 2016) and the more challenging One Billion Word (Imlb) dataset (Chelba al. 2014). All experiments use the standard training, validation, and test split with 90:5:5 ratio as a Child al. 2019)
al.,

---

Published in Transactions on Machine Learning Research (08/2025)
Cuufiguration
Eawiks (BPC)
Text8 (BPC)
WikiTeei 103 (PPL)
Imib (PPL)
Avg. Char level
Avg. Word level
Architecture
Algorichm
Hase
Large
Hase
Large
Base
Large
Base
Large
VQMoE
1.48
1.41
1.47
1.40
38.74
31.98
59.48
49.30
1.44
44.88
SMoF
1.49
1.41
1.49
1.40
39.50
32.30
60.88
51.30
1.45
45.50
Transformer
OMoE-Dropout
1.82
2.22
1.70
1.89
72.62
107 .18
97.45
159.09
1.91
109.59
XMoE
1.51
1.42
1.49
1.42
39.56
32.65
61.17
51.84
1.46
46.06
StableMoE
1.49
1.42
1.49
1.41
39.15
32.31
60.72
50.74
1.45
45.81
VQM&E
1.19
1.08
1.28
1.17
29.48
23.85
56.85
48.70
1.18
39.72
SMoE
1.29
30.16
48.71
40.22
Transformer-X
SMoE-Dropout
1.56
2.21
1.56
1.86
58.37
93.17
65
1.81
65.35
XMcE
1.21
1.09
1.28
1.17
30.34
24.22
58.33
50.64
1.19
40.88
StableMoE
1.20
1.10
1.28
1.19
29.97
24.19
58.25
49.17
1.19
40.40
# Params
20M
210M
20M
210M
20M
210M
20M
210M
Table 1: Bits- Bits-per-character (BPC) on the Enwik8 and Text8 test sets; and perplexity (PPL) on the WikiText- 103 and One Billion Word test sets. Avg. Char-level is the average BPC over Enwik8 and Text8; Avg. Word-level is the average PPL over WikiText-103 and lmlb. Lower is better; best results are in bold.
For parameter-effiiintt fine-tuning, we fine-tune models pre-trained on enwik8 using four widely used NLP datasets: SST-2, SST-5(Socher et al. 2013). IMDB(Maas et al. 2011), and LANKINGG7(Casameeee et al. 2020). TollowingChen et al. (2023a) freeze the router and update only the expert parameters to evaluate fine-tuning efficiency.
For vision tasks, we employ the Vision Transformer (ViT) Dosoviiskyy et al. 2021] and compare our rout- ing method with state-of-the-art alternatives on five benchmark image classification datasets: CIFAR-10, CIFAR-100/(Kizheee 2009), STI-110Coates et al. 2011), SVHN Netzer al. 2011), and ImageNet- 1K (Deng et al. 2009).
5.2 Pre-training Language Models
Models. For the language tasks, we follow the same settings as in SMoE-Dropout (Chen al., 2023ah. We consider two decoder-only architectures: (i) the standard Transformer (Vaswani et al.) 2017); and (ii) and Transformer-XL, (Dai et al.) 2019a) with the same number of parameters Transformer. We evaluate method versus the state of art Sparse Mixture of Expert Layers such as StableMoE (Dai et al. 2022 and XMoE (Chi et al. 2022) with topk = 2 in the experiments. We consider two model configurations: (i) base: with four SMoE blocks and 20M parameters; (ii) large: with twelve SMoE layers and 210M parameters. We emphasize that we are not trying to achiove state-of-tho-art results due to the limited resource constraints. Instead, we evaluate the small and large models on various datasets demonstrate the scalability and efficacy of our algorithm. Lastly, we conduct extensive investigations using the tiny model understand the algorithm behaviours and their robustness different design choiccs.
Baselines. We compare our VQMoE with state-of-the-arl SMoF training strategies for LLMss SMoE (Jian et al. 2024) employs simple router trained end-to-end with the experts. StableMoE (Dai et al. 2022 proposes a two-phase training process where the first phase trains only the router, and then the router is fixed to train the experts in the second phase. XMoE (Chi et al. 2022) implements a deep router that comprises a down-projection and normalization layer and gating network with learnable temperatures. Lastly, motivated by SMoE-Dropout (Chen al. 2023a),, we implement the SMoE-Dropout strategy that employs randomly initialized router and freeze it throughout the training process.
Training procedure. For the language modeling experiments, we optimize the base models and the large models for 100,000 steps. We use an Adam (Kingma & Ba 2017) optimizer with Cosine Annealing learning rate schedule [Loshchilov & Hutter 2017). The lowest val validation loss checkpoint is used to report the final performance on the test set.
Q1: Does VQMoE perform better on Pre-training tasks compared to routing methods? A1: Yes.
8

---

Published in Transactions on Machine Learning Research (08/2025)
162
VOMeE
VENse
200
142
180
Srahiellf
120
Perplouty IPPL
160
102
Perpexity
1.0
80
102
20000 40000
60000
80000
100003
60
20000 40000
60000
80000
100000
Step
Step
(a) Training PPL movement on Wikitext-103 dataset.
(b) Training PPL movement on Imlb dataset.
Figure Perplexity (PPL) over training steps for the Transformer-XI. base model on two datasets: (a) WikiText-108 and (b) lmlb. The results indicate that VQMoE converges faster than the baseline models, demonstrating its efficiency and robustness for language modeling tasks.
Table presents the evaluation metrics comparing VQMoE with state-of-thh-art approaches. We also show the performance progression of the base model on the validation set. Notably, across all methods and datascts, VQMoE consistently outperforms the baseline models for both the Transformer-XL and Transformer architectures on average. Although advanced strategies such as XMoE and StableMoE generally outperform the vanilla SMoE on character-basee datasets such as enwik8 and text8, which involve a small vocabulary size, their improvements tend to diminish or become marginal when trained on more complex, vage-locabularry datasets such as WilaTerr- 109 and One Bilionn Word (lmfb). In contrast, VQMoE consistently outperforms all competitors across benchmarks (keeping in mind that the BPC metric log-scaled), architectures, and also converges more quickly as Figure 2 This highlights VQMoE's effectiveness in learning an efficient routing policy for the language modeling pre-training task.
Q2: Does VQMoE keep outperforming the router method when scaling up? A2: Yes.
Table D also demonstrates that VQMoE maintains consistently strong performance when scaled up to 12- layer Transformer and Transformer- XL architectures, Across all four datasets, the performance gap between VQMoE and other routing methods widens as the dataset size incrcases, from enwik8 to the One Billion Word datasel. This suggests that our approach has the potential to scale effectively with larger language models and bigger datasets. An interesting observation is that EMMo--Dropoe (Chen et al.) 2023a) performs the worst among all methods, indicating that a random routing policy is insufficient and requires updating for effective training. This finding highlights that the suecess of EMoE-Dropout is largely due to its self- slimmable strategy, which linearly increases the number of activated experts (K) during training. However, this approach transforms the sparse network into a dense one, contradicting the original motivation behind using SMoE for largo-scale models.
Q3: Can VQMoE, with only 80% of the total parameter count, achieve better performance than SMoE utilizing the full 100% of parameters? A3: Yes.
To evaluate the robustness of VQMoE, we reduce its hidden dimension to half that of the SMoE baseline, resulting in approximately a. 20% reduction in the total number of parameters. Robustness here denotes the model's ability to maintain strong performance across different parameter scales, particularly with fewer parameters. We then train both models across a range of parameter scales: 1M, 2M, 4M, 8M, and 16M, where M denotes the number of parameters in millions. Despite having only 80% of the parameter count, VQMoE consistently achieves competitive performance compared to SMoE across all scales. This highlights the efficiency and robustness of our approach. The results are illustrated in Figure 3al and Figure 3b which show VQMoE's performance on the Enwik8 and Text8 datasets, respectively.

---

Published in Transactions on Machine Learning Research (08/2025)
1.52-
Model
1:64
Model
1.5
SMoF
1.42
SMor
(204 Punamn Late))
1.4-
VQMkeE(20%) Punumn Lene)
1.42-
1.36
1.4-
1.32-
1.36
1.25
12-
1.3-
1.22
1.28
1.18
10M
1394 20H
10H
15M
20M
(a) Robust VQMoË Benchmark (Enwik8)
(b) Robust VQMoE Benchmark ("Text8)
Figure 3: Illustration of the proposed Robust VQMoE architecture for Pre-training on Enwik8 and Text8 datasct. (a) Robust VQMoE architecture achieves the same performance with the routing methods while only using 80% of the parameters Enwik8 dataset.
(b) Roubust. VQMoF demonstrates robustness the Text8 dataset. Bits-per-character (BPC) the Enwik8 and Text8 datasets, and lower is better.
Architecture
FLOPs(x10'))
Transformer
Transformer-X XL
Avg.
Dataset
SST-2
SST-5
IMDB
BANKING77
SST-2
SST-5
IMDB
BANKING77
VQMoE
5.6145
82.6
41.1
89.5
84.8
83.3
42.0
85.3
74.72
SMoE
7.7620
82.1
39.5
89.3
82.6
80.8
40.4
88.6
80.2
72.94
SMoB-Dropout
7.7620
81.3
39.6
88.9
77.9
81.8
40.0
89.1
77.3
72.00
MMo
7.7620
82.4
39.9
83.1
81.3
40.3
88.7
82.7
73.43
StableMoK
7.7620
82.2
40.4
89.1
82.7
82.5
41.1
88.5
78.6
73.89
Table 2: Accuracy of the model after fine-tuning various datasets. Higher is better; best results are in bold.
5.3 Parameter-Efficient Fine-Tuning
Q4: What is the biggest advantage of VQMoE, compared to the conventional SMoE? A4: Craumeter-fffcieenn Fine-Turee Tuning.
We see that the discrete representation that VQMoE learns at the Pretraning stage 5.2 might consist of rich knowledge. To test. this hypothesis, use only the discrete representation for downstream tasks, allowing VQMoE to save 28% of computational resources compared to SMoE. Table reports the accuracy of the models fine-tuned on the test sets of various datasets. Overall, we observe that VQMoE demonstrates strong transfer learning capabilities by achieving the highest accuracy on all datasets. Notably, on the more challenging datasets of SST-5 and BANKING77, which have fewer training samples or more classes, we observe larger performance gains from VQMoE versus the SMoE baseline (over 2.5% improvements compared to SMoE on average). This result shows that VQMoE can learn a discrete representation that is not only good for pre-training but also exhibits strong transfer capabilities various downstream tasks,
5.4
Vision
Q5: Can VQMoE compete with SMoE in the Vision domain? ? A Yes.
To make performance comparison informative and comprehensive, consider two kinds of baselines that are fairly comparable to VQMoE: (1) Dense Model (Vision Transformer) (Dosovitskiy et al., 2021):
(2) SoftMoE |Puigcerver et al.) configurations for training the Mixture of Experts: (1) small 10 million parameters (10M); (2) large 110 million porameters (110M). The result at Table shows that VQMoE outperforms both Vision Transformer Dense (Dosovitskiy et al., 2021). SoftMoE (Puigcerver al.) 2024), and other routing methods such
2024 the most advanced MoE in Vision domain. We perform two
Dai et al. 2022), (Chi et al. 2022) six out of eight tasks across four image classification datasets. We
10

---

Published in Transactions on Machine Learning Research (08/2025)
Consistent Score (h))
(a) Consistent Score.
(b) VQMoË Representation.
(c) SMoE Representation.
Figure 4: Analysis Inconsistent Expert Selection and Representation Collapse issues when training SMoE. Figure 4a demonstrates consistent score movement from VQMoE, compared with SMoE and XMoE. Figure 1b and Figure 4c visualize the representation by experts in 2D dimension using Principal Component
Analysis (PCA) method.
conduct our experiments three times on four datasets (CIFAR-10, CIFAR-100, STL-10, and SVHN) using different seeds, reporting the average results along with the standard deviation. For the large-scale dataset ImageNet-1K,, perform single run due to resource constraints. The average performance of our method surpasses other baselines and is more stable, as indicated by the low standard deviation.
Architecture
Vision Transformer (Small)
Vision Transformer (Large)
Average
# params
10M
110M
Dataset
Cifar10
Cof
STL-10
SVHN
ImageNet-LK
Cifar10
Cifer100
STL-10
SVHN
ImageNet-
VQMoE
89.7
67.3
66 66.5
95.6
54.8
92.8
67.0go5
64.3
96.0202
76.5
SMoE
88.7
65.4g0.
66.4g0
95.4g0.1
52.8
85.7
55.52ss
64.4
94.5201
71.0
74.0
XMoE
88.820.2
65.5gn.5
66.3
95.420.1
52.5
87.1
55.92n.
64.6g0.
94. 94.1202
70.8
74.2
StableMoE
68 g8.8
65.5ga.
66.5goc2
95. 96.4.
52.5
81.7110.
55.5g1.
6133
94.5g0a
70.6
73.8
SoflMoE
65.6go3
61.4ga3
65.4g0
94.8g0.1
41.6
80.3ga.
42.9
63.22ns
93.5ga.1
68.2
69.7218
VIT (Dense)
89.0go.2
65.7.0.3
66.6ga.2
95.6go.1
52.2
92.2g03
60.2
61.1gos
96.0ga.1
71.1
75.3
Table 3: Accuracy of models evaluated on vision datasets. Higher is better, the best results are in bold.
5.5
In-depth Analysis
Consistent Score. Figure illustrates that expert selections when training SMoF. face inconsistent problems. As the Theorem exceeds that of the Transformer representation. Figure 44]also also shows that our method achieves the highest consistency score compared to the SMoE and XMoE models. However, the VQMoE model's consistency score is around 75%, as our method also requires learning continuous representation during the Pre-training phase.
this inconsistency arises because the router's coverage rate significantly
Representation Collapse issue. To visualize the Representation collapse problem in practice, we apply Principal Component Analysis (PCA) method to reduce from d dimension of the Transformer to 2D for plotting purposes, thanks to Chi et al. 2022). Figures 4b and 4c show the expert representations from the pretrained VQMoE and SMoE models. The results suggest that VQMoE experiences less representation collapse in the expert space compared to SMoE. The analysis is in line with the theorem proof at Section However, projecting the d-dimonsional space onto 2D for visualization may lead to information loss.
5.6 Ablation Study
We examine the effectiveness of VQMoE across various hyper-parameter settings, with all experiments conducted using the base Transformer architecture on the WikiText-103 dataset.
Vector Quantization Method. To learn a discrete representation, we rescarch various types of Vector Quantization methods, inchuding VQVAE (van den Oord al. 2017). VQGAN (Yu et al.) 2022). LFQ (Yu) et al. 2023), and ResidualVQ (Yang et. al. 2023). We observe that VQGAN using cosine similarity for
11

---

Published in Transactions on Machine Learning Research (08/2025)
distance achieves good and stable results in practice as Figure 6a sionality also delivers strong performance and exhibits robustnnss.
Interestingly, VQGAN with lower dimen-
Number of codebook impact. The number of codebook entries is a crucial hyperparameter when training Vector Quantization techniques. As shown in Figure 6b, we can see the best performance when the number of codebook entries matches the number of experts. This aligns with the proof by (Dikkala et al., 2023 which demonstrates that in the optimal case, the number of clusters equals the number of experts.
Sensitiveness of VQ loss contribution contribution of the Vector Quantization loss to the overall loss. If a is too high, it leads to a better discrete representation but may negatively affect the final target. Conversely, if a is too low, it may result in poor discrete representation. Therefore, a should be selected based on the data, lypically within the range of (0.05, 0.15).
Figure 6c illustrates the impact of a, which controls the
6
Conclusion and Future Directions
This study illustrates Vector-Quantized Mixture of Experts (VQMoE), novel and theoretically-grouneed architecture, to overcome challenges in training SMoE such as representation collapse and inconsistency. We evaluate our method on various Pre-training and Fine-tuning tasks, for both language and vision domains. The results show that VQMoË outperforms the routing methods both theoretically and empirically. Further- more, fine-tuning VQMoE with the discrete representation for rownstream lasks could reduce computational resource usage by 28%. We believe that focusing on discrete representation learning will offer a promising strategy for training and testing sparse mixtures of experts (SMoE) at a large scale. Finally, we believe that our approach opens up new research avenues for effectively training SMoE, where cutting-edge techniques in discrete representation learning and vector quantization can be harnessed to enhance their performance.
References
Mikel Artetxe, Shruti Bhosale, Naman Goyal, Todor Mihaylov, Myle Ott, Sam Shleifer, Xi Victoria Lin, Jingfei Du, Srinivasan Iyer, Ramakanth Pasunuru, Giri Anantharaman, Xian Li, Shuohui Chen, Halil Akin, Mandeep Baines, Louis Martin, Xing Zhou, Punit Singh Koura, Brian O'Horo, Jeff Wang, Luke Zettlemoyer, Mona Diab, Zornitsa Kozareva, and Ves Stoyanov. Efficient large scale language modeling with mixtures of experts, 2022.
Iñigo Casamevaa, Tadas Temčinas, Daniela Gerz, Matthew Henderson, and Ivan Vulić. Efficient intent detection with dual sentence encoders. In Proceedings of the 2nd Workshop on Natural Language Processing for Conversational AI, pp. 38-45, Online, July 2020. Association for Computational Linguistics. doi: URL httpps:/aadlaattooooo....//// ott.s/atttooooooooooooooo.aaaaa. mlp4ccovaiii1.5
Ciprian Chelba, Tomas Mikolov, Mike Schuster, Qi Ge, Thorsten Brants, Phillipp Kochn, and Tony Robin- son. One billion word benchmark for measuring progress in statistical language modeling, 2014. URL
httpe://axxiv..r/aabbe. //////112..0000 3005
Tianlong Chen, Zhenyu Zhang, Ajay Jaiswal, Shiwei Liu, and Zhangyang Wang. Sparse moe as the new dropout: Scaling dense and self-slimmnable transformers, 2023a.
Zitian Chen, Yikang Shen, Mingyu Ding, Zhenfang Chen, Hengshuang Zhao, Erik G. Learned-Miller and Chuang Gan. Mod-squad: Designing mixtures of experts modular multi-task learners. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 11828 11837, June 2023b.
Zixiang Chen, Yihe Deng, Yue Wu, Quanquan Gu, and Yuanzhi Li. Towards understanding the mixture- of-experts layer in deep learning. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (ods.), Advances in Neural Information Processing Systems, volume 35, pp. 23049-23062. Cur- ran Associates, Inc., 2022. URL https:::/p:pooooot http:::/pproceeiiggg.oo Pharroo722etbsssaaaaefffffPooooe http://pwwoooiiiggonorrrr.../c/p http://pppp...eeee//////////////
12

---

Published in Transactions on Machine Learning Research (08/2025)
Zewen Chi, Li Dong, Shaohan Huang, Damai Dai, Shuming Ma, Barun Patra, Saksham Singhal, Payal Bajaj. Xia Song, Xian-Ling Mao, Heyan Huang, and Furu Wei. On the representation collapse of sparse mixture of experts, 2022.
Rewon Child, Scott Gray, Alec Radford, and Ily Sutskever. Generating long sequences with sparse trans- formers, 2019. URL https://arxile http:///axxvv.orraaa//ao wwt. http//b////90/// 10509
Adam Coates, Andrew Ng, and Honglak Lee. An Analysis of Single Layer Networks in Unsupervised Feature Learning. In AISTATS, 2011. http:://saaaateeeoooo///oo//ttta
2011.pdf
Damai Dai, Li Dong, Shuming Ma, Bo Zheng, Zhifang Sui, Baobao Chang, and Furu Wei. Stablemoe: Stable routing strategy for mixture of experts, 2022.
Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc Le, and Ruslan Salakhutdinov. Transformer- XL: Attentive language models beyond a fixed-length context. In Anna Korhonen, David Traum, and Lluís Marquez (eds.), Proceedings of the 57th Anmual Meeting of the Associalion for Computational Linguislies, pp. 2978 2988, Florence, Italy, July 2019a. Association for Computational Linguistics. doi: 10.18653/v11 /v/ P19- 1285. URL dettps//accantiiee wwwpaaathooooooooooo..oo... org/P19-1285
Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc V. Le, and Ruslan Salakhutdinov. Transformer-xl: Attentive language models beyond fixed-length context, 2019b.
Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc V. Le, and Ruslan Salakhutdinov. Transformer-xl: Attentive language models beyond a fixed-length context, 2019c. URL https://arxiv. org/ 2t55ab//1990100880
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEFE conference on computer vision and paltern recognition, pp. 2009.
Teee,
Nishanth Dikkala, Nikhil Ghosh, Raghu Meka, Rina Panigrahy, Nikhil Vyas, and Xin Wang. On the benefits of learning route in mixture-of expertsss models. In Houda Bouamor, Juan Pino, and Kalika Bali (eds.), Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pp. 9376 9396, Singapore, December 2023. Association for Computational Linguistics. doi: 10.18633/v1/2023 emnlp-main.583. URL https:://aclanthollee http://claacbbllloooooo/r////../ www.p-mnniiiii..e www.llmmmiinnamm
Giang Do, Khiem Le, Quang Pham, TrungTin Nguyen, Thanh-Nam Doan, Bint T. Nguyen, Chenghao Liu, Savitha Ramasamy, Xiaoli Li, and Steven Hoi. Hyperrouter: Towards efficient training and inference of sparse mixture of experts, 2023.
Giang Do, Hung Le, and Truyen Tran. Simsmoe: Solving representational collapse via similarity measure, 2024. URL htttp:://arat www.vv.org/bbbr httg//bb/24406 15883
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Un- terthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale, 2021. URL https: //arxiv///xxvv...// org/ abs/ abs//a10......///a./... 2010. 11929
Nan Du, Yanping Huang, Andrew M. Dai, Simon Tong, Dmitry Lepikhin, Yuanzhong Xu, Maxim Krikun, Yanqi Zhou, Adams Wei Yu, Orhan Firat, Barret Zoph, Liam Fedus, Maarten Bosma, Zongwei Zhou, Tao Wang, Yu Emma Wang, Kellie Webster, Marie Pellat, Kevin Robinson, Kathleen Meier-Hellstern, Toju Duke, Lucas Dixon, Kun Zhang, Quoe V Le, Yonghui Wu, Zhifeng Chen, and Claire Cui. Glam: Efficient scaling of language models with mfxturo-of-ccpeetss, 2022.
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity, 2022.
13

---

Published in Transactions on Machine Learning Research (08/2025)
Mor Geva, Roci Schuster, Jonathan Berant, and Omer Levy. Transformer feed-forward layers are key- value memories. Tn Marie-Francine Moens, Xuanjing Huang, Lucia Specia, and Scott. Wen-tau Yih (eds.); Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pp. 5484 5495, Online and Punta Cana, Dominican Republic, November 2021. Association for Computational Linguistics. doi: 10. N0.18653/v 10/03/v11///// /2221 Lemnlp-main.445 URL
org 2021 eenlp-main .46
Yunhao Gou, Zhili Liu, Kai Chen, Lanqing Hong, Hang Xu, Aoxue Li, Dit- Yan Young, James T. Kwok, and Yu Zhang. Mixture of cluster-conditional lora experts for vision-language instruction tuning, 2024. URL
https://arxix http:s///arri..ogg/a/////// wwtg/bsg////aeee 2312. 12379
David Ha, Andrew Dai, and Quoe V. Le. Hypernetworks, 2016.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition, 03385
2015. URL http:::://axxi..rr/bbbsss
Weihua Hu, Takeru Miyato, Seiya Tokui, Eicch Matsumoto, and Masashi Sugiyama. Learning discrete representations via information maximizing self-augmented training. In Doina Procup and Yoe Whye Teh (eds.), Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pp. 1568-1567. PMLR, 06-11 Aug 2017. URL httpps://proceedings mlr.
press/v70/hu17b.. html
Changho Hwang, Wei Cui, Yifan Xiong, Ziyue Yang, Ze Liu, Han Hu, Zilong Wang, Rafael Salas, Jithin Jose, Prabhat Ram, Joe Chau, Peng Cheng, Fan Yang, Mao Yang, and Yongqiang Xiong. Tutel: Adaptive mixture-of-expprts at scale, 2023.
Robert A. Jacobs, Michael I. Jordan, Steven J. Nowlan, and Geoffrey E. Hinton. Adaptive mixtures of local experts. Neural Computation, 3(1):79-83 -7, 1991a. doi: 10. wtt.////cc..99........
Robert A. Jacobs, Michael I. Jordan, Steven J. Nowlan, and GeofTrey E. Hinton. Adaptive mixtures of local experts. Newral Computation, 3(1):79-8 87, 1991b. doi: 00.1162/nec..191.1
Albert Q. Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, Gianna Lengvel, Guillaume Bour, Guillaume Lample, Lélio Renard Lavaud, Lucile Saulnier, Maric-Anne Lachaux, Picrro Stock, Sandeep Subramanian, Sophia Yang, Szymon Antoniak, Teven Le Scao, Théophile Gervet, Thibaut Lavril, Thomas Wang, Timothée Lacroix, and William El Sayed. Mixtral of experts, 2024.
Michael Jordan and Robert Jacobs. Hierarchical mixtures of experts and the. Newral compulation, 6:181-, 1994.
Jean Kaddour, Joshua Harris, Maximilian Mozes, Herbie Bradley, Roberta Raileanu, and Robert McHardy. Challenges and applications of large language models, 2023.
Diederík P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization, 2017.
Diederik P Kingma and Max Welling, Auto-encoding variational bayes, 2022. URL
org/
abs htts/3122.6144
Jakub Krajewski, Jan Ludziejewski, Kamil Adamczewski, Maciej Pióro, Michal Krutul, Szymon Anto- niak, Kamil Cicbiera, Krystian Król, Tomasz Odrzygóádá, Piotr Sankowski, Marek Cygan, and Scbastian Jaszezur. Scaling laws for fine-grained mixture of experts, 2024.
Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, UoT, 2009.
Yoohwan Kwon and Soo- Whan Chung. Mole Mixture of language experts for multi-lingual automatic speech recognition. In ICASSP 2023 2023 IEEE International Conference Acoustics, Speech and Signal Processing (ICASSP), pp. 1-5, 2023. doi: 11010///CAASP0033222220....0
Margaret Li, Suchin Gururangan, Tim Dettmers, Mike Lewis, Tim Althoff, Noah A. Smith, and Luke Zettlemoyer. Branch-train-merge: Fmbarrrssinggy parallel training of expert language models, 2022.
14

---

Published in Transactions on Machine Learning Research (08/2025)
Bin Lin, Zhenyu Tang, Yang Ye, Jiaxi Cui, Bin Zhu, Peng Jin, Jinfa Huang, Junwu Zhang, Munan Ning, and Li Yuan. Moe-llava: Mixture of experts for large vision-language models, 2024.
Iyya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts, 2017. URL https://arxiv.orr
1608. 03983
Andrew L. Maas, Raymond E. Daly, Peter T. Pham, Dan Huang, Andrew Y. Ng, and Christopher Potts. Learning Word Vectors for Sentiment, Analysis. In Proceedings of the 49th Annual Meeling of the Associa- tion for Computational Linguistics: Human Language Technologies, pp. Portland, Oregon, USA, June 2011. Association for Computational Linguistics. URL https://aclanttoooogy..oo 1ttp///ogyy.///111011111
Matt Mahoney. Large text compression benchmark, 2011. URI. html
Chengzhi Mao, Lu Jiang, Mostafa Dehghani, Carl Vondrick, Rahul Sukthankar, and Irfan Essa, Discrete
representations strengthen vision transformer robustness, 2022. URL https://arxiv wwg. abs/21111
10493
Fabian Mentzer, David Minnen, Eirikur Agustsson, and Michael Tschannen, Finite scalar quantization: Vq-vae made simple, 2023. URL httpp:://axxi...oo owap/iii...oiiioooo 2tbbb//20055555
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models, 2016. URL
Niklas Muennighoff, Nouamane Tazi, Loic Magne, and Nils Reimers. Mteb: Massive text embedding bench- mark. arXic preprint
2022.
Niklas Muennighoff, Luca Soldaini, Dirk Groeneveld, Kyle Lo, Jacob Morrison, Sewon Min, Weijia Shi, Pete Walsh, Oyvind Tafoord, Nathan Lambert, Yuling Gu, Shane Arora, Akshita Bhagia, Dustin Schwenk, David Wadden, Alexander Wettig, Binyuan Hui, Tim Dettmers, Douwe Kicla, Ali Farhadi, Noah A. Smith, Pang Wei Koh, Amanpreet Singh, and Hannaneh Hajishirzi. Olmoe: Open mixture-of-experts language models, 2025. URL https:////// org/ abs 2409. 02060
Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Y Ng. Reading digits in natural images with unsupervised feature learning. NIPS Workshop, 2011.
Quang Pham, Giang Do, Huy Nguyen, TrungTin Nguyen, Chenghao Liu, Mina Sartipi, Binh T. Nguyen, Savitha Ramasany, Xiaoli Li, Steven Hoi, and Nhat Ho. Competesmoe effective training of sparse mixture of experts via competition, 2024.
Joan Puigcerver, Carlos Riquelme, Basil Mustafa, and Neil Houlsby. From sparse to soft mixtures of experts, 2024. URL https://arxxv..orr org/ abs taabs/2308.0me 22////2000.00000
Carlos Riquelme, Joan Puigcerver, Basil Mustafa, Maxim Neumann, Rodolphe Jenatton, André Susano Pinto, Daniel Keysers, and Neil Houlsby. Scaling vision with sparse mixture of experts, 2021a. URL
https://arxiv/ arg/abs//en 2106. 05974
Carlos
Riquelme, Joan Puigcerver, Basil Mustafa, Maxim Neumann, Rodolphe Jenatton, André Susano Pinto, Daniel Keysers, and Neil Houlsby. Scaling vision with sparse mixture of ex- perts. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 8583 8595. Curran Associates, Inc., 2021b. URL
neurips. hew.pee.carpscc.ppppprrr files//ppeer/22///iie// pdf.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparscly-gated mixture-of-experts layer, 2017.
Sheng Shen, Le Hou, Yanqi Zhou, Nan Du, Shayne Longpre, Jason Wei, Hyung Won Chung, Barret Zoph, William Fedus, Xinyun Chen, Tu Vu, Yucxin Wu, Wuyang Chen, Albert Webson, Yunxuan Li, Vincent Zhao, Hongkun Yu, Kurt Keutzer, Trevor Darrell, and Denny Zhou. Mixture-of-expeers meets instruction tuning: winning combination for large language models, 2023a.
15

---

Published in Transactions on Machine Learning Research (08/2025)
Sheng Shen, Zhewei Yao, Chunyuan Li, Trevor Darrell, Kurt Keutzer, and Yuxiong He. Scaling vision- language models with sparse mixture of experts. In Houda Bouamor, Juan Pino, and Kalika Bali (eds.); Findings of the Association for Computational Lònguistics: EMNLP 2023, pp. 11329 11344, Singapore, December 2023b. Association for Computational Linguistics. doi: 10.18653/v1 httg://////22..nnnnnnmnmmo findings-cmnlp.755 URL httpp:://clanthhooog.. wrw.egy..rg/22333.// Sindings-ennll.755
Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D. Manning, Andrew Ng, and Christopher Potts. Recursive Deep Models for Semantic Compositionality Over a Sentiment Tree- bank. In Proceedings of the 2013 Conference Empirical Methods in Natural Language Processing, pp. 1631-1642, Seattle, Washington, USA, October 2013. Association for Computational Linguistics. URL
Robin Strudel, Ricardo Garcia, Ivan Laptev, and Cordelia Schmid. Segmenter: Transformer for semantic segmentation, 2021. URL https://arxiv: wwwpsriii..rr.ooooooooooo wwt//ogg//////eeee/oo abs//a/5..... 05633
Aaron van den Oord, Oriol Vinyals, and koray kavukcuoglu. Neural discrete representation learn- ing. In Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL bttps://proceeddngg hetprsenurrsp.. meurips.cc/papeeeeeeee wwtppcppppataaaeeeeeeeeeee / Tslaaleeeeaaa0002200000000110000
pdf
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, L ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Newral Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL /attps::/proceediigss httte//pporooo////1////fff266660
Wenxuan Wang, Guodong Ma, Yuke Li, and Binbin Du. Language-routing mixture of experts for multilingual and code-switching speech recognition, 2023.
Fushao Xue, Zian Zheng, Yao Fu, Jinjie Ni, Zangwei Zheng, Wangchunshu Zhou, and Yang You. Openmoe An early effort on open mixture-of-experr language models, 2024.
Dongehao Yang, Songxiang Liu, Rongjie Huang, Jinchuan Tian, Chao Weng, and Yuexian Zou. Hifi-codec: Group-residual vector quantization for high fidelity audio codec, 2023. URL /sttps//arxiv..oo org/abs/ 2305.02765
Hanrong Ye and Dan Xu. Taskexpert: Dynamically assembling multi-task representations with memo- rial mixture-offexperts. In Proceedings of the IEFE/CVF International Conference on Computer Vision (ICCV), pp. 21828 21837, October 2023.
Jiahui Yu, Xin Li, Jing Yu Koh, Han Zhang, Ruoming Pang, James Qin, Alexander Ku, Yuanzhong Xu, Jason Baldridge, and Yonghui Wu. fector-quantizer image modeling with improved vggan, 2022. URL
https: https://arxiv: http:://xxivv.org
2ttp//110..444.. 04627
Lijun Yu, Yong Cheng, Kihyuk Sohn, José Lezama, Han Zhang, Huiwen Chang, Alexander G. Hauptmann, Ming Hsuan Yang, Yuan Hao, Irfan Essa, and Lu Jiang. Magvit: Masked generative vidco transformer, 2023. URL https: /tarxi:///rxi.... org/ 2btp/212....... 05199
Bolei Zhou, Hang Zhao, Xavier Puig, Tete Xiao, Sanja Fidler, Adela. Barriuso, and Antonio Torralba. Semantic understanding of scenes through the ade20k dataset, 2018. URL https://arxiv..r abs//
05442)
Yanqi Zhou, Tao Lei, Hanxiao Liu, Nan Du, Yanping Huang, Vincent Zhao, Andrew M Dai, zhifeng Chen, Quoe Le, and James Laudon. Mfxture-of-expert with expert choice routing. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (ods.), Advances in Neural Information Processing Systems, volume 35, pp. 7103-7114. Curran Associates, Inc., 2022. URL / paper_
16

---

Published in Transactions on Machine Learning Research (08/2025)
Yanqi Zhou, Nan Du, Yanping Huang, Daiyi Peng, Chang Lan, Da Huang, Siamak Shakeri, David So, Andrew Dai, Yifeng Lu, Zhifeng Chen, Quoc Le, Claire Cui, James Laudon, and Jeff Dean. Brainformers: Trading simplicity for efficiency, 2024.
Barret Zoph, Irwan Bello, Sameer Kumar, Nan Du, Yanping Huang, Jeff Dean, Noam Shazeer, and William Fodus. St-moc: Designing stable and transfcrable sparse expert models, 2022.
A
Appendix
Supplementary Material for "On the Role of Discrete Representation in Sparse Mixture of Experts"
This document is organized as follows. Appendix A.1 provides a detailed proof for Section 4 Appendix A.2 presents additional experimental results demonstrating the effectiveness of our method compared to the baselines. Finally, Appendix A.3 offers an in-depth analysis of representation collapse, while Appendix A.4 details the implementation aspects.
A.1 Proof for Results in Section
A.1.1 Jacobian Matrix of VQMoE
To investigate whether VQMoE alleviates this collapse, we derive the Jacobian of the VQMoE output with respect the input x E Raxd;
SSMB dg.(x)
VQMEE gx(x)
j (SMoB (x)
dx
K
0ga(1)) Ox:
+ ga(x:)
J FFN
PP (ii)
1=1
N
K
J1
+
T gme
Σ diej
(12)
Ø=1
me(ad)
1=1
J1
c
die
e
mE (c,d)
N+K+2
Ji
Σ
Oje,
T
j=l
where:
ej
Embedding of the j-th expert the SMoE;
Ji = S(x) FPN
Jacobian of the top-k FFN block;
As in SMoE, the Jacobian of VQMoE consists of two major components: J, the primary contribution from the input and selected expert; and N o,eo additional gradient contributions from both the continuous part and the discrete part..
17

---

Published in Transactions on Machine Learning Research (08/2025)
A.1.2 Proof of Theorem 4.3
In this proof, we use contradiction to establish the theorem. Assume that the expert embeddings e form a consistent router. By Delinition 4. we have:
dist(Tp. u.) I min(dist(zp. uz));
where is the reprosentation corresponding to the closest expert 6;.
According to (Chi et al. 2022). projecting information from a hidden representation space IR to the expert dimension leads to representation collapse. Now, consider three output of Multi-Head Attention (MHA) (MHA) layer: 21, $1.22,43 4y E Rd, belong to experts whose embeddings e1, 62. C3 collapse. Without loss of generality, assume that e2 lies between 61 and eg in the embedding space. Then, have:
dist(x2, u2) & min(dist(z), e1), dist(22, e2), dist(x3, ea))
(13)
I dist(e1, ez).
Let te denote the step at which the embeddings 61 and converge, and tm denote the step at which the Multi-Head Attention (MHA) module converges. From step te, it follows that:
lim dist u2) lim dist(e1, = 0. tm
t Hm
Thus, y (the output of MHA) converges step le-
This directly contradicts the assumption that the MHA converges at step tm: where to tm.
A.1.3 Proof of Proposition 4.4
We use contradiction to prove the proposition. Assume that, at training step t, there exists a set of pairs (C,, Ey) such that i * j. Let 02,.2,.... EN represent a sequence of inputs sampled from N clusters. From step to to step 1m-1, each pair (I,, Ej), where j [I, is updated using the following gradient descent equation:
Wh++ E
Wh: TT(
where Whe is the weight of expert E, at iteration t,, [Ix,) is the Jacobian matrix with respect to input Zy. and 17 is the learning rate, and 0 - 7 < m 1.
Let L denote the loss function during the training process described by Equation g After Im-1 training steps, the following condition holds:
((E((e))) min ((E.(z))). cE[L.N]
Under the assumption of contradiction, there exists a set of pairs, where Ty is assigned to an expert Eg: (zy, ; i,j [] and i + ji where the loss function C is minimized. means:
((E,(z)) s ((E((a))
However, by definition of the loss minimization process, the inequality
((E((,))) ((F(())))
must hold.
This leads to contradiction with our initial assumption.
18

---

Published in Transactions on Machine Learning Research (08/2025)
Scale
TopK
# Experts
SMoF
VQMoE (Discrete Only)
1
16
1.28
1.25
2
16
1.26
Base 20M-50K Steps
4
16
1.26
8
16
1.27
16
16
1.27
1
1.22
1.18
2
1.20
Base 20M-100K Steps
4
1.21
8
16
1.21
16
16
1.21
1
64
1.12
1.14
2
64
1.09
4
64
1.09
Large (210M)
8
64
1.09
16
64
1.10
32
64
1.10
64
64
1.12
Table 4: Performance comparison of SMoE and VQMoE (Discrete Only) the Enwik8 (BPC) dataset.
Scale
TopK
Experts
SMoE
MoCLE
VQMoE
1
16
1.28
1.29
1.25
2
16
1.26
1.28
Base 20M-50K Steps
1.26
1.28
8
1.27
1.28
16
16
1.27
1.28
Table 5: Performance comparison of VQMoE and MoCLE (Clustering approach) the Enwik88 (BPC) dataset.
A.2 Additional Experiment Results
Q6: Can VQMoE learn Discrete Representation Only from scratch? A6: Yes for small scale, but no for large scale.
The answer is yes for small models. However, training discrete representation-only approach is feasible primarily for small-scale models with moderately sized dataset. The results of the Transformer-XL. model in Table 4 on the Enwik8 dataset support this observation. As the model scales up, relying solely discrete representation reaches its limitations, leading to performance below the SMoE baselines.
Q7: Can VQMoE outperform the clustering- approach such as KMean? A7: Yes.
We explored a clustering- approach -MoCLE Gou et al. 2024), but found unsuitable for our method. Unlike MoCLE, Vector Quantization allows the model greater flexibility in learning cluster representa- tions during training, making it competitive in practical applications. The training results using the Transformer-XL model the Enwik8 dataset are presented in Table 5
Q8: Can VQMoE contribute to AI real-world applications? A8: Yes.
We found that VQMoE can directly benefit real-world AI applications, such as image segmentation, demon- strating its strong generalization capabilitics. Specifically, our method outperforms both the bascline and dense models in terms of Mean Accuracy and mloU metries on the ADE20K dataset (Zhou et al. 2018 using the Segmenter (Srrud et. al. 2021 Detailed results are provided in Table
19

---

Published in Transactions on Machine Learning Research (08/2025)
Model
ViT
SoftMoe
SMoE
StableMoR
XMoE
VQMoF
Metrics
20.8
19.0
23.1
22.4
22.3
23.4
Mean accuracy
Segmenter
15.0
14.0
15.5
16.0
15.7
16.6
mloU
Table 6: Comparison of VQMoE versus the baselines on the ADE20K datasct.
Q9: Does VQMoE consistently outperform the baselines across multiple training runs? A9: Yes.
Due to resource constraints, it is challenging train all models across all datasets multiple times and to perform formal statistical significance testing. To illustrate the variance across multiple training runs, we train VQMoE and baseline models on the Text8 dataset three times. The average Bits-Per-Characee and standard deviation for each model are reported in Table 7 The results indicate that VQMoE achieves the best average performance, while also exhibiting a lower standard deviation compared to other models, suggesting greater training stability. The consistency obscrved across repeated runs supports the reliability of the results reported in Table
(BPC)
Table 7: Average BPC and standard deviation across three training runs Text8. Lower is better; best results are in bold.
Model
Algorithm
Dataset
Avg. BPC
Transformer-X
VQMoF
Text8
1.280 + 0.003
SMoE
1.293 + 0.007
XMoE
1.282 + 0.003
StableMoF.
1.285 + 0.005
Q10: Is VQMoE able to consistently surpass SMoE models in large-scale evaluation scenarios? A10: Yes.
We explore a more extensive model variant, OLMoE-IB-7B Mucnnighoff et al. 2025 which comprises 16 layers, 7 billion parameters, 64 experts, and lop-k selection of 8. Due to limitations time and com- putational resources, we utilize the pre-trained routers for codebook embedding and compare our proposed VQMoE with OLMoE in training-free setting. The evaluation is conducted across 6 diverse tasks and 19 datasets from the Massive Text Embedding Benchmark (MTEB) Mueunighoff et al. (022)). The summary of this evaluation is provided in Table X
Interostingly, we find that VQMoE consistently outperforms OLMoE across all tasks and datasets, despite not undergoing additional training or fine-tuning. On average, across six tasks, VQMoE shows relative improvement. of 4.6%. The most significant, gains appear in the Classification and Clustering tasks. These findings support our hypothesis that VQMoE enhances pre-trained models by learning more effective routing policies. Furthermore, by mitigating representation collapse through the use of discrete representations, VQMoE improves the model's overall representational capacity.
A.3 Representation Collapse Analysis
To illustrate Theorem movement of Expert Input Representation in Figure and and Expert Embedding (router) in Figure 5b We analyze the dynamics of the expert input represcntations by tracking their changes across training iterations. The results indicate that the inputs to the experts become increasingly divergent over time. This divergence suggests that the model learns to represent the data a more specialized and diverse manner, allowing each expert to focus on distinct features or patterns within the data. Similarly, we track the changes expert embeddings (router) throughout the training process. However, the trend is the opposite: the expert embeddings appear to converge quickly, stabilizing around 10,000 iterations. The findings align with
perform language model task described in Section A.4.2 examining the
20

---

Published in Transactions on Machine Learning Research (08/2025)
Table 8: Zero-shot performance comparison between OLMoE and VQMoE on MTEB. The best score per
dataset is highlighted in bold. Improvement (Imp.) calculated as (GAP / OLMoE) 100
Task
Dataset
Params
#Exp
Top-k
OLMoE
VQMoE
Imp. (%)
Emotion
7B
64
49.9
52.5
5.2
Classification
Toxic
7B
64
65.2
67.2
3.1
Tweet
7B
64
58.0
59.8
3.1
Medrxiv
7B
64
8
23.9
25.8
7.5
Clustering
20Groups
64
25.7
28.4
10.5
SemEval
7B
64
46.7
49.5
6.0
Pair Classification
URLCorpus
7B
64
8
77.4
79.4
2.6
Ask
7B
64
51.9
53.3
2.7
Reranking
SciDocs
7B
64
69.6
72.3
3.7
StackOver
7B
64
32.5
33.9
Biosses
7B
64
61.8
68.7
11.2
SickR
7B
64
65.7
66.5
1.4
STS12
7B
64
53.8
56.0
4.1
STS
STS13
7B
64
66.5
74.0
11.3
STS14
7B
64
56.8
59.5
4.6
STS15
7B
64
69.3
71.5
3.2
STS16
7B
64
70.1
70.5
0.6
Summarization
Medrxiv
7B
64
8
28.9
29.8
3.1
Average
54.1
56.6
4.6
4.3 indicating that Expert Embedding converges more quickly than Expert Input Representation. These results provide further evidence supporting the Theorem
assumption stated in Theorem
A.4
Experiments implementation details
This section provides detailed parameters of our experiments in Section
A.4.1
General Settings
The experiments are based on the publicly available SMoE-Dropout implementation (Chen et al. 2023a However, the pre-training was conducted on two H100 GPUs, so results might differ when using paralle training on multiple GPUs.
A.4.2 Pre-training Experiments
Table 9 provides the detailed configurations for pre-training Transformer Vaswani et al. 2017 Transformer-
XL Dai et al. 2019b) on Enwik8, Text8, WikiText -aiToxt-103,and One Billion Word.
A.4.3 Fine-tuning Experiments
For fine-tuning experiments, we employ the identical model architecture as in pre-training. Table I0 presents the detailed configurations utilized for fine-tuning experiments SST-2, SST-5, IMDB, and BANKING77 datasets. We start with the pretrained checkpoint of the base model on enwik8, remove the final layer, and replace it with two randomly initialized fully connected layers to serve as the classifier for each fine-tuning dataset. All methods are fine-tuned for 5,000 steps with a uniform learning rale.
http:///hhhtuuccco//////aarraarp Dropot
21

---

Published in Transactions on Machine Learning Research (08/2025)
Dataset
Input length
Batch size Optimizer
Lr
# Training Step # Experts TopK
Enwik8
512
48
Adam
3.5e-4
100k
16
2
Text
512
48
Adam
3.50-4
100k
16
2
WikiText-103
512
22
Adam
3.5e-4
100k
16
2
One Billion Word
512
11
Adam
3.5c-4
100k
16
2
Table 9: Hyperparameter settings for pre-training experiments on Enwik8, Text8 WikiText- and One
Billion Word.
Dataset
Input length
Batch size
Optimizer
Lr
# Epochs
SST-2
512
16
Adam
le4
5
SST-5
512
16
Adam
1e-4
5
IMDB
512
4
Adam
lc-4
5
BANKING77
512
16
Adam
1c-4
5
Table 10: Detail settings for fine-tuning experiments on the evaluation datasels.
(b) Training Router Representation (Expert embed- ding).
(a) Training Input Token Representations.
Figure& 5: Comparison of Token Representation and Expert Representation across Training Iteration.
631
(a) Vector Quantization method.
(b) Number of codebook.
(c) Impact of a for VQMoE.
Figure 6: Pre-training small Transformer-Xi XL on WikiText-103 across different hyperparameters.
22


Reviews:
[{"review_id": "oLKiuOuLb5", "reviewer": "Reviewer_SLrR", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths:**\n\nS1. Re-framing routing as vector-quantised clustering is conceptually clean and differs from previous “better-router” work. The idea of freezing the discrete path for parameter-efficient fine-tuning is appealing.\n\nS2. Results cover diverse datasets and fine-tuning tasks, with both base (20 M) and large (210 M) scales. VQMoE usually wins or ties on BPC/PPL and top-1 accuracy, and converges slightly faster.\n\nS3. The paper is generally well structured and key formulas are explicit. The ablation plots in the appendix help.\n\n---\n\n**Weaknesses and suggestions**:\n\nW1. My first concern is that the main performance results reported in Table 1 are generally marginal. I understand training multiple models and show variance is infeasible, but at least variance of PPL and BPC over the final training steps can be reported to see how noisy they are. \n\nW2. The paper is mainly motivated by resolving representation collapse in standard MoE. There is not enough evidence in the experiments to support this though. Figure 3.a shows a better consistent score but it’s really not obvious how it would impact the downstream performance. Can the author devise an experiment that cause collapse in standard MoE which can be mitigated by VQMoE? \n\nW3. Regarding the \\alpha parameter, the authors say “if $\\alpha$ is too high, it leads to a better discrete representation but may negatively affect the final target”. Can you elaborate on this? The proposed method is on the premise that discrete representation can be helpful. How it can have a negative affect? And what do you mean by final target?\n\nW4. I assume by $W_{FFN}$ you mean the fead forward network in Eq. 1. Please explain the notations.\n\nW5. Theorem 4.3 hinges on a strong separation of convergence times between MHA and experts (assumption $ t_m \\geq t_e$). Is this true in practical training regimes?\n\nW6. Can you please formally define robustness? It’s somehow vague in Figure 2.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper applies discrete representations for sparse mixture-of-experts to solve the problem of representation collapse in SMoEs. The paper suggests replacing the traditional learned router with a discrete codebook obtained via Vector Quantization. The authors argue that selecting experts through an intermediate discrete representation yields an optimal assignment. The paper provides some theoretical analysis on the optimality of their expert selection, and how representation collapse is alleviated. They also provide experimental result on language modelling and vision tasks.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "please see the above.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "qT6nmksGMb", "reviewer": "Reviewer_SGvX", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths: the paper proposes a novel SMoE architecture that doesn’t require a router.\n\nWeaknesses: some of the contributions are unclear and some of the claims aren’t supported by the data", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper presents a novel SMoE architecture that doesn’t require a router along with a theoretical and experimental analysis of the novel architecture showing that it performs on par with the standard methods while requiring fewer FLOPs for fine-tuning.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Critical:\n\n1. The contribution of the paper. The abstract states that VQMoE achieves 28% improvement of robustness while p. 9 states that it saves 28% of computational resources during fine-tuning. Is robustness equated with the amount of computational resources? \n\n2. My confusion partially stems from my lack of understanding of Q3 that appears to focus on robustness. First, I don’t understand the difference between VQMoE and VQMoE_robust in Fig. 2. As far as I can tell, VQMoE_robust is never mentioned anywhere in the paper. Second, it’s unclear to me how the answer ‘A3: The lower hidden size of FFN’ follows from Fig. 2 or the discussion in Q3. \n\n3. Q1 — I disagree with the authors’ claim that VQMoE performs better on pre-training tasks compared to routing method; at least not in all instances. From Table 1, it performs the same as SMoE on Enwik8 and Text8 (Large) and the same as XMoE and StableMoE on Text8 (Base). Generally, performance gains of VQMoE are very small. The authors note that router method performance increases are often marginal but VQMoE performance seems just as marginal to me. I’d like to see multiple runs over different random seeds with standard deviations (similar to Table 3) to put the performance gain in perspective. That said, even if this analysis reveals no performance gains, fewer FLOPs in fine-tuning still present an advantage.\n\n4. The authors need to provide more information about the pre-training tasks — how the data was tokenized, the motivation for including both character and word-level tasks, the motivation for conducting the fine-tuning experiments only on the models trained on enwik8 (I think they should be repeated at least on one word-level task to show generalizability, preferably on all datasets from Table 1), the motivation for including Transformer and Transformer-XL, how was topk=2 was chosen and whether the results hold with other topk choices. \n5. In Q1, the authors draw comparisons about model performance across character- and word-level dataset sizes. I don’t think this comparison can be made since the task is different if I understand it correctly (see my questions in #4 above).\n\nImprovements:\n\nVQMoE learns an additional representation. Does it come at a cost? \n\nIn Q1, the authors refer to increasing model complexity from small to medium. What does medium refer to in Table 1?\n\nI don’t follow Fig. 3 b and c —perhaps unpack in text why they’re showing representation collapse.\n\nTheorem 4.3 and Section 4.2 focus on the analysis of inconsistent router and representational collapse in SMoE. Would these considerations apply to XMoE, SMoE-dropout and StableMoE?\n\nSmaller comments:\nTable 1 — what is lm1b?\nFig. 2: the comparison across the two benchmarks would be easier if the y-axis was the same on both graphs.\nFig. 3 (a) — what is the x-axis? Can the tick labels be increased? It’s too small to read\nQ4: ‘advantage of SMoE, compared to the conventional SMoE’ —> advantage of VQMoE", "broader_impact_concerns": "n/a", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "p2oPZnjoBC", "reviewer": "Reviewer_Szmq", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n\n- The integration of vector quantization into the SMoE routing framework is novel and interesting.\n- Experiments span both language and vision domains, demonstrating the generality of the proposed approach.\n- The authors provide theoretical analysis and visualization that attempts to clarify the effectiveness of the method.\n\nWeakness\n\n- Clarity and readability are major concerns of this article. There are many inconsistent and messy notation and unclear symbol usage.\n- Figures are hard to interpret, same problem as the main paragraphs.\n- The experimental setting is limited to small-scale models, making it difficult to assess the scalability of VQMoE to large-scale deployments typical of MoE applications.\n- The claim of mitigating representation collapse is not convincingly supported by the provided analysis or visualization (e.g., Figure 3b is ambiguous and still show representation overlayed with each other).\n- Several theoretical results and propositions are under-explained or inconsistently labeled, and contain notational errors or unclear definitions (e.g., Proposition A.1 vs Proposition 4.4).", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes a novel architecture, which incorporates vector quantization techniques into sparse mixture of experts (MoE) models. The proposed method is evaluated on both language and vision tasks, and shows improved or comparable performance against existing MoE variants on a range of benchmarks. The paper also provides theoretical insights and attempts to analyze expert utilization behavior through visualization and propositions.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "No", "audience_explanation": "", "requested_changes": "I list my questions and requested changes here by the paper structure. I have read the whole paper including the appendix, and there are many confusing points to me.\n\nTypos\n\n1. Page 3, 3.1 Preliminaries, Discrete Representation Learning Line 2. Vector Quantisation → Vector Quanti**z**ation\n2. Page 3, 3.1 Preliminaries, Discrete Representation Learning Line 4. Let denote $z_v(x)\\in\\mathbb{R}^{n\\times d}$ denotes the output of the VQVAE and 1() is the indicator function → Let $z_v(x)\\in\\mathbb{R}^{n\\times d}$ **denote** the output of the VQVAE and 1() **be** the indicator function.\n\nSection 3 (Method)\n\n- Figure 1\n    - It’s better to add arrows when a line is between blocks to clarify data flow directions (e.g., tokens → gating layer).\n    - Please explain the difference between \"Top1 Router\" and \"Gating Layer\" in relation to Equation 4.\n    - Why is the width of discrete representations inconsistent?\n    - Please clarify the use of the $\\bigotimes$ symbol in Figure 1(b).\n    - How are codebook vectors mapped to experts if the codebook size ≠ number of experts?\n    - Please consider redesigning the figure for clarity and consistency.\n- Section 3.2\n    - Please use **consistent notation** throughout: e.g., $f_{SMoE}, f^{VQMoE}$.\n    - Please avoid bullet points; provide explanations in plain text.\n    - If the codebook size equals the number of experts, won't the same FFN see the **s**ame discrete input every time during inference?\n- Section 3.3\n    - Please consider to write the loss function explicitly, even if based on prior work.\n    - Please define $L^{l2}, L^{commitment}$ as they do not appear in the cited work.\n    - What does $L_{task}$ represents (cross-entropy or regression?) Please give some clues for pre-training and fine-tuning tasks.\n    - Also, avoid to have both $L^{l2}$ and $L_{task}$ at the same time, what is the difference of using subscript and superscript? If not difference, please stay consistent.\n\nSection 4 (Theoretical Analysis)\n\n- Section 4.1\n    - Please use consistent notation for k, K, and x across sections of whole paper, avoid to redefine them anywhere.\n    - Please clarify whether Proposition A.1 is the same as Proposition 4.4. If it is the same, please modify the parts to make it consistent.\n    - Please use consistent terminology: e.g., “Mixture-of-Experts” vs “mixture of experts”.\n    - I would like to know whether the propositions and theorems can be extended to top-k = 2, as used in experiments.\n    - Please avoid redundancy: “In practice, … In practice, …” (Page 6, Lines 7-8).\n- Section 4.2\n    - what is $o_j$?\n    - In equation 11, what is $m\\in\\{c, d\\}$? What is $d_l$?\n\nSection 5 (Experiments)\n\n- Section 5.2\n    - Please eliminate redundant sentences between \"Training tasks\" and \"Experimental settings\".\n    - Which model configuration is used in Table 2? (base or large?)\n    - Please consider to expand abbreviations (e.g., BPC, PPL) in Table 1.\n    - The claim that VQMoE converges faster is unsupported — Table 1 shows only final scores.\n    - In SMoE-Dropout, they split the dense model’s FFN to several networks, so that eventually the size is equals to the original dense model. What configuration do you use in SMoE-Dropout?\n    - Please improve clarity of Figure 2: distinguish icons, clarify legend, and explicitly state sizes used. You seems to use same icons for different models, and i can’t see other methods on the board, it seems like the number of experiments did not baseline is not equal?\n    - For Q3 robustness experiments, specify reduced FFN dimensions and explain the relationship with parameter count.\n    - I wonder  the inference process for top-k = 2 when using discrete tokens; this is not made clear in the current paper or in Figure 1.\n- Section 5.3:\n    - Please clarify score improvements in Table 2. The claim of 5% improvement on SST-5 and BANKING77 is inconsistent with actual deltas (0.7, 1.7, 0.9, and 2.6).\n- Section 5.5:\n    - In Figure 3(b) and 3(c), what does circle size represent? Is it the parameter size or output magnitude?\n    - The PCA visualization of Figure 3 does not clearly demonstrate that VQMoE avoids representation collapse; results in [1] are more compelling.\n- Section 5.6:\n    - Why use VQVAE instead of VQGAN, if VQGAN achieves lower PPL?\n    - Please clarify how many codebook vectors and FFNs are used during pretraining and fine-tuning.\n\nAppendix\n\nAppendix A.1.1\n\n- In Page 17, you said the proof of Proposition A.1 can be found in Appendix A, which is very weird. And actually it’s in A.1.3.\n\nAppendix A.1.2\n\n- Please use consistent notation: $\\mathbb{R}$ vs $\\mathcal{R}$\n- Please explain expert symbols $x, y, z$ clear, are they the weights (parameters) of three expert network?\n- Please clarify that what is expert embedding? Is it the codebook vector you use for each expert? Or the weights of linear router that correspond to each expert? Or the weights of that FFN expert itself?\n- I don’t understand equation 12 and the equation after it. Could you explain it clear to me? I can’t see why this prove Thorem4.3, also in practice, SMoE’s router won’t choose the expert that corresponding router embedding is closest to the token.\n\nAppendix A.1.3\n\n- Why there is step $t_0$ to step $t_{k-1}$ but also have iteration $l$? Is iteration and step represent different thing?\n- What does “pair” mean?\n\nAppendix A.2\n\n- Q7 Please briefly describe MoCLE, or it is not clear that why question mention K-means but you did not directly compare to it.\n\nFigure 4\n\n- What does the y-axis value fo Figure 4(a) and (b) mean? I can’t see why the inputs to the experts become divergent overtime.\n- Also, is Figure 4(b) indicates that all expert embedding become the same after training?", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)