ReMix: Roteamento por Reforço para Misturas de LoRAs no Ajuste Fino de LLMs

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

Resumo

Os adaptadores de baixo posto (LoRAs) são uma técnica de ajuste fino com eficiência de parâmetros que injeta matrizes de baixo posto treináveis em modelos pré-treinados para adaptá-los a novas tarefas. Os modelos de Mistura-de-LoRAs expandem redes neurais de forma eficiente ao rotear cada entrada de camada para um pequeno subconjunto de LoRAs especializados da camada. Os roteadores de Mistura-de-LoRAs existentes atribuem um peso de roteamento aprendido a cada LoRA para permitir o treinamento end-to-end do roteador. Apesar de seu potencial empírico, observamos que os pesos de roteamento são tipicamente extremamente desbalanceados entre os LoRAs na prática, onde frequentemente apenas um ou dois LoRAs dominam os pesos de roteamento. Isso essencialmente limita o número de LoRAs efetivos e, portanto, prejudica severamente o poder expressivo dos modelos existentes de Mistura-de-LoRAs. Neste trabalho, atribuímos essa fraqueza à natureza dos pesos de roteamento aprendíveis e repensamos o design fundamental do roteador. Para abordar essa questão crítica, propomos um novo design de roteador que chamamos de Roteamento por Reforço para Mistura-de-LoRAs (ReMix). Nossa ideia central é usar pesos de roteamento não aprendíveis para garantir que todos os LoRAs ativos sejam igualmente efetivos, sem que nenhum LoRA domine os pesos de roteamento. No entanto, nossos roteadores não podem ser treinados diretamente via gradiente descendente devido aos nossos pesos de roteamento não aprendíveis. Por isso, propomos ainda um estimador de gradiente não enviesado para o roteador empregando a técnica de reforço leave-one-out (RLOO), onde consideramos a perda de supervisão como recompensa e o roteador como a política no aprendizado por reforço. Nosso estimador de gradiente também permite escalar o poder computacional de treinamento para impulsionar o desempenho preditivo do nosso ReMix. Experimentos extensivos demonstram que nosso ReMix proposto supera significativamente os métodos state-of-the-art de ajuste fino com eficiência de parâmetros sob um número comparável de parâmetros ativados.

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.