MoLEx: Mixture of Layer Experts for Fine-tuning with Sparse Upcycling

Q7.The MoLEx method proposed in this paper appears to be adaptable to various models and training methods. However, it was only validated under the LoRA method. I suggest that the authors include experiments with full parameter training to provide a more comprehensive evaluation of the method's effectiveness.

Answer: Thank you for your comments. As the reviewer noted, MoLEx is adaptable to various models and training methods. Following the reviewer's suggestion, we have conducted additional experiments for RoBERTa using full parameter fine-tuning on 4 GLUE tasks. We present the results in Table 3 below. Our MoLEx consistently outperforms the full parameter fine-tuning across all tasks. These findings further confirm MoLEx's adaptability to different models and training methods. We have included these results in Appendix E.1, Table 10, in our revision. We are currently running additional results for this experiment on other GLUE tasks and will report the results in our next reply.

Table 3: RoBERTa with full parameter fine-tuning and with MoLEx when fine-tuned on 4 GLUE tasks, RTE, MRPC, STS-B and CoLA. We report accuracy for RTE and MRPC, Pearson correlation for STS-B and Matthew's correlation for CoLA. A higher value reflects a better performance of the model.

Q8. Current sMoE models applied to LLMs typically operate at the token level, where different experts are activated for each input token. However, the appendix of this paper states that MoLEx employs a sentence-level approach, with all tokens in a batch going through the same experts. It would be valuable for the authors to further explore how MoLEx performs in a token-level activation context, as this could provide insights into its adaptability and effectiveness in such scenarios.

Answer: We agree with the reviewer that investigating MoLEx's performance in a token-level setting would offer valuable insights into its adaptability and effectiveness. However, the implementation of this token-level approach can be inefficient due to the complexities involved in scattering and gathering the tokens to different experts. A highly optimized code written at the CUDA level is needed to efficiently implement the token-level approach for MoLEx. We leave this as future work.

Q9. Line 148 of the paper concludes that "preserving the useful information in the pre-trained model" is achievable, supported by extensive mathematical derivations. However, I believe this conclusion is not adequately reflected in the experiments. As noted in [1], parameters themselves can retain knowledge, so a catastrophic forgetting experiment is necessary to validate this claim. I suggest using benchmarks like TRACE [2] to assess whether MoLEx can mitigate the issue of catastrophic forgetting.

Answer: We believe there is a misunderstanding of the advantage of our MoLEx method in preserving the useful information in the pre-trained model. Please allow us to clear this misunderstanding by clarifying that our MoLEx is designed for parameter-efficient finetuning, but not continual learning. By claiming that MoLEx "preserves the useful information in the pre-trained model", we mean that MoLEx, as a new parameter-efficient finetuning method, can preserve the useful information in the pre-trained model with its richer structure for obtaining better performance on a downstream task but not to mitigate the catastrophic forgetting. However, following the reviewer's suggestion, we have conducted an additional experiment on the TRACE benchmark and include our results in Table 4, 5 and 6 below.

Table 4: Overall Performance (OP) and Backward Transfer Score (BWT) of Llama-2-7B on the continual learning TRACE benchmark when fine-tuned sequentially using LoRA and MoLEx.

Table 5: Detailed results of the performance of Llama-2-7B at each round of the continual learning TRACE benchmark when fine-tuned sequentially using LoRA. We report accuracy for all tasks except ROUGE-L for MeetingBank, SARI for 20Minuten and Edim similarity for Py150. Each column represents the task that the model is trained on and the rows are the tasks that the model is evaluated on after being trained on the task in the that round.

Table 6: Detailed results of the performance of Llama-2-7B at each round of the continual learning TRACE benchmark when fine-tuned sequentially using MoLEx. We report accuracy for all tasks except ROUGE-L for MeetingBank, SARI for 20Minuten and Edim similarity for Py150. Each column represents the task that the model is trained on and the rows are the tasks that the model is evaluated on after being trained on the task in the that round.

References

[1] HFT: Half Fine-Tuning for Large Language Models

[2] TRACE: A Comprehensive Benchmark for Continual Learning in Large Language Models

We hope we have cleared your concerns about our work. We have also revised our manuscript according to your comments, and we would appreciate it if we can get your further feedback at your earliest convenience.

Q1. The introduction of this paper does not effectively convey the motivation for the proposed MoLEx method. I would appreciate it if the authors could briefly clarify the limitations of existing sMoE methods and explain the specific problems that the proposed approach aims to address.

Answer: Thanks for your comments. Please allow us to clarify the motivation of MoLEx.

1. On one hand, MoLEx is a novel sparse mixture of experts (SMoE) whose experts are layers in a pre-trained model. In particular, given a dense pre-trained model, MoLEx combines layers in the pre-trained model for fine-tuning on downstream tasks. Compared with other SMoE methods, MoLEx does not introduce new expert parameters and is memory efficient. In order to turn a dense pre-trained model into an SMoE, other existing SMoE methods introduces new experts with their own parameters at each layer, thus increasing the total number of parameters of the model. MoLEx reuses the layers in the dense pretrained models as additional experts at each layer. During fine-tuning, the parameters of these experts are shared with those of the corresponding layers in the dense model. Compared with the dense model, MoLEx only introduces new parameters for the routers, but not the experts. As the only additional parameters introduced is a single linear layer for the routing function, the memory required to load MoLEx onto a GPU or CPU is almost the same as the dense checkpoint and less than an SMoE model.

2. On another hand, MoLEx is an innovative parameter efficient fine-tuning method. By combining pre-trained layers, MoLEx enhances the model's predictive capabilities on downstream tasks, achieving better fine-tuning performance without increasing the effective number of parameters, as shown in Table 1 and 2 for accuracy enhancement, Table 3 for zero-shot evaluation, and Table 5 for robustness improvement, in our manuscript.

Q2. The methods section of this paper relies heavily on mathematical formulas, whereas the MoLEx method itself is relatively straightforward. I suggest that the authors consider providing a method diagram to complement the mathematical expressions, as this could facilitate a better understanding of the proposed method.

Answer: Thanks for your suggestion. While the concept of MoLEx is novel, the method itself is simple to understand and easy to implement, which is indeed an advantage of MoLEx. We rigorously formulate MoLEx in Section 2.2 of the main text to provide a theoretical guarantee of its improved robustness compared to the baseline model, as demonstrated in Theorem 1 in Section 2.3. Additionally, our mathematical formulation of MoLEx allows us to analyze MoLEx's capability to retain valuable information from the pre-trained model and achieve compositional representations later in Section 2.2. As the reviewer suggests, we have included a method diagram in Section 2.2, Figure 1 to demonstrate our proposed MoLEx method.

Q3. I believe there are some unclear aspects in the implementation of the method presented in this paper. Specifically, when mixing layers, is the paper using: $x_{t+1} = x_t + \alpha u_t(x_t) + (1 - \alpha) u_l(x_t)$, where $u_l$ is the top-1 layer selected by the router, to obtain $u_{t+1}$? Or is it treating the low-rank matrices $A$ and $B$ from LoRA as experts to derive $x_{t+1}$ through $x_{t+1} = x_t + u_t(x_t) + \alpha B_t(A_t(x_t)) + (1 - \alpha) B_l(A_l(x_t))$? I also hope the authors can add relevant training details in the methods or implementation section to clarify this.

Answer: We are using $x_{t+1} = x_t + \alpha u_t(x_t) + (1 - \alpha) u_l(x_t)$, where $u_l$ is the top-1 layer selected by the router. We only employ LoRA to adapt our MoLEx model to downstream tasks. We have added relevant training details in Section 2.2 of our manuscript to clarify this and summarize the training details below.

To clarify the implementation of our method, we first insert the relevant parameter efficient fine-tuning method into the pre-trained model to obtain each layer $u_j$. Next, we initialize a trainable gate in the model, to be shared across all layers. This gate will determine the top-1 layer selected, $u_l$, to be mixed with $u_j$ at each $j=0,1,\cdots,T-1$.

Q4 .Referring to point three, I believe the authors' assessment of efficiency may be inaccurate regardless of the method used. The claim that "these layer experts can be computed in parallel and do not increase the computational time" (lines 494-495) seems misleading. Since each input must undergo forward computations through two layers in the case of top-1, the total computational load will inevitably increase, not just due to the additional computations from the gate layer. Therefore, it may not be appropriate to assert that there is no impact on efficiency solely based on parallel computation. I encourage the authors to reassess their evaluation of efficiency, particularly by considering results from single-GPU inference.

Answer: We agree with the reviewer that each input must undergo forward computations through two layers, hence increasing the total computational load. However, as computational efficiency refers to the amount of time required for a given step in a calculation, we maintain that MoLEx is as efficient as the original method used without MoLEx. This is due to the parallelization of the two forward computations which maintains the computational time while the computational load is doubled. We provide a more detailed efficiency analysis in the Table 1 below and included it in Appendix E.3 of our revised manuscript.

Table 1: Efficiency analysis of Llama-3.2-1B fine-tuned using LoRA during inference on the Alpaca dataset with and without MoLEx implemented.

As we see in Table 1, while the number of flops per sample is almost doubled, the amount of time per sample is only slightly increased. Hence, MoLEx is as efficient as the naive parameter efficient fine-tuning model.

Regarding single-GPU inference, it is also possible to compute two forward computations in parallel on a single-GPU due to the Pytorch asynchronous nature. We refer to a reference in pytorch.

Q5. The descriptions of the trainable parameters for LoRA and MoLEx in Tables 1 and 2 may be misleading. It could lead readers to believe that MoLEx has no additional trainable parameters compared to LoRA. However, the inclusion of the gate layer, regardless of its parameter count, should be represented to accurately reflect the differences. I encourage the authors to provide a clear comparison of the trainable parameters in the tables.

Answer: We only use 1 gate network, a.k.a, router, for all layers in MoLEx, i.e., layers in MoLEx share the same gate network. The number of parameters of this gate network is 9,228, which is neglegible when compared with the total number of the trainable parameters in LoRA, which is 0.3M. We have updated Table 1 and 2 in our revision to include these additional parameter counts.

Q6. The authors utilized eight A100 GPUs for their experiments, but models like RoBERTa and GPT-2 seem somewhat outdated at this point. I suggest that the authors consider expanding their experiments to include more recent LLMs, such as Qwen-2.5 or Llama-3.1, to enhance the credibility of their findings. Additionally, evaluating these models using popular current assessment methods like MMLU, GSM8K, and HumanEval would provide valuable insights.

Answer: Thank you for your suggestion. As you suggested, we have conducted additional experiments to fine-tune Llama-3.2-1B on the Alpaca dataset using LoRA. We use the publicly available repository https://github.com/meta-llama/llama3 for our experiments and employ MoLEx to fine-tune the model in comparison with LoRA. As shown in Table 2 below, on this task with Llama-3.2-1B, MoLEx achieves better train and validation PPL than LoRA. These results further justify the effectiveness of our method. We have also included these results in Appendix E.2, Table 11, in our revision.

Table 2: Train and validation perplexity (PPL) when fine-tuning Llama-3.2-1B on Alpaca using LoRA and LoRA + MoLEx. Lower PPL is indicative of better performance.

We would like to thank the reviewer again for your thoughtful reviews and valuable feedback.

We would appreciate it if you could let us know if our responses have addressed your concerns and whether you still have any other questions about our rebuttal and answers to your previous questions.

We would be happy to do any follow-up discussion or address any additional comments.

Q1. My major concern of this paper is the linguistic analysis part. One of the main claim of this paper is "relevant linguistic information is captured by selected experts". I think it's very intuitive and hard to explain. In Figure1, the author tried to connect some linguistic feature with different layers in a probing task. But I am not quite convinced by the analysis, I think most NLP task involves different level of linguistic information, so it's hard to say one task focus more on certain part of information and the proposed method could facilitate the process. So would like to see more discussion on this direction or if there is another angle to explain where the performance gain is coming from.

Answer: Thank you for your question. We agree with the reviewer that the claim of "relevant linguistic information is captured by selected experts" is intuitive but might be difficult to understand. We will provide a more detailed explanation of how the claim can be empirically validated using the probing tasks as in seen Table 4 of our manuscript and its connection with NLP tasks.

Each probing task is a classification task consisting of sentences and labels that represent a certain type of linguistic information. For example, BShift is a binary classification task where half the dataset has sentences intact and another half has sentences with 2 random adjacent words inverted. The goal is to predict if the sentence has a legal word order or if it has been inverted, which is a form of syntatic information. We input each sentence in the dataset into a language model and extract out the hidden representations, the outputs, of the sentence in each layer. Then, a classifier is trained on these representations and their labels. A high classifier score indicates that the representations at that layer effectively encodes the type of linguistic information that the dataset represents. This is a standard analysis performed to understand the linguistic properties in a model as proposed in [1].

This is the primary reason for bolding the top 2 results of each probe task in Table 4 of our manuscript, to highlight the type of linguistic information captured by each layer of the model. From our results, in a RoBERTa model, the early layers contain superficial information, the middle layers, syntactic information and later layers, semantic information.

Now that we have established the linguistic feature in each layer, we can form the connection between that and our proposed method. We agree with the reviewer that each NLP task should involve different levels of linguistic information. That is the reason we believe our method works well, as a broader variety of information is provided to the model at each layer by combining layers as experts. At each layer in MoLEx, we combine different experts, each capturing distinct linguistic information, to provide the model with a more comprehensive understanding of the input data to make better predictions. Hence, generally, we see in Figure 1 of the paper, most of the chosen experts include layers that represent different types of linguistic information i.e. on the off-diagonal blocks.

However, there are certain tasks that require only specific types of information about the sentence to complete the task. These are the tasks that should prioritize choosing layers that encode for the information needed. CoLA is one example as it is a classification task on grammatical acceptibility. To solve the task, the model should mainly require information about grammer like tenses and noun or verb placements. These are captured by the later layers and as seen in the heatmap for CoLA in Figure 4 (left subplot), these are the layer blocks chosen the most frequently (last column of the grid).

References

[1] Conneau, Alexis, German Kruszewski, Guillaume Lample, Loïc Barrault, and Marco Baroni. "What you can cram into a single vector: Probing sentence embeddings for linguistic properties." arXiv preprint arXiv:1805.01070 (2018).

Q2. I feel the real computational cost of the proposed model is larger than Lora, but the parallelism makes the final speed the same. Is there any statics on this issue?

Answer: We agree with the reviewer that real computational cost of the proposed model is larger than the naive PEFT method, but the parallelism makes the final speed the same.

While MoLEx almost doubles the overall computational load (flops), there is only a minimal increase in inference time due to parallelization of the forward computations through two layers. We provide a more detailed efficiency analysis in the Table 1 below for a more accurate comparison with baseline models and have included this in Appendix E.3 of our revised manuscript.

Table 1: Efficiency analysis of Llama-3.2-1B fine-tuned using LoRA during inference on the Alpaca dataset with and without MoLEx implemented.

As we see in Table 1, while the number of flops per sample is almost doubled, the amount of time per sample is only slightly increased, maintaining the efficiency of the model in terms of run time per computational load.

Q3. One small issue is the presentation of Lemma and Theorem looks a little messy, the alignment and layout could be adjusted to make it easy to read.

Answer: Thank you for your comments. We have centered the assumptions in Theorem 1 and prevent line breaks on equations in the revised manuscript. We would greatly appreciate any additional suggestions you could provide on improving the alignment and layout of the Lemma and Theorem in our paper to enhance readability. We will incorporate your recommendations to refine the presentation of the Lemma and Theorem in our final revision.

We hope we have cleared your concerns about our work. We have also revised our manuscript according to your comments, and we would appreciate it if we can get your further feedback at your earliest convenience.

Q1. I am a bit confused by the inference time. I can understand MoE inference can be accelerated with parallel/distributed computation. But it also means using more resources. It is better to clarify what is the fair comparison in terms of the efficiency with naive PEFT models.

Answer: We agree with the reviewer that while more resources are required during MoE inference, these can be accelerated with parallel/distributed computation.

As computational efficiency refers to the amount of time required for a given step in a calculation, we maintain that MoLEx is as efficient as the original method used without MoLEx. While MoLEx almost doubles the overall computational load (flops), there is only a minimal increase in inference time due to parallelization of the forward computations through two layers. We provide a more detailed efficiency analysis in the Table 1 below for a more accurate comparison with naive PEFT models. We also include our detailed explanation in Appendix E.3 of our revised manuscript.

Table 1: Efficiency analysis of Llama-3.2-1B fine-tuned using LoRA during inference on the Alpaca dataset with and without MoLEx implemented.

As we see in Table 1, while the number of flops per sample is almost doubled, the amount of time per sample is only slightly increased, maintaining the efficiency of the model in terms of run time per computational load.

Q2. How to determine K in top-K and the number of experts in the upcycling process. I know Top-1 is widely used number. But I am not sure in this paper, how will it impact the performance.

Answer: Thanks for your comments. As mentioned in Section 4.3 of our manuscript, in Table 9 in Appendix D, we evaluate 3 GLUE tasks—CoLA, QQP, and SST-2—using both Top-1 and Top-2 routing. Our results indicate that Top-1 routing consistently performs better. Therefore, we adopt Top-1 routing for MoLEx. For convenience, we display our results below as well.

Table 9: Comparison of RoBERTa-base on GLUE tasks, CoLA, QQP and SST-2 when fine-tuned with MoLEx using Top-1 and Top-2 routing. We report accuracy for all tasks in the table below.

We hope we have cleared your concerns about our work. We have also revised our manuscript according to your comments, and we would appreciate it if we can get your further feedback at your earliest convenience.

We would like to thank the reviewer again for your thoughtful reviews and valuable feedback.

We would appreciate it if you could let us know if our responses have addressed your concerns and whether you still have any other questions about our rebuttal.

We would be happy to do any follow-up discussion or address any additional comments.