Resumen Técnico: Aprender a Seleccionar, No a Reaprender: Mezclas de LoRAs de Razonamiento con Enrutamiento Rígido

Declaración del Problema

El artículo aborda el desafío de componer adaptadores de Adaptación de Bajo Rango (LoRA) entrenados de forma independiente en un único Modelo de Lenguaje Grande (LLM) para la adaptación multidominio, particularmente cuando los datos de entrenamiento originales no pueden compartirse debido a restricciones de privacidad o regulatorias.

Un problema crítico surge al aplicar el enrutamiento suave (soft routing) de la Mezcla de Expertos (MoE) estándar a adaptadores LoRA entrenados de forma independiente y congelados. Los módulos LoRA se entrenan bajo la suposición de una actualización aditiva de escala unitaria ($W_{updated} = W_{pretrained} + AB$). Sin embargo, el enrutamiento MoE suave estándar combina los expertos mediante sumas ponderadas ($W_{updated} = W_{pretrained} + \sum \omega_i A_i B_i$), donde los pesos de enrutamiento $\omega_i < 1$. Esto crea un desajuste de escala: la actualización efectiva aplicada durante la integración difiere de la actualización de escala unitaria bajo la cual se entrenó originalmente el experto.

Las soluciones existentes, como LoRA-Mixer, intentan resolver esto reentrenando los adaptadores con pérdidas de preservación auxiliares. Este enfoque aumenta el número de parámetros entrenables, requiere más datos y corre el riesgo de alterar el conocimiento codificado en los expertos originales. Además, este desafío es particularmente pronunciado para los modelos con ajuste de instrucciones y Cadena de Pensamiento (CoT), donde el ajuste fino convencional puede degradar los comportos de razonamiento internos, y los métodos de composición existentes a menudo fallan en preservar estos rastros de razonamiento especializados.

Metodología

Los autores proponen Hard-Routed MoR-LoRA, un marco de dos etapas diseñado para componer expertos LoRA de razonamiento congelados mediante una selección rígida de escala unitaria, evitando así la necesidad de reentrenar los expertos.

Etapa I: Entrenamiento de Expertos de Razonamiento

Los adaptadores LoRA específicos de dominio se entrenan de forma independiente utilizando Aprendizaje por Refuerzo con Retroalimentación Verificable (RLVF). A diferencia del ajuste fino supervisado (SFT), que fomenta la imitación de salidas de referencia, el RLVF optimiza la corrección utilizando recompensas automáticamente verificables (por ejemplo, comprobando la respuesta final). Este paradigma fomenta la emergencia de un CoT estructurado sin sobrescribir los patrones de razonamiento alineados del modelo base. Los expertos resultantes se congelan posteriormente.

Etapa II: Integración con Enrutamiento Rígido

Los expertos congelados se integran en un modelo compartido utilizando una Mezcla de Enrutamiento Rígido (Hard-Routed Mixture).

• Enrutamiento Rígido Top-1 (Hard Top-1 Routing): Para cada token, el enrutador selecciona exactamente un experto. Esto asegura que el experto seleccionado se aplique con escala unitaria, preservando la suposición original de LoRA ($W_{updated} = W_{pretrained} + 1 \cdot A_j B_j$).
• Estimador de Paso Directo (Straight-Through Estimator - STE): Dado que la selección discreta (argmax) no es diferenciable, los autores emplean STE. En la pasada hacia adelante (forward pass), se utiliza la decisión rígida; en la pasada hacia atrás (backward pass), los gradientes se propagan a través de las probabilidades suaves. Esto permite el entrenamiento de extremo a extremo del enrutador mientras se mantiene el comportamiento discreto de escala unitaria durante la inferencia.
• Eficiencia de Parámetros: Solo se entrenan un enrutador lineal ligero compartido y un pequeño módulo de atención LoRA. Los expertos preentrenados permanecen congelados.
• Destilación: El enrutador se entrena mediante Ajuste Fino Supervisado (SFT) en un pequeño conjunto de datos de rastros de razonamiento destilados de los expertos congelados. El objetivo es aprender qué experto seleccionar, no reaprender la capacidad de razonamiento en sí misma.

Contribuciones Clave

1. Identificación del Desajuste de Escala: El artículo identifica que aplicar el enrutamiento MoE suave a adaptadores LoRA congelados entrenados de forma independiente introduce un desajuste de escala que altera la función efectiva de los adaptadores.
2. Enrutamiento Rígido de Escala Unitaria: Los autores proponen el uso de enrutamiento rígido Top-1 con un Estimador de Paso Directo para preservar la aplicación de escala unitaria de los expertos, permitiendo la composición modular sin reentrenamiento.
3. Marco de Dos Etapas: Un flujo de trabajo novedoso que separa la adquisición de razonamiento (vía RLVF) de la selección de expertos (vía enrutamiento rígido), manteniendo los expertos congelados durante la integración.
4. Validación Empírica: El estudio demuestra que este enfoque preserva el comportamiento del experto mientras requiere sustancialmente menos parámetros entrenables que las líneas base de enrutamiento suave.

Resultados Experimentales

Los experimentos se realizaron en cinco bancos de pruebas (GSM8K, ARC-Challenge, MedQA, BoolQ, CoLA) y múltiples escalas de modelos (1B, 3B, 8B) utilizando las familias LLaMA-3, Gemma-3 y Qwen-2.5.

• Superioridad de RLVF: Los expertos entrenados con RLVF superaron significativamente a los expertos entrenados con SFT en tareas de razonamiento intensivo para modelos capaces (3B y 8B), mientras que el SFT a veces degradó el rendimiento del razonamiento en modelos grandes con ajuste de instrucciones.
• Rendimiento vs. Parámetros: Hard-Routed MoR-LoRA logró un rendimiento comparable o mejor que las líneas base de enrutamiento suave (LoRAMixer, MoLE) mientras entrenaba sustancialmente menos parámetros. Por ejemplo, en el modelo de 3B con expertos RLVF, el método propuesto entrenó ~73M de parámetros para alcanzar una precisión promedio de 72.07, mientras que LoRAMixer requirió ~606M de parámetros para alcanzar 65.67.
• Análisis del Comportamiento de Enrutamiento: El análisis del enrutamiento suave Top-2 reveló que a menudo colapsa hacia una selección cercana a Top-1 (peso dominante promedio ~0.71), lo que sugiere que el enrutamiento rígido captura el comportamiento deseado de manera más directa y eficiente.
• Eficiencia de Datos: La etapa de integración (Etapa II) requirió solo ~1,000 muestras por conjunto de datos para lograr un rendimiento cercano al entrenamiento con el conjunto de datos completo, lo que indica que el enrutador aprende patrones de selección en lugar de conocimiento de dominio.
• Capacidad de Dominio Mixto: En evaluaciones de dominio mixto (por ejemplo, un prompt que contiene tanto matemáticas como comprensión lectora), el enrutamiento rígido a nivel de token superó al enrutamiento a nivel de prompt, demostrando flexibilidad para seleccionar diferentes expertos dentro de una misma entrada.

Significado y Afirmaciones

El artículo afirma que la adaptación modular puede lograrse aprendiendo cuándo usar cada experto en lugar de modificar qué sabe cada experto.

Los autores argumentan que, para los expertos LoRA preentrenados y congelados, el enrutamiento rígido de escala unitaria proporciona una abstracción simple y eficiente. Al tratar la integración principalmente como un problema de enrutamiento, el método evita el costo computacional y la posible degradación del conocimiento asociados con el reentrenamiento de los expertos. Los resultados sugieren que el beneficio principal de la composición proviene de seleccionar el experto congelado apropiado para el contexto, en lugar de modificar extensamente los parámetros de los expertos.

El trabajo destaca que, si bien el enrutamiento a nivel de prompt es suficiente para entradas limpias de un solo dominio, el enrutamiento rígido a nivel de token ofrece la flexibilidad necesaria para entradas complejas de dominios mixtos. Los autores señalan modestamente que el método no es enteramente libre de entrenamiento, ya que la Etapa II requiere algunos datos etiquetados para aprender la política de enrutamiento, y actualmente no soporta la integración zero-shot real de nuevos dominios sin datos de enrutamiento adicionales.