ReMix: Reinforcement routing voor mengsels van LoRA's in fijnafstemming van grote taalmodellen

ReMix: Reinforcement routing for mixtures of LoRAs in LLM finetuning

March 10, 2026

Auteurs: 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

cs.AI

Samenvatting

Laagrangepasstukken (LoRA's) zijn een parameter-efficiënte fijnafstemmingstechniek die trainbare laagrangematrices injecteert in vooraf getrainde modellen om ze aan nieuwe taken aan te passen. Mixture-of-LoRA's-modellen breiden neurale netwerken efficiënt uit door elke laaginvoer door te sturen naar een kleine subset van gespecialiseerde LoRA's van de laag. Bestaande Mixture-of-LoRA's-routers kennen een geleerd routeringsgewicht toe aan elke LoRA om end-to-end training van de router mogelijk te maken. Ondanks hun empirische potentie merken we dat de routeringsgewichten in de praktijk typisch extreem onevenwichtig zijn over de LoRA's, waarbij vaak slechts één of twee LoRA's de routeringsgewichten domineren. Dit beperkt in wezen het aantal effectieve LoRA's en belemmert daardoor ernstig de expressieve kracht van bestaande Mixture-of-LoRA's-modellen. In dit werk schrijven we deze zwakte toe aan de aard van leerbare routeringsgewichten en heroverwegen we het fundamentele ontwerp van de router. Om dit kritieke probleem aan te pakken, stellen we een nieuwe routerontwerp voor dat we Reinforcement Routing for Mixture-of-LoRA's (ReMix) noemen. Ons kernidee is het gebruik van niet-leerbare routeringsgewichten om ervoor te zorgen dat alle actieve LoRA's even effectief zijn, zonder dat één LoRA de routeringsgewanten domineert. Onze routers kunnen echter niet direct worden getraind via gradient descent vanwege onze niet-leerbare routeringsgewichten. Daarom stellen we verder een zuivere gradientenschatter voor de router voor door gebruik te maken van de reinforce leave-one-out (RLOO) techniek, waarbij we de supervisieverlies beschouwen als de beloning en de router als het beleid in reinforcement learning. Onze gradientenschatter maakt het ook mogelijk de rekenkracht voor training op te schalen om de voorspellende prestaties van onze ReMix te verbeteren. Uitgebreide experimenten tonen aan dat onze voorgestelde ReMix state-of-the-art parameter-efficiënte fijnafstemmingsmethoden significant overtreft bij een vergelijkbaar aantal geactiveerde parameters.

English

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.