ReMix: Enrutamiento por refuerzo para mezclas de LoRAs en el ajuste fino de modelos de lenguaje grandes

ReMix: Reinforcement routing for mixtures of LoRAs in LLM finetuning

March 10, 2026

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

Resumen

Los adaptadores de bajo rango (LoRAs) son una técnica de ajuste fino eficiente en parámetros que inyecta matrices de bajo rango entrenables en modelos preentrenados para adaptarlos a nuevas tareas. Los modelos de Mezcla-de-LoRAs expanden las redes neuronales de manera eficiente enrutando cada entrada de capa a un pequeño subconjunto de LoRAs especializados de dicha capa. Los enrutadores existentes de Mezcla-de-LoRAs asignan un peso de enrutamiento aprendido a cada LoRA para permitir el entrenamiento de extremo a extremo del enrutador. A pesar de su potencial empírico, observamos que en la práctica los pesos de enrutamiento suelen estar extremadamente desequilibrados entre los LoRAs, donde frecuentemente solo uno o dos LoRAs dominan los pesos de enrutamiento. Esto esencialmente limita el número de LoRAs efectivos y, por lo tanto, obstaculiza severamente el poder expresivo de los modelos existentes de Mezcla-de-LoRAs. En este trabajo, atribuimos esta debilidad a la naturaleza de los pesos de enrutamiento aprendibles y replanteamos el diseño fundamental del enrutador. Para abordar este problema crítico, proponemos un nuevo diseño de enrutador que denominamos Enrutamiento por Refuerzo para Mezcla-de-LoRAs (ReMix). Nuestra idea clave es utilizar pesos de enrutamiento no aprendibles para garantizar que todos los LoRAs activos sean igualmente efectivos, sin que ningún LoRA domine los pesos de enrutamiento. Sin embargo, nuestros enrutadores no pueden entrenarse directamente mediante descenso de gradiente debido a nuestros pesos de enrutamiento no aprendibles. Por lo tanto, proponemos además un estimador de gradiente insesgado para el enrutador empleando la técnica de refuerzo "leave-one-out" (RLOO), donde consideramos la pérdida de supervisión como la recompensa y al enrutador como la política en el aprendizaje por refuerzo. Nuestro estimador de gradiente también permite escalar el cómputo de entrenamiento para impulsar el rendimiento predictivo de nuestro ReMix. Experimentos exhaustivos demuestran que nuestro ReMix propuesto supera significativamente a los métodos state-of-the-art de ajuste fino eficiente en parámetros con un número comparable de parámetros activados.

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.