HMoRA: Making LLMs More Effective with Hierarchical Mixture of LoRA Experts

In our paper, $T^{\text{in}}$ and $T^{\text{tg}}$ refer to the input sequence and the output sequence, respectively, as introduced in the Preliminary section. This corresponds to the practical use of LLMs, where $T^{\text{in}}$ represents the input question given to the LLM, and $T^{\text{tg}}$ represents the expected response. This is a common practice in instruction fine-tuning for LLMs. During training, we concatenate $T^{\text{in}}$ and $T^{\text{tg}}$ and feed them together into the LLM, but we only compute the loss and perform backpropagation on the tokens in $T^{\text{tg}}$. Since $T^{\text{tg}}$ is unknown during the LLM's actual deployment, it does not participate in the task routing calculations. Task information is derived solely from $T^{\text{in}}$.

Thank you for your valuable feedback and thoughtful questions. We truly appreciate the time and effort you have invested in reviewing our work. Below, we provide detailed responses to your comments, addressing each point.

• The key innovation in our approach is the discovery of a hierarchical mechanism, where using token-level routing in the shallow layers and gradually transitioning to task-level routing in the deeper layers proves to be more effective. Our hierarchical mechanism allows the model to capture information at different granularities more effectively. The experiments in Appendix E.5 also validate the effectiveness of our hierarchical hybrid routing.
• Although GJS may have applications in other fields, we are the first to apply it to MoE routing. GJS perfectly addresses the issues of imbalance and lack of certainty in routing. Our study provides significant insights and a valuable reference for future MoE research. Furthermore, we propose a new Constrained GJS divergence (CGJS) to enhance its effectiveness as a loss function. Directly applying the GJS divergence did not yield satisfactory results, as it severely impacted the router's generalization ability and failed to achieve balanced routing. The experiments in Appendix E.4 further validate the necessity and effectiveness of this adjustment. When $\gamma_c$ is set to 0.4, the model achieves optimal performance. In contrast, when $\gamma_c$ is set to 0, equivalent to optimizing the original GJS, the routing imbalance is as severe as when no auxiliary function is applied.
• Additionally, in Appendix D, we provide a unique explanation of how the auxiliary function operates similarly to a clustering method, which clarifies why our auxiliary function enhances the task router's routing capabilities. The auxiliary function brings similar tasks closer in the latent space and pushes dissimilar tasks further apart, thereby achieving task differentiation. This clustering-like approach ensures that the task router is capable of generalizing to unseen tasks.

• Ablation experiments in Section 4.3 demonstrate that the task router can differentiate unseen tasks in MMLU when using the auxiliary function, whereas without it, the router is unable to make such distinctions. This indicates that the auxiliary loss enhances the task router's ability to differentiate between tasks, even for those not encountered during training.

• In Appendix E.5, we conducted ablation studies on the balance of token routing and task routing. The results show that hierarchically combining task-level routing with token-level routing outperforms using either one alone. A constant $\alpha$ indicates that the proportion of task-level routing remains fixed. In contrast, a hierarchical $\alpha$ refers to $\alpha$ either increasing or decreasing across layers. Experimental results show that an increasing $\alpha$ achieves the best performance.The experimental results are as follows:

• First, since we do not use task labels during training but instead extract semantic information directly from the input for task-level routing, we are not constrained by a predefined set of tasks in the training data, unlike some previous supervised task-level routing methods.
• Second, the auxiliary function serves a clustering-like role, bringing semantically similar tasks closer in the latent space and pushing dissimilar tasks further apart. This enhances the task router's ability to distinguish between tasks in an unsupervised manner. Similar to clustering methods, when faced with a new task, the router can direct it to the nearest cluster based on its semantics. A detailed explanation of this is provided in Appendix D.

Thank you very much for the response and revision, which addresses some of my concerns. I will maintain my current rating and advocate acceptance of this paper.

Thank you very much for your thoughtful feedback. We truly appreciate the time and effort you have invested in reviewing our paper. Below, we address each of your concerns in detail and hope to clarify the points you raised:

The main body of the paper focuses on the core structure of the method, which we consider indispensable. Due to the 10-page limit, we had to place many key details in the appendix, despite their importance for understanding certain aspects of our approach. In the upcoming revised version, we will provide a more detailed description of these aspects in the main text to ensure that readers gain a better understanding of our approach.

LoRA+MoE serves as the underlying framework to validate the effectiveness of our hierarchical hybrid routing and auxiliary loss. Standard MoE models generally require training LLMs from scratch, and verifying our method on such models would incur significant computational costs. In contrast, the LoRA+MoE architecture is commonly used for fine-tuning, which is far more cost-efficient compared to pretraining. Therefore, we initially validated our approach using the LoRA+MoE architecture. However, it is important to emphasize that LoRA+MoE is not the focus of our work. We will clarify this relationship in the revised version of the paper.

The key contribution of our work is not merely the combination of MoE and LoRA, but rather the following highlights:

• The key innovation in our approach is the discovery of a hierarchical mechanism, where using token-level routing in the shallow layers and gradually transitioning to task-level routing in the deeper layers proves to be more effective. Althogh PanGu-Σ [1] also combines task-level and token-level routing, first selecting a group of experts based on domain labels, and then randomly routing tokens to one of the experts in that group. This method is limited to predefined domains and cannot generalize to new domains. And it is inefficient because the shallow layers of an LLM primarily capture token-level information, making task-level routing unnecessary at that stage. Our hierarchical mechanism allows the model to capture information at different granularities more effectively. The experiments in Appendix E.5 demonstrate that our hierarchical hybrid routing is more effective than using token-level and task-level routing at every layer. Additionally, J. Zhu et al.[2] proposed various granularity routing methods but did not combine these different granular approaches.
• The auxiliary function is another key contribution of our work. Our auxiliary function enhances the certainty of routing decisions, resulting in more stable routing and promoting expert specialization. While StableMoE [3] achieves stable routing (in terms of balance and certainty) through a two-stage training and distillation process, we are able to achieve similar stability by simply optimizing our auxiliary function.
• Additionally, the combination of the auxiliary function and the task router improves the router's ability to differentiate between tasks. While S. Kudugunta et al.[4] uses sentence representations for routing, thereby eliminating the dependency on task lables, our experiments reveal that directly using sentence representations for routing struggles to effectively differentiate unseen tasks. Our auxiliary function operates in a clustering-like manner, pulling similar tasks closer together in the latent space while pushing dissimilar tasks further apart. This clustering-like approach ensures that the task router’s ability to distinguish between tasks generalizes to unseen tasks as well.

A comparison of our method with other representative MoE approaches is shown in the table below.

We will include a discussion of related work and highlight the strengths of our approach in the revised version of the paper.

[1] X. Ren et al., PanGu-Σ: Towards Trillion Parameter Language Model with Sparse Heterogeneous Computing, arXiv: arXiv:2303.10845, 2023.

[2] J. Zhu et al., Uni-Perceiver-MoE: Learning Sparse Generalist Models with Conditional MoEs, NeurIPS. 2022.

[3] D. Dai et al., StableMoE: Stable Routing Strategy for Mixture of Experts, ACL, 2022. 

[4] S. Kudugunta et al., Beyond Distillation: Task-level Mixture-of-Experts for Efficient Inference, EMNLP 2021

Additional Computational Resources for HMoRA

Thank you for your fairness. We have identified that the task encoder is indeed the primary contributor to the parameter overhead. Optimizing its efficiency will be a key focus of our future work.

Ablation Studies of Hyperparameters

We aim to ensure balanced routing, so we fixed $\gamma_b$ at 1. In Appendix E.4, we conducted experiments on $\gamma_c$, selecting values from {0, 0.2, 0.4, 0.6, 0.8}. Lower values of $\gamma_c$ result in higher certainty, but this makes achieving balance more difficult, leading to performance degradation. Our experiments showed that a $\gamma_c$ value around 0.4 yields the best results. To further investigate the effect of this hyperparameter, we visualized the distribution of the routing results in the Appendix E.4 and compared it with other auxiliary functions. The experimental results are as follows:

Thank you for your valuable feedback and thoughtful questions. We truly appreciate the time and effort you have invested in reviewing our work. Below, we provide detailed responses to your comments, addressing each point. Since baselines compare include full fine-tuning, larger models require significantly more resources for fine-tuning. Currently, we lack the resources to conduct experiments on larger models, such as 7B. However, HMoRA is compatible with other architectural models, and we conducted baseline comparison experiments on LLaMA 3.2 1B, and the results are as follows:

The experimental results on LLaMA differ somewhat from those on Qwen. Compared to Qwen, LLaMA requires more fine-tuning steps for the average accuracy across multiple benchmarks to gradually converge. With full fine-tuning, an increased number of fine-tuning steps may impair the knowledge learned during pre-training. Additionally, since our fine-tuning data does not include training data specific to the benchmarks, full fine-tuning performs worse than other PEFT methods. Full fine-tuning shows almost no improvement in accuracy on knowledge-intensive benchmarks like MMLU and MMLU-Pro. It only achieves improvements in simpler reasoning tasks such as ARC-E, OpenBookQA, and CommonsenseQA. Furthermore, MoLoRA fine-tuning proved completely ineffective, which we suspect is due to the fact that MoLoRA does not fine-tune the attention layers. Other fine-tuning methods, which freeze the parameters of the initial LLM, preserve the pre-trained knowledge and outperform full fine-tuning. Surprisingly, HMoRA+lw achieved the best average accuracy and performed best on 4 out of 7 benchmarks, even surpassing HMoRA, which has a larger parameter count.

Thank you for your valuable feedback and thoughtful questions. We truly appreciate the time and effort you have invested in reviewing our work. Below, we provide detailed responses to your comments, addressing each point.

It is important to clarify that generalization to unseen tasks refers to the task router's ability to distinguish between different tasks, rather than the inference capabilities of the LLM itself. LLMs inherently possess strong generalization abilities, and our task routing mechanism is designed to enhance the LLM's performance in multi-task scenarios. Therefore, our primary focus is on whether the task router itself can effectively differentiate tasks that were not encountered during training.

In the ablation experiments of the paper, we conducted tests on 6 unseen tasks from MMLU and used visualizations to verify the task router's ability to differentiate between unseen tasks. Furthermore, tests on multiple benchmarks also demonstrate that using the auxiliary function enhances the performance of the LLM.

Here, we delve deeper into this capability through quantitative analysis. We sampled 100 examples from each of all tasks (57) in MMLU and analyzed the task router's routing results for these samples. First, we recorded the number of expert activations, as shown in the table below:

It is important to note that, since we use top-2 routing, each sample activates two experts (an expert pair). From the results, we can observe that without the auxiliary function, all tasks are routed to experts 1 and 2.

Further, we analyzed all expert pairs activated for each task. For each task, we identified the most frequently activated expert pair from the 100 samples as the main expert pair for that task. We then examined the distribution of main expert pairs across the 57 tasks, calculating the proportion of samples for each task in which the main expert pair was activated, and whether this proportion exceeded a certain threshold. The results are presented in the table below:

We define tasks with a main expert pair proportion exceeding 0.8 as recognizable by the router, as this indicates consistent and reliable routing. Although this proportion is consistently 100% without any auxiliary function, it simply routes all tasks to a pair of experts (1,2), which we do not consider indicative of its ability to differentiate between tasks.

We further analyzed the main expert pairs for all tasks, and the results with the auxiliary function are as follows:

As shown in the table, the task router grouped all MMLU tasks into three clusters, with 42 tasks (73.68%) being effectively recognized.

The results with the load balancing loss are as follows:

Although the number of primary expert pairs increased, only two clusters, (1,3) and (3,5), were useful, and only 7 tasks (12.28%) were effectively recognized overall.

The above quantitative analysis demonstrates that our auxiliary function enhances the task router's ability to differentiate between tasks, even for those that were not encountered during training.

I'm extremely appreciative of the response and the revisions carried out. Given that my assessment score is positive, I shall accordingly retain my current rating and advocate for the acceptance of this paper.