ReMix: Reinforcement-Routing für Mischungen von LoRAs beim Feinabstimmen von LLMs

ReMix: Reinforcement routing for mixtures of LoRAs in LLM finetuning

March 10, 2026

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

Zusammenfassung

Low-Rank-Adapter (LoRAs) sind eine parameter-effiziente Feinabstimmungstechnik, die trainierbare Low-Rang-Matrizen in vortrainierte Modelle injiziert, um diese an neue Aufgaben anzupassen. Mixture-of-LoRAs-Modelle erweitern neuronale Netze effizient, indem sie jede Schichteingabe an eine kleine Teilmenge spezialisierter LoRAs der Schicht weiterleiten. Bestehende Mixture-of-LoRAs-Router weisen jedem LoRA ein gelerntes Routing-Gewicht zu, um ein End-to-End-Training des Routers zu ermöglichen. Trotz ihrer empirischen Aussichtsfähigkeit beobachten wir in der Praxis, dass die Routing-Gewichte typischerweise extrem unausgeglichen über die LoRAs verteilt sind, wobei oft nur ein oder zwei LoRAs die Routing-Gewichte dominieren. Dies begrenzt im Wesentlichen die Anzahl der effektiven LoRAs und behindert somit die Ausdruckskraft bestehender Mixture-of-LoRAs-Modelle erheblich. In dieser Arbeit führen wir diese Schwäche auf die Natur der erlernbaren Routing-Gewichte zurück und überdenken das grundlegende Design des Routers. Um dieses kritische Problem zu adressieren, schlagen wir einen neuen Router-Entwurf vor, den wir Reinforcement Routing for Mixture-of-LoRAs (ReMix) nennen. Unsere Schlüsselidee ist die Verwendung von nicht-erlernbaren Routing-Gewichten, um sicherzustellen, dass alle aktiven LoRAs gleichermaßen effektiv sind, ohne dass ein LoRA die Routing-Gewichte dominiert. Allerdings können unsere Router aufgrund der nicht-erlernbaren Routing-Gewichte nicht direkt über Gradientenabstieg trainiert werden. Daher schlagen wir weiterhin einen unverzerrten Gradientenschätzer für den Router vor, indem wir die Reinforce-Leave-One-Out (RLOO)-Technik anwenden, wobei wir den Überwachungsverlust als Belohnung und den Router als Policy im Reinforcement Learning betrachten. Unser Gradientenschätzer ermöglicht es auch, die Trainingsrechenleistung zu skalieren, um die Vorhersageleistung unseres ReMix zu steigern. Umfangreiche Experimente belegen, dass unser vorgeschlagenes ReMix state-of-the-art parameter-effiziente Feinabstimmungsmethoden bei einer vergleichbaren Anzahl aktivierter Parameter signifikant übertrifft.

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.