Resumo Técnico: Aprendendo a Selecionar, Não a Reaprender: Misturas de LoRAs de Raciocínio com Roteamento Rígido (Hard-Routed)

Declaração do Problema

O artigo aborda o desafio de compor adaptadores de Adaptação de Baixo Rank (LoRA) treinados independentemente em um único Grande Modelo de Linguagem (LLM) para adaptação multidomínio, particularmente quando os dados de treinamento originais não podem ser compartilhados devido a restrições de privacidade ou regulatórias.

Um problema crítico surge ao aplicar o roteamento suave (soft routing) padrão de Mistura de Especialistas (MoE) a adaptadores LoRA congelados e treinados de forma independente. Os módulos LoRA são treinados sob a premissa de uma atualização de escala unitária ($W_{updated} = W_{pretrained} + AB$). No entanto, o roteamento MoE suave padrão combina especialistas usando somas ponderadas ($W_{updated} = W_{pretrained} + \sum \omega_i A_i B_i$), onde os pesos de roteamento $\omega_i < 1$. Isso cria um descompasso de escala: a atualização efetiva aplicada durante a integração difere da atualização de escala unitária sob a qual o especialista foi originalmente treinado.

Soluções existentes, como o LoRA-Mixer, tentam resolver isso retranejando os adaptadores com perdas de preservação auxiliares. Essa abordagem aumenta o número de parâmetros treináveis, exige mais dados e corre o risco de alterar o conhecimento codificado nos especialistas originais. Além disso, esse desafio é particularmente pronunciado para modelos com ajuste de instrução (instruction-tuned) que utilizam Cadeia de Pensamento (CoT), onde o ajuste fino convencional pode degradar os comportamentos de raciocínio internos, e os métodos de composição existentes frequentemente falham em preservar esses traços especializados de raciocínio.

Metodologia

Os autores propõem o Hard-Routed MoR-LoRA, um framework de dois estágios projetado para compor especialistas LoRA de raciocínio congelados através de seleção rígida de escala unitária (unit-scale hard selection), evitando assim a necessidade de retreinar os especialistas.

Estágio I: Treinamento de Especialistas de Raciocínio

Adaptadores LoRA específicos de domínio são treinados independentemente usando Aprendizado por Reforço com Feedback Verificável (RLVF). Diferente do ajuste fino supervisionado (SFT), que incentiva a imitação de saídas de referência, o RLVF otimiza a correção usando recompensas automaticamente verificáveis (ex: checar a resposta final). Esse paradigma incentiva a emergência de um CoT estruturado sem sobrepor os padrões de raciocínio alinhados do modelo base. Os especialistas resultantes são então congelados.

Estágio II: Integração com Roteamento Rígido

Os especialistas congelados são integrados em um modelo compartilhado usando uma Mistura de Roteamento Rido (Hard-Routed Mixture).

• Roteamento Rígido Top-1 (Hard Top-1 Routing): Para cada token, o roteador seleciona exatamente um especialista. Isso garante que o especialista selecionado seja aplicado com escala unitária, preservando a suposição original do LoRA ($W_{updated} = W_{pretrained} + 1 \cdot A_j B_j$).
• Estimador de Passagem Direta (Straight-Through Estimator - STE): Como a seleção discreta (argmax) não é diferenciável, os autores empregam o STE. Na passagem direta (forward pass), a decisão rígida é utilizada; na passagem de retropropagação (backward pass), os gradientes são propagados através das probabilidades suaves. Isso permite o treinamento de ponta a ponta do roteador enquanto mantém o comportamento discreto de escala unitária durante a inferência.
• Eficiência de Parâmetros: Apenas um roteador linear compartilhado leve e um pequeno módulo de atenção LoRA são treinados. Os especialistas pré-treinados permanecem congelados.
• Destilação: O roteador é treinado via Ajuste Fino Supervisionado (SFT) em um pequeno conjunto de dados de traços de raciocínio destilados dos especialistas congelados. O objetivo é aprender qual especialista selecionar, não reaprender a capacidade de raciocínio em si.

Principais Contribuições

1. Identificação do Descompasso de Escala: O artigo identifica que a aplicação de roteamento MoE suave a adaptadores LoRA congelados e treinados independentemente introduz um descompasso de escala que altera a função efetiva dos adaptadores.
2. Roteamento Rígido de Escala Unitária: Os autores propõem o uso de roteamento rígido Top-1 com um Estimador de Passagem Direta para preservar a aplicação de escala unitária dos especialistas, permitindo a composição modular sem retreino.
3. Framework de Dois Estágios: Um pipeline inovador que separa a aquisição de raciocínio (via RLVF) da seleção de especialistas (via roteamento rígido), mantendo os especialistas congelados durante a integração.
4. Validação Empírica: O estudo demonstra que esta abordagem preserva o comportamento do especialista enquanto requer substancialmente menos parâmetros treináveis do que as linhas de base de roteamento suave.

Resultados Experimentais

Os experimentos foram conduzidos em cinco benchmarks (GSM8K, ARC-Challenge, MedQA, BoolQ, CoLA) e múltiplas escalas de modelos (1B, 3B, 8B) utilizando as famílias LLaMA-3, Gemma-3 e Qwen-2.5.

• Superioridade do RLVF: Especialistas treinados com RLVF superaram significativamente os especialistas treinados com SFT em tarefas de raciocínio intenso para modelos capazes (3B e 8B), enquanto o SFT às vezes degradou o desempenho de raciocínio em modelos grandes ajustados por instrução.
• Desempenho vs. Parâmetros: O Hard-Routed MoR-LoRA alcançou desempenho comparável ou superior aos baselines de roteamento suave (LoRAMixer, MoLE) enquanto treinava substancialmente menos parâmetros. Por exemplo, no modelo de 3B com especialistas RLVF, o método proposto treinou ~73M de parâmetros para atingir 72,07 de acurácia média, enquanto o LoRAMixer exigiu ~606M de parâmetros para atingir 65,67.
• Análise de Comportamento de Roteamento: A análise do roteamento suave Top-2 revelou que ele frequentemente colapsa para uma seleção próxima ao Top-1 (peso dominante médio de ~0,71), sugerindo que o roteamento rígido captura o comportamento pretendido de forma mais direta e eficiente.
• Eficiência de Dados: A etapa de integração (Estágio II) exigiu apenas ~1.000 amostras por dataset para atingir um desempenho próximo ao treinamento no dataset completo, indicando que o roteador aprende padrões de seleção, não o conhecimento do domínio.
• Capacidade de Domínios Mistos: Em avaliações de domínio misto (ex: um prompt contendo matemática e compreensão de leitura), o roteamento rígido ao nível de token superou o roteamento ao nível de prompt, demonstrando flexibilidade na seleção de diferentes especialistas dentro de uma única entrada.

Significância e Alegações

O artigo afirma que a adaptação modular pode ser alcançada aprendendo quando usar cada especialista, em vez de modificar o que cada especialista sabe.

Os autores argumentam que, para especialistas LoRA pré-treinados e congelados, o roteamento rígido de escala unitária fornece uma abstração simples e eficiente. Ao tratar a integração primariamente como um problema de roteamento, o método evita o custo computacional e a potencial degradação do conhecimento associados ao retreino dos especialistas. Os resultados sugerem que o principal benefício da composição vem da seleção do especialista apropriado para o contexto, em vez da modificação extensiva dos parâmetros dos especialistas.

O trabalho destaca que, embora o roteamento ao nível de prompt seja suficiente para entradas limpas de domínio único, o roteamento rígido ao nível de token oferece a flexibilidade necessária para entradas complexas de domínios mistos. Os autores observam modestamente que o método não é inteiramente livre de treinamento, pois o Estágio II requer alguns dados rotulados para aprender a política de roteamento, e atualmente não suporta a integração zero-shot real de novos domínios sem dados de roteamento adicionais.