A Stronger Mixture of Low-Rank Experts for Fine-Tuning Foundation Models

We appreciate your reviews and thank you for acknowledging our efforts on theoretical and experimental analysis. For your concerns mentioned in the review, we provide corresponding response below:

Response to your concern of AdamW performances under Multi-Task scenarios

Our supplementary material does not indicate an out-performance of our method for AdamW optimizer under multi-task scenarios. We considered it may due to lacking a sufficient exploring across different multi-task scenarios, since we only conducted the multi-task experiments under a single mixture of two tasks (ScienceQA and MRPC). Consequently, we conducted more experiments on multi-task scenarios in our revision, including two mixtures of tasks and two various configurations of MoE structure.

In our revision, we grouped six tasks from the GLUE Benchmark into two mixtures. The first mixture consists of CoLA, SST-2 and MRPC tasks, which serves as a multi-task scenario involving both grammar checking, sentiment classification, and equivalent sentences judging; The second mixture consists of STS-B, QQP and QNLI tasks, which serves as another multi-task scenario involving both sentence similarity scoring, equivalent questions judging, and question-answering NLI. For evaluation, we tested candidates on each of the tasks individually and then averaged per-task performances within the mixture as the overall evaluation for that mixture.

For sufficiently assessing the multi-task performance of our proposed gate-based rescaling method, we conducted experiments under two different MoE configurations, i.e., $20/10/4$ and $10/5/4$. We trained each candidate for 2000 steps under RAdamW and gRAdamW (RAdamW with our proposed gate-based rescaling method) optimizers. The following two tables illustrate our performances under the first and the second mixtures respectively.

Mixture 1: CoLA + SST-2 + MRPC

Mixture 2: STS-B + QQP + QNLI

As a result, we concluded that overall our proposed method is still effective to boost AdamW optimization under multi-task scenarios. We have revised our supplementary material to integrate these experiments and conclusions.

Thanks for your valuable reviews and your agreements on our proposed method and our efforts on literature review. We have conducted several new experiments and provided responses to all your concerns:

Response to W1 about the limitation of LLaVA experiments

In our revision, we conducted more LLaVA experiments on both Visual7W and VMCBench using LLaVA-v1.5-7B. Specifically, we implemented candidates under two new MoE configurations with different expert numbers: $16/8/4$ and $10/5/4$. The following table illustrates the results. An overall enhancement of our method can still be witnessed under different configurations, especially for SGD.

Our performance boosts for LLaVA under AdamW might not be that remarkable. Therefore, to conduct a further significance analysis of our AdamW boosting, we also implemented more candidates and trained them by AdamW. Please refer to the following table. Similar phenomenon can be invariably witnessed.

Response to W2 about performances on different LLaMA models

We conducted more experiments on LLaMA models besides Llama-3.2-3B. Specifically, among all the LLaMA 3.2 models there are only 1B and 3B models that are purely textual. Therefore, we decided to include further experiments on Llama-3.2-1B. Four QA benchmarks have been tested, each trained for 2000 steps. Due to the limited resources and time, we set a relatively smaller MoE configuration to speed up training, which is $10/5/1$. Results are illustrated in the following table. We also included this table in the revision of our paper.

Besides LLaMA 3.2 models, we also tested Llama-3.1-8B. Until now we only conducted four ScienceQA evaluations with each training for only 600-800 steps. Results are below:

Response to W3 about the sum-to-1 constraint in MoE

In general MoE, constraining the sum of all gate values to 1 contributes to the model stability and probabilistic interpretation. Firstly, it normalizes gate outputs to get rid of some uncontrolled issues during training and inferring, such as gradient explosion or vanishing, overflow of variables from accumulated forwarding, etc. During training, by fixing the sum to 1, the gating network focuses solely on allocating relative importance among experts, rather than learning absolute weight magnitudes. This also reduces the complexity of the optimization problem and stabilizes training; Secondly, sum-to-1 constraint can be interpreted as the probabilistic distribution of selecting each expert, aligning with the requirements of probabilistic models and enabling clearer theoretical foundation of MoE.

Response to W4 and your suggestion of prominently highlighting our contributions

Thank you for your suggestions. Our contributions include integrating mixture of LoRAs with Riemannian preconditioners to alleviate both limited representation and sub-optimality issues; emphasizing the distortion issue behind per-expert preconditioning; and proposing a gate-based rescaling method and its engineering approximation to further boost MoE-LoRA training. We have already revised our abstract, introduction and conclusion sections, to explicitly involve these statements.

Response to W5 about small errors

We revised the mentioned "table 4" to "Table 4" in Section 4.5 in our revision, together with checking and fixing some other small notation errors.

Thank you for providing postive feedbacks on our presentations, derivations, and also experiments. For the weaknesses and suggestions you pointed out, we highly value them and provide responses below:

Response to W1 about further explaning the convergence and fixing the issues in Figure 2

We have added in our revision further explanations for the converging figures in Figure 2. For example, the x-axis represents training steps, the left y-axis in each figure represents the training or validation losses, while the right y-axis in each figure represents the accuracy metrics of test sets; Before implementing our gate-based rescaling method, the training and validation losses of RSGD optimizer across four tasks are significantly reduced around 100-200 steps, while after implementing our method, they are significantly reduced earlier around 0-100 steps; In addition to converging speed, we also notice an out-performance of our method in terms of converged loss and QA accuracy; Finally, to address the axis overlapping issue, we re-arranged each subplot in the figure to be less tight with each other and re-drew the Figure 2 in our revision.

Response to W2 and W3 about further elaborating Eq.13

We further elaborate why we implement Eq.13 ($X=\hat{W}+\sum_{i=1}^{N_{Expert}}\hat{\sqrt{g_i}}B_iA_i+(g_i-\hat{\sqrt{g_i}})\hat{B_i}\hat{A_i}$), which is our engineering approximation for achieving Eq.11 and Eq.12. Since by forwarding as Eq.13, the gradient updating process of $X$ can be derived as the following (Similar as Eq.9 and Eq.10, we treat gate value $g_i$s as constants when focusing on the gradients of $A_i$s and $B_i$s):

$$\begin{aligned} X_{new}=&\hat{W}+\sum_{i=1}^{N_{Expert}}[\hat{\sqrt{g_i}}(B_i-\eta\nabla_{B_i}\mathcal{L})(A_i-\eta\nabla_{A_i}\mathcal{L})+(g_i-\hat{\sqrt{g_i}})\hat{B_i}\hat{A_i}] \\\\ =& (\hat{W}+\sum_{i=1}^{N_{Expert}}\hat{\sqrt{g_i}}B_iA_i+(g_i-\hat{\sqrt{g_i}})\hat{B_i}\hat{A_i})-\eta\sum_{i=1}^{N_{Expert}}\hat{\sqrt{g_i}}[B_i(\nabla_{A_i}\mathcal{L})+(\nabla_{B_i}\mathcal{L})A_i] \\\\ =& X-\eta\sum_{i=1}^{N_{Expert}}\hat{\sqrt{g_i}}(B_i\nabla_{A_i}\mathcal{L}+\nabla_{B_i}\mathcal{L}A_i) \\\\ =& X-\eta\sum_{i=1}^{N_{Expert}}(\hat{\sqrt{g_i}})^2Proj_{col(B_i)}(\nabla_{X}\mathcal{L})^T-\eta\sum_{i=1}^{N_{Expert}}(\hat{\sqrt{g_i}})^2Proj_{row(A_i)}(\nabla_{X}\mathcal{L}) \text{ (same as the derivation of Eq.10)} \\\\ =& X-\eta\sum_{i=1}^{N_{Expert}}g_iProj_{col(B_i)}(\nabla_{X}\mathcal{L})^T-\eta\sum_{i=1}^{N_{Expert}}g_iProj_{row(A_i)}(\nabla_{X}\mathcal{L}), \end{aligned}$$

so that Eq.12 can be achieved.

The advantages of implementing Eq.13 can be elaborated from two aspects: Firstly, it achieves Eq.12 when doing gradient updating of $X$ while still keeps the original behavior of training gates, because it holds the same gate gradient, $\nabla_{g_i}X={A_i}^T{B_i}^T$, as normal forwarding; Secondly, it provides equivalent behavior and the same result of normal module forwarding $X=\hat{W}+\sum_{i=1}^{N_{Expert}}g_iB_iA_i$, and only requires a relatively low overhead.

Response to W4 about the presentation issue of MoLA experiments

We found that the candidates order presented in our baseline comparison experiments (Table 4) may not appropriate. As you suggested, we have moved the $MoLA-SGD (2,4,6,8)$ candidate to the lower part of the table, alongside with other MoLA candidates.

Response to W5 about the legend of Figure 3

As you mentioned, we notice that the legend of Figure 3 should be streamlined since it has a chance to cover part of the blue line (actually it does not cover, but that still leads to unclarity). As a result, we shortened the names of lines in the legend by deleting the "Loss" word and drew the figure again.

Response to your suggestion of discussing practical implications

Thank you for your suggestions. The practical implications of this work are mainly focused on boosting the training of MoE-LoRA, which may be applied to the fields like efficient and low-resource model training, continual or multi-task learning, stablized training and modular task adaptation, etc. For example, some current works use MoE-LoRA structure to distillate knowledge from a much larger dense model, such as Xu et al.[1] Our proposed method may enhance their distillation process. We have added these practical implications to our paper, as a new section before the conclusion section.

Response to your suggestion about open-sourcing

Yes, we will open-source our implementation on Github after this paper is published.

[1] Xu, Haiyang, et al. "Sparse Mixture of Experts Language Models Excel in Knowledge Distillation." CCF International Conference on Natural Language Processing and Chinese Computing. Singapore: Springer Nature Singapore, 2024.

Thank you for your acknowledgment on our innovations and theoretical value. We’ve checked our paper again carefully to address any issues you mentioned. Here are our responses to your valuable concerns:

Response to W1 and W2 about the notations and abbreviations issues

We have carefully reviewed our paper again to address those undeclared or unclear notations and abbreviations. For example, as you mentioned, the symbol $X$ represents the overall weight matrix after integrating pretrained weights $W$ and LoRA modules $A$s and $B$s; $Proj_V(M)$ represents a projection function which projects a given matrix $M$ onto a sub space constructed by all vectors in set $V$. And when we treat all vectors in $V$ as the rows of another matrix, such as $P$, then the projecting action can be calculated as $Proj_{row(P)}(M) = MP^T(PP^T)^{-1}P$; FFN and FFT represent different concepts. FFN is the abbreviation for Feed-Forward Network; and FFT is the abbreviation for Fully Fine-Tuning. In our revision, we have addressed all the above notation issues, as well as some other unclear notation issues we found during our re-check.

Response to W3 and Q1 about more explanations of Eq.12 on its fully-finetuning equivalency

Eq.12 is the refined Riemannian-preconditioned backpropagating equation in MoE cases, after applying our proposed gate-based rescaling method. It further approaches global fully finetuning because of two basic reasons: Firstly, it is derived by implementing Riemannian preconditioners to calibrate each LoRA expert's gradient (given by Eq.6), thus ensuring each LoRA expert can get close to their respective full-rank training behavior locally (i.e., per-expert fully-finetuning equivalency), according to Zhang et al.[1]; Secondly, we notice a further distortion of each expert space introduced by their respective gate values, leading to an inconsistency between per-expert local optimals and the global optimal. Therefore, we further introduce a respective gate value $g_i$ as a re-scaler to each expert's Riemannian preconditioner (given by Eq.11), to relieve the expert distortion resulting from the multiplication of gate value during forwarding. As a result, Eq.12 can further approach global fully-finetuning equivalency (e.g., larger gate values introduce less distortion. Thus, through Eq.12, experts with larger gate values are re-scaled less than those with smaller ones). We have already integrated above discussions about fully-finetuning equivalency in our revision.

Response to W4 about insufficient $n/k/r$ analysis

We already presented our $n/k/r$ analysis in Table 5 in Section 4.6, which consists of seven different candidates tested under SGD and AdamW optimizers. Llama-3.2-3B serves as their foundation model. Table 5 already demonstrates our overall effectiveness across various $n/k/r$ configurations. To make the investigation further sufficient, we added two new experiments under LLaVA-v1.5-7B with two different $n/k/r$ configurations (16/8/4 and 10/5/4). Please refer to the results in our rebuttal to Reviewer 1wAK. An overall enhancement of our method can still be witnessed.

We have already emphasized our preliminary conclusion in Section 4.6 that, the value of $k$ is more important to performance boosting of our method under SGD optimizers. Now we provide a further analysis about both our performance boosting and the final overall performance: Firstly, for performance boosting, we calculate its correlation with $n$, $k$ and $r$ from Table 5, which are $0.357$, $0.912$ and $0.093$ respectively under SGD optimizers, demonstrating our preliminary conclusion is correct. While under AdamW optimizers they are $0.001$, $0.197$ and $0.806$, indicating that $r$ might be more important for boosting AdamW; Secondly, for final performance, we claim that it results from both our boosting effectiveness and the MoE fundamental features (such as a too-large MoE structure may lead to underfitting, while a too-small one may lead to overfitting). After the mixed influence of both aspects, the optimal configuration for LLama-3.2-3B in our ScienceQA experiments is the $10/5/1$ with our re-scaled SGD optimization; while for LLaVA-v1.5-7B in VMCBench experiments, for example, it is $10/5/4$ with our re-scaled AdamW optimization.

Response to your suggestion of reviewing each derivation step

Thank you for your suggestion. We have reviewed all the derivation steps in our paper again to make sure each step is theoretically correct.

[1] Zhang, Fangzhao, and Mert Pilanci. "Riemannian preconditioned lora for fine-tuning foundation models." arXiv preprint arXiv:2402.02347 (2024).