On the Effectiveness of Discrete Representations in Sparse Mixture of Experts

We sincerely appreciate your valuable feedback. Please see our responses to your concerns below.

W1. Theorem 4.1 – The assumption that MHA and expert embeddings converge, converging with $t_m>>t_e$ does not appear to account for the possibility of representation collapse. A specific mathematical approach addressing the occurrence of representation collapse is missing.

A1.

As theorem 4.1 The assumption that MHA and expert embeddings converge, converging with $t_m>>t_e$ lead to inconsistent expert selections. From the consistent expert selections to expert representation collapse has been proved by [1].

W2. Proposition 4.2 – Here, an even less agreeable assumption is made. It assumes that input data can be divided into k clusters; however, it is well-known in the deep learning field that high-dimensional data, such as text or images, does not cluster cleanly. Furthermore, this claim expects the number of clusters to match the number of experts, which is unlikely to hold for real-world data.

A2.

• We build upon the assumptions of [2] and [3] that the data exhibits a clustering structure. Besides that, other works in MoE, such as [4] and [5], assume that the data contains specific patterns.
• From a theoretical perspective, [3] demonstrated that the Mixture of Experts (MoE) can be interpreted as a clustering problem. If the number of clusters is smaller than the number of experts, collapse may occur; conversely, if it is larger, the model could underfit. In practice, selecting an appropriate number of experts can be achieved through hyperparameter tuning tailored to the specific dataset.

W3. Additional experiments on high-res datasets, such as ImageNet, would be necessary to strengthen the findings.

A3.

We are currently training on the ImageNet dataset and will report the results once the process is complete.

Reference:

[1] Chi, Z., Dong, L., Huang, S., Dai, D., Ma, S., Patra, B., Singhal, S., Bajaj, P., Song, X., Mao, X., Huang, H., & Furu Wei, F. (2022) On the representation collapse of sparse mixture of experts https://arxiv.org/abs/2204.09179

[2] Chen, Z., Deng, Y., Wu, Y., Gu, Q., & Li. Y. (2022) Towards understanding mixture of experts in deep learning. https://arxiv.org/abs/2208.02813

[3] Dikkala, N., Ghosh, N., Meka, R., Panigrahy, R., Vyas, N., & Wang. X. (2023) On the benefits of learning to route in mixture-of-experts models. https://aclanthology.org/2023.emnlp-main.583.pdf

[4] Chowdhury, M., Zhang, S., Wang, M., Liu, S., & Chen. P. (2023) Patch-level routing in mixture-of-experts is provably sample-efficient for convolutional neural networks. https://arxiv.org/pdf/2306.04073

[5] Chowdhury, M., Wang, M., Maghraoui, K., Wang, N., Chen, P., &Carothers. C. (2024) A provably effective method for pruning experts in fine-tuned sparse mixture-of-experts. https://arxiv.org/abs/2405.16646

Dear Reviewer WyD6,

Thank you so much for your question. We would like to answer your questions as below:

For A1-1. I understand that you are trying to make the claim that representation collapse is caused by $t_m >> t_e$ .You need an experimental or theoretical proof of this.

A4.

The core theoretical contribution of our work is presented in Section 4.2, “VQMOE Solves Representation Collapse by Design,” where we demonstrate that our method outperforms other routing approaches in addressing representation collapse. Furthermore, we identify inconsistencies in routing selection in practice and support this observation through Theorem 4.1, under the assumption $t_m >> t_e$. This issue is also illustrated in Figures 4 and 5. Regarding the relationship between representation collapse and inconsistencies in routing, [1] provides visual evidence in Section 4.5 Analysis (Figures 2 and 3).

For A1-2. Additionally, I noticed a typo and missing notation explanation in the proof that follows.

A5.

Thank you for your suggestions. We have revised it and updated it in our new revision.

For A2-1. if, as you claim, real data already has a well-clustered structure,.

A6.

We think there is a misunderstanding here, our assumption is that data has cluster structure in a latent space, not well-clustered defined structure. Our method will learn the structure by the Vector Quantionzation approach. For example, images or text are unstructured data but they have a cluster structure in latent spaces as illustrated by [1], [2], [3], [4].

*For A2-2. collapse can be solved by tuning the number of experts through hyperparameter tuning?

A7.

No, we did not say that, our answer for your first part is "In practice, selecting an appropriate number of experts can be achieved through hyperparameter tuning tailored to the specific dataset."

For A2-4. Is it just that existing MoE methodologies suffer from collapse because they fail to properly tune the number of experts?,.

A8.

No, that is not correct. Even if the true number of experts equals the number of clusters in the hidden space, conventional SMoE still faces the collapse issue due to the projection of information from d_model (dimension of Transformer model) to number_of_experts as proof by [5].

For A2-5. What is the novelty of your paper?,.

A9.

While conventional SMoE models rely on a router to allocate input to each expert, we propose an alternative approach: learning a cluster structure and directing inputs to experts based on this structure, which we find to be more efficient. This is supported by theoretical results in Section 4.2 and experimental results in Section 5.

For A2-6. Proposition 4.2 oversimplifies the problem and consequently fails to highlight,.

A10.

Proposition 4.2 aligns with our method, as it demonstrates that learning a cluster structure and assigning each cluster consistently to the corresponding experts leads to an optimal solution.

W3. Additional experiments on high-res datasets, such as ImageNet, would be necessary to strengthen the findings.

A11.

VQMoE consistently outperforms both dense models and other routing methods on large-scale datasets like ImageNet, across various scales, from small to large. For further details, please refer to the General Response.

Reference:

[1] Li, Y., Hu, P., Peng, D., Lv, J., Fan, J., & Peng. X. (2024) Image clustering with external guidance. https://arxiv.org/abs/2310.11989

[2] Stephan, A., Miklautz, L., Sidak, K., Wahle, J., Gipp, B., Plant, C., & Roth, B. (2024) Text-guided image clustering. https://arxiv.org/abs/2402.02996

[3] Tipirneni, S., Adkathimar, R., Choudhary, N., Hiranandani, G., Amjad, R., N. Ioannidis, V., Yuan, C., & Reddy, C. (2024) Context-aware clustering using large language models. https://arxiv.org/abs/2405.00988

[4] Viswanathan, V., Gashteovski, K., Lawrence, C., Wu, T., & Neubig, G.(2023) Large language models enable few-shot clustering.https://arxiv.org/abs/2307.00524

[5] Chi, Z., Dong, L., Huang, S., Dai, D., Ma, S., Patra, B., Singhal, S., Bajaj, P., Song, X., Mao, X., Huang, H., & Furu Wei, F. (2022) On the representation collapse of sparse mixture of experts https://arxiv.org/abs/2204.09179

Thank you for your thoughtful review! Please find our responses below.

W1. In VQVAE, only one codebook is used. Why are multiple codebooks used in VQMoE, and what difference does this make?

A1.

We apologize for the confusion caused by the writing mistake. In our method, We only learn one codebook, $K$ is denoted for the number of codes in the codebook.

W2. If two tokens are very similar, how does VQMoE ensure they are not quantized to the same codebook?

A2.

Proposition 4.2 proves that, for input data with a clustering structure, the optimal expert selection is to route each cluster to its corresponding expert. Thus, if two tokens are very similar, consistently routing them to the same expert is the optimal choice. In contrast, traditional SMoE methods use a router to direct tokens, but as shown in Figure 5, the router converges quickly (typically around 10k iterations), resulting in inconsistent expert selections.

W3. The proofs of the theorem and definitions in the appendix B, particularly Figures 4 and 5, are unclear. It would be helpful if the authors provided a more detailed explanation of these figures.

A3.

Thank you for your suggestion and we have updated our revision. Figures 4 and 5 illustrate Theorem 4.1. In Figure 4, we visualize the change in expert representations (FFN) over training iterations, while Figure 5 shows the change in expert embeddings (router) over training iterations. Since the expert embeddings converge faster than the expert representations, this results in inconsistent expert selections.

W4. Performance testing on ImageNet is missing from the classification task evaluations.

A4.

We are training the baselines and our method on ImageNet and we will report the result to you as soon as possible.

W5. Although the paper states that experiments were conducted in the Vision domain, they are limited to classification tasks, with no experiments in areas like object detection or segmentation.

A5.

Our paper focuses on addressing the representation collapse and inconsistency problems in SMoE foundation models. Applications of the model to tasks such as object detection and segmentation are left for future work. However, [1] demonstrates that a stronger image classification backbone leads to improved performance in both object detection and image segmentation.

W6. The writing and formula expressions in the paper are sometimes unclear, with some statements potentially leading to misunderstandings.

A6.

Thank you for your valuable suggestions. We have reviewed and revised our paper to enhance its clarity and readability.

Reference:

[1] Goldblum, M., Souri, H., Ni, R., Shu, M., Prabhu, V., Somepalli, G., Chattopadhyay, P., Ibrahim, M., Bardes, A., Hoffman, J., Chellappa, R., Wilson, A., & Goldstein, T. (2023) Battle of the backbones: A large-scale comparison of pretrained models across computer vision tasks. https://arxiv.org/abs/2310.19909

Dear Reviewer aPBa,

Thank you so much for your constructive discussions. We would like to answer your questions as below:

1. W2-p1 In the quantization mechanism of VQMoE, if two data tokens are very similar, will they be forced to be quantized into the same codebook?

A7.

Our answer is yes. The Vector Quantization mechanism will learn a codebook to group similar elements by assigning them to the same code within the codebook. Proposition 4.2 proves that if there is an input data with a clustering structure, the optimal expert selection is to route each cluster to its corresponding expert.

1. W2-p2 will this result in some information loss or model capability degradation?

A8.

The discrete representation can help a model to learn a true data pattern and improve model generation, however, it can lead to information loss in some cases such as our answer for question 2 of the reviewer ckSP. To overcome the challenge, we propose to learn both continuous representation and discrete representation for the pretraining model (for large dataset). For downstream tasks (small datasets), we suggest that only discrete representations are efficient in terms of both performance and inference speed.

2. W4. I would like to see the effect on ImageNet.

A9.

VQMoE consistently outperforms the dense model and other routing methods on large-scale datasets such as ImageNet, across both small and large-scale versions. Please refer to the General Response for more detail.

3. W5. Are there relevant experiments on detection or segmentation? Can there be some simple experiments to demonstrate the superiority of your method?

A10.

For the image segmentation task, our method surpasses the baselines and the dense model in terms of Mean Accuracy and mIoU metrics on the ADE20K dataset using the Segmenter model[1]. Detailed results are provided below.

Reference:

[1] Robin Strudel, Ricardo Garcia, Ivan Laptev, and Cordelia Schmid. 2021. Segmenter: Transformer for semantic segmentation https://arxiv.org/abs/2105.05633

We sincerely appreciate the reviewer’s constructive suggestions and have provided our responses below.

W1. The definition of symbols is confusing, especially on Page 3.

A1.

We apologize for the confusion caused by the symbols and have corrected it at our new revision. We have revised them as follows: $g(x)_c(x) = col_0(G(x))$, $g(x)_d(x) = col_1(G(x))$ is gating function for continuous and discrete representation with $G(x) = \operatorname{softmax}(W_g^T \times x)$. $W_g^T \in \mathbb{R}^{2 \times d}$ is a learnable weight.

W2. The lacked comparison. Can the authors provide a comparison with the MoCLE method?

A2.

We experimented with a clustering-based approach similar to MoCLE[1], but found it unsuitable for our method. Compared with the MoCLE[1] approach, Vector Quantization provides the model with greater flexibility to learn cluster representations during training, making it more competitive than the clustering approach in practical applications. The training results on the Enwik8 dataset are presented in the Table below.

W3. The additional training need for the pre-training stage.

A3.

For the pretraining task, the continuous representation is learned using a top-1 router, while the discrete representation is learned through vector quantization. As a result, the total training resources are equivalent to those required for training a Top-2 router, as in [2], [3]. Therefore, no additional training is needed during the pretraining stage.

W4. This work must train the codebook in the pre-training stage. Thus, I kind of disagree with the claim made in Line 367. Although not training the routing in fine-tuning, the additional training time is needed for pre-training.

A4.

As stated in Theorem 4.1 of our paper, the conventional MoE approach, which selects experts based on a router, can lead to inconsistent expert selections. To address this, we propose selecting experts by learning a discrete representation and assigning each input to its corresponding expert. For the pretraining task, the total computational resources remain unchanged; however, our method saves 28% of computational resources when applied to downstream tasks.

W5. Besides, I would like to know that whether this technique works for unbalanced datasets. When most samples correspond to one vector in the codebook, this technique may be hard to prevent the collapse. It will be helpful to provide some explanations about it.

A5.

Yes, our method can handle unbalanced datasets. However, we need to employ some collapse-prevention techniques from the VQVAE field, such as using a frozen codebook and implicitly generating codes through a linear projection [4] or Online clustered codebooks as [5].

Reference:

[1] Gou, Y., Liu, Z., Chen, K., Hong, L., Xu, H., Li, A., Yeung, D., Kwok, J., & Zhang, Y. (2024) Mixture of cluster-conditional lora experts for vision-language instruction tuning. https://arxiv.org/abs/2312.12379

[2] Chi, Z., Dong, L., Huang, S., Dai, D., Ma, S., Patra, B., Singhal, S., Bajaj, P., Song, X., Mao, X., Huang, H., & Furu Wei, F. (2022) On the representation collapse of sparse mixture of experts https://arxiv.org/abs/2204.09179

[3] Dai, D., Dong, L., Ma, S., Zheng, B., Sui, Z., Chang, B., & Wei, F. (2022) Stablemoe: Stable routing strategy for mixture of experts. https://arxiv.org/abs/2204.08396

[4] Zhu, Y., Li, B., Xin, Y., & Xu, L. (2024) Addressing representation collapse in vector quantized models with one linear layer. https://arxiv.org/abs/2411.02038

[5] Zheng, C., & Vedaldi, A. (2023) Online clustered codebook. https://arxiv.org/abs/2307.15139

Dear Reviewer YzaE,

Thank you for your reponse. We would like to address your concen as below

C1. Regarding W2, what is the specific configuration of the experiment? Why can't I find relevant data from the manuscript, such as 1.25 for VQMoE

A1.

For the configuration, we trained Transformer-XL at the base scale for 50k steps across VQMoE, SMoE, and MoCLE to ensure a fair comparison within the limited rebuttal time. Since VQMoE learns a discrete representation where each cluster corresponds directly to an expert, its results for every $K$ are equal to $top1$: 1.25.

C2. this technique works for unbalanced datasets

A2.

Unbalanced datasets are not the focus of our research. Our paper has demonstrated efficiency in the general setting through both theoretical analysis and experiments. Addressing unbalanced datasets is planned for future research applications.

Could you kindly point out any concerns that have not been fully addressed? If you have any further questions or issues, we would be happy to discuss them.

Thank you for your review. We hope our response below fully addresses your concerns and questions. Please let us know if you have any additional questions.

W1. Very confusingly written and not a clear explanation of the method.

A1.

• We theoretically demonstrate that learning a discrete representation is a more optimal approach for expert selection compared to relying on a router. However, relying solely on training a discrete representation may lead to underfitting when dealing with large-scale data in practice.
• For the pretraining task, we propose simultaneously learning both discrete and continuous representations. The continuous representation is learned using a top-1 router, while the discrete representation is learned through vector quantization.
• For the downstream tasks, we propose fine-tuning the pretraining model using only the discrete representation, which achieves better performance and improves speed by over 28%.

W2. Not sure the scale of the data and experiments are enough to draw any conclusions.

A2.

• While we agree with your concerns, scaling up to billion-scale models is not the primary focus of our work due to resource limitations. Instead, our contribution lies in proposing a discrete representation mechanism that optimally routes tokens to experts and proves effective for downstream tasks.
• For language tasks, we progressively scale the model from a few million parameters to several hundred million parameters and consistently observe that our method outperforms existing approaches. For vision tasks, we demonstrate the effectiveness of our method across various classification datasets. We are currently conducting additional experiments on the ImageNet dataset and will share the results with you once available.

Q1. Please provide clear explanations of what is done at pre-training, finetune and inference so one can clearly see what it attempts to be done?

A3.

• As illustrated in Figure 1 and detailed in Equations 4 and 5, we propose a pretraining approach that simultaneously learns both discrete and continuous representations. The continuous representation is obtained using a Top-1 Router, while the discrete representation is learned via vector quantization.
• For the downstream tasks using the pretraining model, we propose fine-tuning the pretraining model using only the discrete representation, which achieves better performance and speed.

*Q2. What is the setup for image tasks? What was the model pretrained on?

A4.

For image tasks, we conduct experiments on training image classification models from scratch across various datasets. Our training setup is based on the repository available at: https://github.com/horrible-dong/QTClassification.

W3. XMoE shows wikitext-103 results where a dense transformer reaches a pplx (22.61) smaller than any reported by this work (23.85 for large VQMoE)

Q3. What does your dense baseline obtains in wikitext-103 and why is it so different from XMoE reported numbers?

A5.

• Are you referring to 22.61 PPL on the OpenWebText that is reported by Yange et al. 2024[1] ?. The model XMoE which we mentioned in our work is proposed by Chi et al. 2022[2].
• The study by [2] trained their model with more than twice the number of parameters and significantly more iterations compared to our approach, leading to results that differ from those presented in our paper.

[1] Yang, Y., Qi, S., Gu, W., Wang, C., Gao, C., & Xu, Z. (2024) Xmoe: Sparse models with fine-grained and adaptive expert selection. https://arxiv.org/abs/2403.18926

[2] Chi, Z., Dong, L., Huang, S., Dai, D., Ma, S., Patra, B., Singhal, S., Bajaj, P., Song, X., Mao, X., Huang, H., & Furu Wei, F. (2022) On the representation collapse of sparse mixture of experts https://arxiv.org/abs/2204.09179

We are grateful for the reviewer’s constructive suggestions and provide our responses below.

W1. Proposition 4.2: I have concerns in this statement’s mathematical rigorousness. Optimality in what metrics is discussed in the proposition? Training error or loss? What assumption is used in the proof?

A1.

• For optimality we discussed in Proposition 4.2 is an expert selection that minimizes training loss function.
• About metrics or loss functions, it can be applied for any training loss function, for instance: Cross Entropy Loss or The negative log likelihood loss.
• About proof assumption, please refer to our above general response.
• The question about optimal training selection vs optimal inference selection is out of scope of this Section 4.1 Theory and Section B proof. However, it is related to the generation theory of a mixture of experts, which was proved by [1] and [2].

W2. On the two-stage training: The necessity of the two-stage training is not experimentally confirmed.

A2.

• For the question about two-stage training required, the answer is No. Table 1 and Table 3 we report the training model from scratch and our method outperforms the existing baseline.
• Besides that, we finger out that the discrete representation that we learning from massive of data can be useful for a downstream task which was illustrated as the fine-tuning results at Table 2.

Q1. How good is the first-stage-only training (i.e., VQ loss is used as a regularizer but the conventional Top-k routers are kept and used in inference) results?

A3.

• For the pretraining task, we utilize both Top-1 Router for continuous representation learning and VQ for learning discrete presentation.
• For the finetuning task, we excluded Top-1 Router for both finetuning and inference. And we showed that discrete presentation is efficient in terms of both performance and inference speed.

Q2. Is it impossible to start the second-stage training (i.e,. train VQ routers without Top-k routers from scratch) from scratch?

A4.

The answer is yes, however, training the discrete representation only approach might only work for tiny/base scale with a medium dataset in practice. The below results illustrates that point with our experiments as a below Table results on the Enwik8 dataset.

[1] Chen, Z., Deng, Y., Wu, Y., Gu, Q., & Li. Y. (2022) Towards understanding mixture of experts in deep learning. https://arxiv.org/abs/2208.02813

[2] Zhao, J., Wang, P., & Wang. Z. (2024) Generalization error analysis for sparse mixture-of-experts: A preliminary study. https://arxiv.org/abs/2403.17404

Dear Reviewer ckSP,

Thank you so much for your reponse. We would like to adress your concern about the Proposition 4.2 as below:

1. Our idea for Proposition 4.2 is a optimal to training experts rather than selection by each training steps:

• At the first training iteration, since all experts are initialized randomly, so any pair $(t_j, E_i)$ can be assigned, where $t_j \in C_j$ is a token belong to $j_{th}$ cluster, and $E_i$ is the $i_{th}$ expert.
• From the second training iteration onward, the optimal strategy is to maintain the assignment $j_{th}$ cluster to $i_{th}$ expert. This is because the $i_{th}$ expert has already started learning knowledge specific to the $j_{th}$ cluster during the previous step, making it the most effective choice among $N$ possible experts (where $N$ is the total number of experts).

2. Additionally, we want to emphasize that the core theoretical contribution of our work is presented in Section 4.2, "VQMOE Solves Representation Collapse by Design." In this section, we demonstrate that our method surpasses other routing approaches in effectively addressing representation collapse. Consequently, our approach proves to be more efficient than alternative routing methods, as supported by empirical results as Section 5.

If you have any further questions or concerns, we would be delighted to discuss them.

Thank you so much.

We would like to express our sincere gratitude to all the reviewers for their insightful comments and constructive feedback, which have greatly helped us refine and improve our approach. In the following response, we address each reviewer’s questions and concerns individually. Here, we provide a summary of the revisions made to the manuscript and respond to the common questions raised by multiple reviewers.

Summary of Paper Revisions

• Revised Formulas 4 and 5, as well as Figure 1, to improve clarity and accuracy.
• Updated Subsection 4.1 (Theoretical Analysis) and Section B (Theoretical Proof) with enhanced mathematical notation and more detailed explanations.
• Added further explanations to Figures 4 and 5 for better interpretability.
• Corrected several writing errors for improved accuracy and coherence.

Responses to Common Questions

Q1. Is this method required additional training or two stage training ?

A1. No, our method requires only single-stage training, similar to a conventional SMoE. For the pretraining phase, we leverage both a Top-1 Router for continuous representation learning and a VQ module for learning discrete representations. This results in no additional computational overhead compared to a Top-2 Router MoE. When applying the pretrained model to downstream tasks, we propose excluding the Top-1 Router during both fine-tuning and inference, showing that discrete representations are efficient in terms of both performance and inference speed.

Q2. What is the assumption for the section 4.1 Theory Analysis and is this possible for real data?

A2. We build upon the assumptions of [1] and [2] that the data exhibits a clustering structure. Besides that, other works in MoE, such as [3] and [4], assume that the data contains specific patterns. From a theoretical perspective, [2] demonstrated that the Mixture of Experts (MoE) can be interpreted as a clustering problem. If the number of clusters is smaller than the number of experts, collapse may occur; conversely, if it is larger, the model could underfit. In practice, selecting an appropriate number of experts can be achieved through hyperparameter tuning tailored to the specific dataset.

Q3. Can this method scale up language models to billion parameters?

A3. Scaling up to billion-scale language models is not the primary focus of our work due to resource limitations. Instead, our contribution lies in proposing a discrete representation mechanism that optimally routes tokens to experts and proves effective for downstream tasks. For language tasks, we progressively scale the model from a few million parameters to several hundred million parameters and consistently observe that our method outperforms existing approaches. For vision tasks, we demonstrate the effectiveness of our method across various classification datasets. We are currently conducting additional experiments on the ImageNet dataset and will share the results with you once available.

Q4. Can this method work with a large scale vision dataset such as Imagenet?

A4. We are training the baselines and our method on ImageNet and we will report the result as soon as possible.

Reference:

[1] Chen, Z., Deng, Y., Wu, Y., Gu, Q., & Li. Y. (2022) Towards understanding mixture of experts in deep learning. https://arxiv.org/abs/2208.02813

[2] Dikkala, N., Ghosh, N., Meka, R., Panigrahy, R., Vyas, N., & Wang. X. (2023) On the benefits of learning to route in mixture-of-experts models. https://aclanthology.org/2023.emnlp-main.583.pdf

[3] Chowdhury, M., Zhang, S., Wang, M., Liu, S., & Chen. P. (2023) Patch-level routing in mixture-of-experts is provably sample-efficient for convolutional neural networks. https://arxiv.org/pdf/2306.04073

[4] Chowdhury, M., Wang, M., Maghraoui, K., Wang, N., Chen, P., &Carothers. C. (2024) A provably effective method for pruning experts in fine-tuned sparse mixture-of-experts. https://arxiv.org/abs/2405.16646