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

Authors: Rachel Teo, Tan Minh Nguyen
Primary area: foundation or frontier models, including LLMs
Keywords: Parameter efficient fine-tuning, mixture of experts, sparse upcycling

Abstract:
Large-scale pre-training of deep models, followed by fine-tuning them to adapt to downstream tasks, has become the cornerstone of natural language processing (NLP). The prevalence of vast corpses of data coupled with computational resources has led to large models with a considerable number of parameters. While the massive size of these models has led to remarkable success in many NLP tasks, a detriment is the expense required to retrain all the base model's parameters for the adaptation to each task or domain. Parameter Efficient Fine-Tuning (PEFT) provides a highly effective solution for this challenge by minimizing the number of parameters required to be trained in adjusting to the new task while maintaining the quality of the model. While existing methods have achieved impressive results, they mainly focus on adapting a subset of parameters using adapters, weight reparameterization, and prompt engineering. In this paper, we study layers as extractors of different types of linguistic information that are valuable when used in conjunction with each other. We then propose the Mixture of Layer Experts (MoLEx), a novel Sparse Mixture of Experts (SMoE) whose experts are layers in the pre-trained model. In particular, MoLEx is applied at each layer of the pre-trained model. It performs a conditional computation of a mixture of layers during fine-tuning to provide the model with more structural knowledge about the data. By providing an avenue for information exchange between layers, MoLEx enables the model to make a more well-informed prediction for the downstream task, leading to better fine-tuning results with the same number of effective parameters. As experts can be processed in parallel, MoLEx introduces minimal additional computational overhead. We empirically corroborate the advantages of MoLEx when combined with popular PEFT baseline methods on a variety of downstream fine-tuning tasks, including the popular GLUE benchmark for natural language understanding (NLU) as well as the natural language generation (NLG) End-to-End Challenge (E2E).

Decision: accept (poster)
Reviewer scores: 8, 6, 5

Reviews and discussion:
<thread>
creation date:
1729848260248

thread_type:
official_review

<thread_object>
creator: reviewer_uihh

<content>
summary:
This paper presents MoLEx, a simple and effective method that does not alter the model's structure or add extra parameters. Instead, it treats all layers as experts and introduces a router at each layer to select the top k layers for activation. Extensive experiments on RoBERTa and GPT-2 demonstrate the strong performance of MoLEx.

strengths:
This paper introduces a novel sMoE model whose experts are layers in the pre-trained model. MoLEx is a simple yet effective approach without additional expert parameters.

weaknesses:
Weaknesses & Questions

1. 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.

2. 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.

3. 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.

4. 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.

5. 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.

6. 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.

7. 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.

8. 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.

9. 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.

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

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

questions:
See Weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730516393409

thread_type:
official_review

<thread_object>
creator: reviewer_wweu

<content>
summary:
This paper introduced a Mixture of Layer Experts (MoLEx), which leverages layers as experts to facilitate the exchange of linguistic information and improve a model’s fine-tuning and transfer knowledge ability. 
In particular, the author proposed a structural change to the architecture of the model which is compatible with PEFT method while maintaining the same number of effective parameters. 
Experimental results demonstrated that MoLEx significantly improves performance across a range of downstream tasks, including the GLUE benchmark and the E2E Challenge, while incurring minimal additional computational overhead and scales well with model size.

strengths:
I think this paper is a good upgrade in the field of Parameter Efficient Fine-Tuning.
Comparing with predominant method (Lora),  it achieve significant improvement on a few data-set with almost the same cost (speed, memory, parameter size).
Also I like the theoretical proof to justify the robustness of MoLEx in a simplified model and the empirical evidence provided.

weaknesses:
1.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. 

2.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?

3.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.

questions:
please see weakness part
</content>
</thread_object>

</thread>
<thread>
creation date:
1730591915984

thread_type:
official_review

<thread_object>
creator: reviewer_d753

<content>
summary:
This paper introduces a new approach to fine-tune LLMs. PEFT is becoming a mainstream method to adapt general-purpose model to specific tasks.MoLEx, proposed by the authors, aims at improving PEFT through leveraging pre-trained model layers and MoE to provide the model with more structural knowledge about the data. This method enhances information exchange between layers, leading to better predictions with minimal computational overhead. Experiments show that MoLEx is a more efficient technique to fine tune models.

strengths:
1. The paper has clearly explained the benefits of using MoLEx to sparse upcycling models and its robustness as an ensemble model. Theoretical proof has laid good foundation for the idea.
2. The idea of exploiting MoE to improve the efficiency of LLM fine-tuning is a novel idea and a problem worth to research.
3. The experiments are well designed and show comprehensive results on the proposed idea. Although I do think it is necessary to further expand to more models and tasks, the promising results are a good start.

weaknesses:
1. 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.

questions:
1. 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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731983643794

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Global Reply

comment:
Dear AC and reviewers,

Thanks for your thoughtful reviews and valuable comments, which have helped us improve the paper significantly. We are encouraged by the endorsements that: 1) (Reviewers D753, WWeU, UiHH) MoLEx is a novel, simple and effective method to upgrade paramter efficient fine-tuning methods, 2) (Reviewers D753, WWeU) the theoretical proof in a simplified model has laid good foundation for the idea and to justify its robustness, and 3) (Reviewers D753, WWeU, UiHH) comprehensive experiments have demonstrated that MoLEx has significant performance gains with almost the same computational speed and memory. 

One of the main concerns from the reviewers is to provide a more detailed and fair comparison with naive PEFT models as MoLEx does increase computational cost but maintains its computational efficiency through parallelism. Another concern is to expand our experiments to include more recent LLMs, such as Qwen-2.5 or Llama-3.1, and other fine-tuning methods, to enhance the credibility of our findings. We address these questions here and have updated our submission based on the reviewers' feedback. Our revisions are highlighted in blue. 

**Q1: Detailed efficiency analysis of MoLEx and naive PEFT models**

**Answer:** We agree with the reviewers 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. 

| Model \ Metric | Total Parameters | Trainable Parameters| Trainable Parameters (%) |Memory (MB)| Flop/Sample  | Sec/Sample | Flop/Sec |
| -------- | -------- | -------- |-------- |-------- |-------- |-------- |-------- |
| Llama-3.2-1B (Baseline) | 1,236,666,368  | 851,968  |0.0689  |10,442     | 12.329 T | 0.506 | 24.366 T |
| Llama-3.2-1B (MoLEx)  | 1,236,699,152  | 884,752     |0.0715     |10,442     |22.511 T | 0.557 | 40.415 T |

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: Experiments on more recent LLMs (Llama-3.2-1B) and other fine-tuning methods (full parameter fine-tuning)**

**Answer:** 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 and demonstrate our compatibility with recent LLMs. 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. 

| Model \ Metric| Train PPL ($\downarrow$) | Validation PPL ($\downarrow$) | 
| -------- | -------- | -------- |
| LoRA  | 4.18    | 4.11     | 
| MoLEx  | **4.05**    | **4.02**   |

Further, as Reviewer UiHH 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. 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 also included these results in Appendix E.1, Table 10, in our revision. 

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. 

| Model \ Task| RTE | MRPC | STS-B |CoLA |
| -------- | -------- | -------- |--|--|
| RoBERTa (full parameter)  | 80.1    | 88.7     | 90.9 |62.6|
| RoBERTa (MoLEx)  | **80.9**    | **89.5**   | **91.1**|**63.3**|


-----

We are glad to answer any further questions you have on our submission.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731983697647

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Summary of revision:

comment:
Incorporating the insightful comments from the reviewers into our paper has helped us improve its quality. Here, we provide a summary of the revisions (highlighted in blue in the revised manuscript) for convenience:

1. Section 2.2: We clarify the implementation of MoLEx with parameter-efficient fine-tuning methods to avoid confusion. A diagram is also included in Figure 1 for easy visualization of our method. 
2. Section 2.3: We fix some alignment issues and equation line breaks in  Theorem 1, Lemma 1 and Corollary 1 and 2 for an easier read and neater presentation. 
3. Section 3.1, Table 1; Section 3.2, Table 2: We make the total number of trainable parameters in MoLEx more exact to reflect the slight increase in the parameter count due to an additional implementation of a gating function. 
4. Appendix E.1: We provide additional experimental results on full parameter fine-tuning for RoBERTa on GLUE to demonstrate that MoLEx is adaptable to various training methods, besides LoRA. 
5. Appendix E.2: We provide additional experimental results on fine-tuning Llama-3.2-1B on the Alpaca dataset to illustrate that MoLEx is compatible with recent large language models. We also evaluate the baseline naive LoRA model and MoLEx on the standard benchmark MMLU for further comparison. 
6. Appendix E.3: We include a more detailed efficiency analysis of MoLEx with naive parameter-efficient fine-tuning models for a fair comparison between them.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732286260160

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Additional Experimental Results

comment:
Dear AC and reviewers,

We would like to thank all reviewers again for their thoughtful reviews and valuable feedback. We have obtained additional results on GLUE and MMLU and added more statistics for our efficiency analysis. We summarize these results below:

(1) As the reviewer UiHH noted, MoLEx is adaptable to various models and training methods. Following the reviewer's suggestion, we conducted  experiments for RoBERTa using full parameter fine-tuning on 4 GLUE tasks, RTE, MRPC, STS-B and CoLA. 

We have finished running additional results for this experiment on the remaining GLUE tasks and report our full results below in Table 1. We include results on the previous tasks reported as well for convenience. The same results can be found in Table 10 in Appendix E.3 of our revised manuscript. Our MoLEx consistently surpasses full parameter fine-tuning on each task, highlighting its versatility across various models and training methods. 

Table 1: RoBERTa with full parameter fine-tuning and with MoLEx when fine-tuned on the GLUE benchmark. We report accuracy for all tasks except for, Pearson correlation for STS-B, Matthew's correlation for CoLA and the overall (matched and mismatched) accuracy for MNLI. A higher value reflects a better performance of the model. 
 | Model \ Task| RTE | MRPC | STS-B |CoLA |MNLI | QNLI | SST-2 |QQP |Ave.|
| -------- | -------- | -------- |--|--|--|--|--|--|--|
| RoBERTa (full parameter)  | 80.1    | 88.7     | 90.9 |62.6| 87.7|92.9 | 95.0|91.8 |86.2|
| RoBERTa (MoLEx)  | **80.9**    | **89.5**   | **91.1**|**63.3**|**87.8** | **93.1**|**95.1** | 91.8| **86.6**|
    
(2) From reviewer UiHH's feedback, PPL is generally not used to evaluate a model's performance during the fine-tuning stage. Instead, benchmarks like MMLU are typically used for evaluation. Following the reviewer's suggestion, we have also run the evaluation on MMLU using Llama-3.2-1B. We report our results in Table 2 below. We improved by about 1% on MMLU, further verifying the effectiveness of MoLEx. We also include these results in Table 12 in Appendix E.2 of our revision. 

Table 2: Accuracy when evaluating Llama-3.2-1B on MMLU using LoRA and LoRA + MoLEx. A higher value is indicative of better performance. 
| Model \ Metric| Acc. % ($\uparrow$) |
 | -------- | -------- |
| LoRA  | 30.42 |
| MoLEx  | **31.51**    |

(3) We add one more statistic, mins/epoch of training, to our detailed efficiency analysis on MoLEx and naive PEFT methods. Fine-tuning on the Alpaca dataset with Llama-3.2-1B using LoRA takes 4:59 minutes per epoch, whereas using MoLEx requires 6:00 minutes per epoch. We summarize these results in Table 3 below. Same results can be found in Table 13 in Appendix E.3 of our revised manuscript.

Table 3: Time in minutes when training Llama-3.2-1B on the Alpaca dataset using LoRA and LoRA + MoLEx for 1 epoch. 
| Model \ Metric| min/epoch |
| -------- | -------- |
| LoRA  | 4:59 |
| MoLEx  | 6:00   |

----
We would be happy to do any follow-up discussion or address any additional comments.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper introduces MoLEx (Mixture of Layer Experts) for parameter-efficient fine-tuning of LLMs. The key idea is to treat existing layers in pre-trained models as experts and employ a routing mechanism to selectively combine them during fine-tuning. The authors demonstrate improved performance across multiple benchmarks while maintaining efficiency through parallel computation.

The idea and the proposed approach is novel through repurposing the layers as experts (without introducing additional parameters). The empirical experiments are also quite comprehensive to support the major claim. Reviewers also raised some concerns regarding the computational efficiency claims, as well as limited exploration of token-level routing approaches, which authors need to consider later.

While there were some concerns regarding the presentation and computation overhead, the core contribution - repurposing existing layers as experts for improved fine-tuning - is novel and interesting, and might inspire future studies. Therefore, I'm inclined to accept this paper and recommend the authors carefully consider those suggestions to revise the paper's presentation logic for the final version.