PaperlayerView exact revision on arXiv

Source-linked AI summary

ReMix: Reinforcement routing for mixtures of LoRAs in LLM finetuning

Ruizhong Qiu, Hanqing Zeng, Yinglong Xia, Yiwen Meng, Ren Chen, Jiarui Feng, Dongqi Fu, Qifan Wang, Jiayi Liu, Jun Xiao, Xiangjun Fan, Benyu Zhang, Hong Li, Zhining Liu, Hyunsik Yoo, Zhichen Zeng, Tianxin Wei, Hanghang Tong

arXiv:2603.10160v1cs.LGcs.CL

TL;DR

Existing Mixture-of-LoRAs routers can concentrate routing weight on very few adapters, limiting effective LoRA utilization and model expressivity. ReMix uses equal non-learnable weights with an RLOO-based reinforcement-learning estimator, and experiments show consistent gains over parameter-efficient baselines with competitive parameter budgets.

  • Problem

    Existing Mixture-of-LoRAs routers can concentrate routing weights on one or two LoRAs, limiting the number of effective adapters and the mixture’s expressive power.

  • Method

    ReMix assigns equal non-learnable weights to activated LoRAs and trains its router with an unbiased RLOO-based reinforcement-learning gradient estimator.

  • Results

    ReMix consistently outperforms state-of-the-art parameter-efficient finetuning methods across diverse benchmarks while maintaining strong parameter efficiency.

  • Takeaways & Limitations

    ReMix achieves superior predictive power and computational efficiency through reinforcement-based routing under comparable parameter budgets.

Abstract

from arXiv · show

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.

1 INTRODUCTION

Mixture-of-LoRAs extend parameter-efficient finetuning by routing inputs among specialized LoRAs, but learned routing weights can collapse onto a few adapters. ReMix uses equal non-learnable weights and RLOO-based reinforcement training to address this limitation.

  • LoRAs inject trainable low-rank matrices into pretrained weights, enabling downstream adaptation with relatively few trainable parameters.
  • Mixture-of-LoRAs route each input through a small pool of layer-specific LoRAs to improve specialization, efficiency, and expressive power.
  • Learned routing weights often become extremely imbalanced, with one or two LoRAs dominating and the remaining adapters effectively disabled.
  • ReMix assigns constant weights to all activated LoRAs, ensuring equal contribution and avoiding collapse into one dominant adapter.
  • RLOO reformulates router training as reinforcement learning, treating supervised finetuning loss as negative reward and the router as the policy.
  • Across diverse benchmarks, ReMix consistently outperforms state-of-the-art parameter-efficient finetuning methods under the same parameter budgets.

2 MOTIVATION: ROUTING WEIGHT COLLAPSE

Existing Mixture-of-LoRAs routers use learned weights to combine specialized adapters, but theoretical and empirical analyses show that these weights concentrate on very few LoRAs. The effective support can rapidly fall to one during finetuning, limiting adapter utilization and expressive power.

  • 2.1 PRELIMINARIES: MIXTURE OF LORAS: Mixture-of-LoRAs dynamically selects a small number of LoRAs per input through a lightweight router, increasing model capacity with few adapters.
  • 2.1 PRELIMINARIES: MIXTURE OF LORAS: Routers predict categorical distributions over LoRAs, whose weights combine the corresponding low-rank transformations with frozen layer weights.
  • 2.2 THEORETICAL ANALYSIS: Theoretical analysis shows that softmax-normalized learned weights tend to concentrate most mass on only one or two LoRAs.
  • 2.2 THEORETICAL ANALYSIS: ESS measures the number of LoRAs with relatively large routing weights, ranging from 1 for one-hot weights to n for uniform weights.
  • 2.2 THEORETICAL ANALYSIS: With σ = 1, n = 8, and a Rademacher input in D = 1024 dimensions, Theorem 1 gives at least 84.19% probability that at most two LoRAs have relatively large weights.
  • 2.3 EMPIRICAL ANALYSIS: Routing collapse limits realized expressive power because small weights reduce LoRA contributions and can suppress their back-propagated gradients.
  • 2.3 EMPIRICAL ANALYSIS: Empirical tracking in MixLoRA shows one dominant LoRA in deeper layers, while the other seven receive negligibly small routing weights.
  • 2.3 EMPIRICAL ANALYSIS: ESS often decreases to 1 quickly during finetuning, even when it begins near 4, and remains there afterward.

3 SIMPLE YET EFFECTIVE METHOD: REMIX

ReMix replaces learnable routing weights with equal constant weights for activated LoRAs, then trains the router through sampling and an unbiased RLOO gradient estimator. During inference, it uses theoretically motivated top-k selection.

  • 3.1 Adapter Architecture: Non-Learnable Weight: ReMix assigns the same constant weight to all k activated LoRAs and zero weight to non-activated LoRAs, preventing routing-weight collapse.The constant weight uses either the LoRA or rsLoRA scaling form, and the method is reported to be insensitive to this choice.
  • 3.1 Adapter Architecture: Non-Learnable Weight: The resulting routing weights have effective sample size ESS(π^(l)) = k, and the layer output is a weighted sum over the k activated LoRAs.The sparse weights simplify the computation of the layer output.
  • 3.2 Finetuning Procedure: RLOO: Because constant routing weights are not differentiable with respect to router parameters, ReMix samples LoRA subsets from the router distribution to define a differentiable expected SFT loss.LoRAs are sampled without replacement, ensuring that the k activated LoRAs are distinct.
  • 3.2 Finetuning Procedure: RLOO: ReMix treats router training as reinforcement learning, using the SFT loss as negative reward and an RLOO estimator to approximate the surrogate router gradient.The estimator is unbiased, and M independent selections represent the training compute budget.
  • 3.3 Inference Procedure: Top-k Selection: During inference, ReMix uses top-k LoRA selection because Theorem 2 guarantees optimality when the router samples the optimal subset with probability above 50%.Under that condition, top-k selection improves the probability of choosing the optimal subset to 100%.
  • 3.3 Inference Procedure: Top-k Selection: Table 1 compares methods under the same parameter-count budget, reporting each method at its best searched parameter count.The caption states that ReMix consistently outperforms all baseline methods while maintaining strong parameter efficiency.

4 EXPERIMENTS

Across diverse benchmarks, ReMix consistently outperforms parameter-efficient baselines while maintaining competitive parameter efficiency. Ablations and diversity comparisons support the contributions of RLOO, top-k selection, and varied LoRA subsets.

  • Main results: ReMix consistently outperforms all baselines across HumanEval, GSM8K, and ARC-c while maintaining strong parameter efficiency.The evaluation covers code generation, mathematical reasoning, and knowledge recall.
  • Main results: 2.82 average accuracy improvement over the strongest competing approach is achieved across the three tasks.ReMix also exceeds the strongest Prefix Injection, Weight Modulation, and Mixture baselines by 25.88, 2.82, and 3.34, respectively.
  • Parameter efficiency: 0.070B trainable parameters deliver these gains, including a 31% reduction versus MixLoRA’s 0.101B.Compared with VB-LoRA’s 0.675B, this is a 90% reduction; versus rsLoRA, ReMix improves average accuracy by 2.82 with 0.042B more parameters.
  • Ablation studies: Removing either RLOO or top-k selection lowers accuracy relative to full ReMix, indicating both components contribute to performance.The ablation study compares the complete method with variants excluding each component.
  • Diversity of activated LoRA subsets: 64.22 accuracy for k = 4 surpasses rank-32 LoRA’s 59.21, demonstrating that ReMix selects diverse LoRA subsets.If the same subset were always activated, the mixture would match a single rank-kr LoRA with the same parameter count.

4.5 TRAINING EFFICIENCY

ReMix retains a training-time advantage over MixLoRA under comparable compute and improves as activated LoRAs or sampled selections increase. Its RL-based estimator enables performance gains from additional training compute.

  • Training efficiency: 58.38% accuracy in 1:28:21 outperforms MixLoRA’s 50.34% in 1:12:56 under similar training time.ReMix uses only 10% more training time and yields a 15.97% relative accuracy improvement.
  • Scaling activated LoRAs: Larger k consistently improves ReMix results when k ≤ n/2 on GSM8K with n = 8.The study varies the number of activated LoRAs to examine the efficiency–accuracy tradeoff.
  • Scaling training compute: Increasing sampled selections M from 2 to 32 raises accuracy from 56.03% to 58.83%.The RL-based gradient estimator allows ReMix to scale training compute through additional sampled selections, unlike methods with fixed deterministic training compute.
  • Routing-weight sensitivity: ReMix shows only a very small performance difference between LoRA-type and rsLoRA-type routing weights under k = 3 on GSM8K.The comparison indicates limited sensitivity to the routing-weight choice.

5 RELATED WORK

Related work develops multi-LoRA architectures that reorganize or combine adapters to improve flexibility, representation balance, or efficiency.

  • LoraHub dynamically composes multiple LoRAs into a unified pipeline for cross-task generalization without retraining.
  • MultiLoRA structurally initializes LoRA subspaces and horizontally expands adapters across layers to mitigate top-singular-vector dominance.
  • HydraLoRA uses an asymmetric architecture that decouples projection and update pathways to improve parameter and training efficiency.

6 CONCLUSION

The paper addresses imbalanced routing weights with ReMix and reports consistent improvements over state-of-the-art parameter-efficient finetuning methods across diverse benchmarks.

  • ReMix targets imbalanced routing weights that hinder effective LoRA utilization through a reinforcement-based router design.
  • ReMix consistently outperforms state-of-the-art parameter-efficient finetuning methods across diverse benchmarks.
  • The reported improvements include superior predictive power and computational efficiency.

A RELATED WORK (CONT’D)

The related work situates the paper within neural-network and Transformer research, citing broad developments across these architectures.

  • Neural networks have been prevalent across various domains.
  • Transformers have become the de facto neural architecture.
  • The discussion cites extensive prior work on neural networks and Transformers.

B.1 PROOF OF THEOREM 1

The proof of Theorem 1 combines Gaussian distribution properties, order statistics, and several technical integral and bound lemmas to establish a routing-logit gap.

  • Technical lemmata: The proof relies on Gaussian gap, inverse, upper-tail, integral, and integer-function estimates before applying the resulting bounds to Theorem 1.
  • Proof of Theorem 1: The proof represents routing logits as ξ = P^(l)x^(l), with routing weights obtained by applying softmax to ξ.
  • Proof of Theorem 1: Because P^(l) has i.i.d. Gaussian entries, the routing logits are also modeled as independent Gaussian variables with scale determined by σ²∥x^(l)∥².
  • Proof of Theorem 1: A Gaussian order-statistics estimate yields, with probability at least 1 − δ, a lower bound on the gap between the largest and second-largest routing logits.

B.2 PROOF OF THEOREM 2

The proof establishes Theorem 2 by formalizing ordered and unordered sampling probabilities and showing that replacing a selected LoRA with one having at least as large q_i increases the unordered sampling probability. An inductive sequence of replacements then yields a contradiction to the assumed premise.

  • Definitions: The proof defines q(I) for ordered size-k subsets sampled without replacement and Q(I) for unordered subsets sampled without replacement.It also introduces permutation actions to relate ordered samples to unordered sampling probabilities.
  • Swapping lemma: Lemma 10 states that replacing a LoRA i in a size-k subset with an outside LoRA i† satisfying q_i ≤ q_i† increases the unordered sampling probability.The lemma proof represents the replacement across all permutations of the ordered subset.
  • Contradiction setup: Theorem 2 assumes a premise and aims to derive a contradiction, with ties in argtop broken arbitrarily.The proof then recalls the relevant definition before constructing the contradiction.
  • Inductive replacement: When I* and I† differ, their intersection has size less than k, so the proof replaces each differing LoRA in I* inductively to obtain I†.Each replacement inherits all other LoRAs from the previous subset.
  • Contradiction: Applying Lemma 10 at every replacement gives Q(eI_j) ≥ Q(eI_{j−1}), and the resulting endpoint probabilities imply Q(I†) + Q(I*) > 1.The proof identifies this inequality as contradicting a previously established fact and concludes by falsifying the premise.
Loading 2603.10160v1…