Organizations: University of Illinois Urbana-Champaign, IL, USA · Meta AI Menlo Park, CA, USA · Washington University St. Louis, WA, USA · Meta AI Sunnyvale, CA, USA
Low-rank adapters (LoRAs) are a parameter-efficient finetuning technique that injects trainable low-rank matrices into pretrained models to adapt them to new tasks. Mixture-of-LoRAs models expand neural networks efficiently by routing each layer input to a small subset of specialized LoRAs of the layer. Existing Mixture-of-LoRAs routers assign a learned routing weight to each LoRA to enable end-to-end training of the router. Despite their empirical promise, we observe that the routing weights are typically extremely imbalanced across LoRAs in practice, where only one or two LoRAs often dominate the routing weights. This essentially limits the number of effective LoRAs and thus severely hinders the expressive power of existing Mixture-of-LoRAs models. In this work, we attribute this weakness to the nature of learnable routing weights and rethink the fundamental design of the router. To address this critical issue, we propose a new router designed that we call Reinforcement Routing for Mixture-of-LoRAs (ReMix). Our key idea is using non-learnable routing weights to ensure all active LoRAs to be equally effective, with no LoRA dominating the routing weights. However, our routers cannot be trained directly via gradient descent due to our non-learnable routing weights. Hence, we further propose an unbiased gradient estimator for the router by employing the reinforce leave-one-out (RLOO) technique, where we regard the supervision loss as the reward and the router as the policy in reinforcement learning. Our gradient estimator also enables to scale up training compute to boost the predictive performance of our ReMix. Extensive experiments demonstrate that our proposed ReMix significantly outperform state-of-the-art parameter-efficient finetuning methods under a comparable number of activated parameters.
Figures & tables
Figure 1: Finetuning procedure of our proposed ReMix.
Figure 2: We often observe that only one LoRA has a dominating routing weight that is close to one at deeper layers even when we activate k>1 LoRAs during finetuning. Consequently, the computation of the other k−1 LoRAs are essentially wasted because using k>1 would have similar accuracy to k=1 . (The dominating LoRA can differ on different inputs .)
Figure 3: Routing weight collapse even worsens as finetuning progresses . The effective support size ESS(π(l)) often drops to 1 quickly during finetuning.
Type
Method
GSM8K
HumanEval
ARC-c
Average
Accuracy
Params
Pass@1
Params
Accuracy
Params
Accuracy
Params
No Tuning
Zero-Shot
04.78
N/A
13.41
N/A
22.03
N/A
13.41
N/A
Few-Shot
55.95
N/A
17.68
N/A
81.36
N/A
51.66
N/A
Prefix Injection
Prefix Tuning
02.65
0.034B
00.00
0.034B
28.47
0.004B
10.37
0.024B
Prompt Tuning
04.70
0.000B
26.22
0.000B
23.73
0.000B
18.22
0.000B
P-Tuning
34.19
0.001B
27.44
0.001B
43.05
0.001B
34.89
0.001B
Table 1: Comparison with existing parameter-efficient finetuning methods. Our ReMix consistently outperforms all baseline methods while maintaining strong parameter efficiency. We use the same parameter count budget for all methods to ensure fair comparison; for each method, we search for the best parameter count and report the accuracy under the best parameter count.
Figure 4: Both of our proposed RLOO and top- k selection contribute significantly to the strong performance of our ReMix.
Method
k=1
k=2
k=4
Rank- kr LoRA
56.10
54.51
59.21
Ours: k rank- r LoRAs
56.18
59.67
64.22
Table 2: The fact that our method significantly outperforms rank- kr LoRA clearly shows that our ReMix can activate diverse LoRA subsets: If the activated subset were always the same subset, then the results would be the same as rank- kr LoRA.
Method
LoRA-type ω
rsLoRA-type ω
ReMix (Ours)
53.30
55.72
Table 3: Our ReMix is not sensitive to routing weight ω .
Method
Per-Step Time
Total Time
Accuracy
MixLoRA
8.95 s
1:12:56
50.34
ReMix (Ours)
9.87 s
1:28:21
58.38
Table 4: Our ReMix significantly outperforms MixLoRA under similar training time.
Method
k=1
k=2
k=3
k=4
ReMix (Ours)
56.18
59.67
61.33
64.22
Table 5: Our ReMix consistently improves as we scale up the number k of activated LoRAs.
Figure 5: Our proposed ReMix can further benefit from scaling up the training compute. In contrast, HydraLoRA and MixLoRA have fixed training compute and thus cannot benefit from scaling up training compute.