Resumen Técnico: Cuando la fusión de modelos rompe el enrutamiento: Calibración sin entrenamiento para MoE

1. Declaración del problema

La fusión de modelos se ha convertido en una estrategia rentable para consolidar las capacidades de múltiples Modelos de Lenguaje Grande (LLM) sin los costes prohibitivos del reentrenamiento conjunto. Mientras que las técnicas existentes (por ejemplo, Task Arithmetic, TIES-Merging) fusionan con éxito arquitecturas densas mediante aritmética lineal de parámetros, fallan cuando se aplican a arquitecturas de Mezcla de Expertos (Mixture-of-Experts, MoE).

Los autores identifican un modo de fallo crítico denominado ruptura del enrutamiento (routing breakdown). En los modelos MoE, un enrutador despacha dinámicamente los tokens a un subconjunto de expertos utilizando una activación softmax no lineal seguida de una selección discreta Top-k. El artículo sostiene que los métodos de fusión lineal estándar no son adecuados para este entorno porque:

1. Desajuste de no linealidad: La interpolación lineal de los pesos del enrutador ($W_m = \sum \lambda_i W_i$) no se traduce en una interpolación lineal de las probabilidades de enrutamiento debido a la convexidad de la función softmax.
2. Sensibilidad: Pequeñas perturbaciones en los pesos del enrutador se amplifican exponencialmente por los mecanismos de softmax y Top-k.
3. Restricciones de equilibrio de carga: Los objetivos de preentrenamiento que imponen el equilibrio de carga hacen que las asignaciones de expertos sean altamente sensibles a los cambios de compuerta (gating).

Evidencia empírica: Experimentos preliminares al fusionar dos modelos OLMoE (ajustados para matemáticas y código) revelaron que el 50,11% de las selecciones de expertos se desviaron de los modelos originales inmediatamente después de la fusión. Esta inestabilidad se acumula a través de la profundidad de la red, donde el 99,98% de los tokens experimentan una divergencia de enrutamiento a lo largo de 16 capas. Crucialmente, la evaluación cruzada mostró que restaurar el enrutador original en el modelo fusionado recuperó el rendimiento, mientras que aplicar el enrutador fusionado al modelo original lo degradó, demostrando que preservar los parámetros de los expertos es inútil sin un comportamiento de enrutamiento consistente.

2. Metodología: Calibración del Enrutador Basada en el Hessiano (HARC)

Para abordar la ruptura del enrutamiento, los autores proponen HARC, un marco de trabajo sin entrenamiento que realinea el enrutador fusionado utilizando información de curvatura de segundo orden.

2.1. Formulación teórica

El objetivo es encontrar una matriz de pesos de enrutador unificada $W_m$ que minimice la divergencia de Kullback-Leibler (KL) entre la distribución de salida del enrutador fusionado y las distribuciones de los enrutadores de origen en todas las tareas.
$$W^*_m = \arg \min_{W_m} \sum_{i=1}^D \mathbb{E}_{x \sim \mathcal{D}_i} [\text{KL}(r_i(x) \parallel r_m(x; W_m))]$$

La optimización directa es intratable debido a la no linealidad de la función softmax. Los autores derivan una aproximación cuadrática de la divergencia KL utilizando una expansión de Taylor de segundo orden (Lema 4.1):
$$\text{KL}(r_i \parallel r_m) \approx \frac{1}{2}(z_m - z_i)^\top H_i (z_m - z_i)$$
donde $H_i = \text{diag}(r_i) - r_i r_i^\top$ es el Hessiano de la función de partición logarítmica. Este Hessiano captura la dinámica competitiva entre expertos (términos fuera de la diagonal) y la sensibilidad de los expertos activos (términos de la diagonal).

2.2. Solución de forma cerrada

La sustitución de la relación lineal $z = Wx$ en la aproximación cuadrática transforma el problema en un sistema de mínimos cuadrados ponderados. Los pesos óptimos $w^*_m$ (forma vectorizada de $W_m$) satisfacen el sistema lineal:
$$\left( \sum_{i=1}^D \Sigma_i \right) w^*_m = \sum_{i=1}^D \Sigma_i w_i$$
donde $\Sigma_i = \mathbb{E}_{x \sim \mathcal{D}_i} [H_i \otimes xx^\top]$ es la matriz de precisión específica de la tarea.

2.3. Computación eficiente (libre de matrices)

La construcción explícita e inversión de la matriz de precisión agregada $\Sigma$ es computacionalmente prohibitiva (complejidad de espacio $O(K^2d^2)$). HARC emplea dos estrategias clave para asegurar la escalabilidad:

1. Regularización diagonal: Para manejar la deficiencia de rango en el Hessiano (debido a la compuerta dispersa Top-k), los autores introducen un término de regularización que interpola entre la matriz de precisión completa y su aproximación diagonal. Esto estabiliza la solución sin asumir independencia entre expertos.
2. Gradiente Conjugado (CG) libre de matrices: En lugar de materializar $\Sigma$, el algoritmo utiliza un operador libre de matrices para calcular productos matriz-vector $\Sigma w$ sobre la marcha utilizando identidades de producto de Kronecker. Esto reduce la complejidad de espacio a $O(Kd)$ y permite el uso del método de Gradiente Conjugado para resolver el sistema de manera eficiente.

2.4. Algoritmo

HARC opera como un paso de post-procesamiento en un modelo fusionado no calibrado. Para cada capa MoE:

1. Realizar un pase hacia adelante con los datos de calibración a través del modelo fusionado hasta la capa actual para generar las características de entrada.
2. el cálculo de las salidas de enrutamiento y los Hessianos para los modelos de origen.
3. Acumular estadísticas para formar el sistema lineal.
4. Resolver los pesos del enrutador calibrado utilizando el resolvedor CG libre de matrices.
5. Actualizar los pesos del enrutador en el modelo fusionado.

3. Contribuciones clave

1. Identificación de la ruptura del enrutamiento: El artículo establece que la inestabilidad del enrutamiento, más que la interferencia de los parámetros de los expertos, es la causa principal de la degradación del rendimiento en la fusión de MoE.
2. Marco de trabajo HARC: Un novedoso método de calibración sin entrenamiento que aprovecha la información del Hessiano de segundo orden para realinear las distribuciones de enrutamiento, ofreciendo una solución de forma cerrada al problema de fusión no lineal.
3. Optimización escalable: La derivación de un resolvedor de Gradiente Conjugado libre de matrices que permite que HARC escale a grandes arquitecturas MoE sin la necesidad de materializar el Hessiano explícitamente.
4. Validación empírica: Demostración de que HARC mitiga eficazmente la ruptura del enrutamiento en diversos baselines de fusión (Weight Averaging, TIES-Merging, WUDI, etc.) y tareas (Matemáticas, Código).

4. Resultados experimentales

Los experimentos se realizaron en el modelo OLMoE-1B-7B-0125, fusionando variantes ajustadas para razonamiento matemático (GSM8K, MATH500) y generación de código (HumanEval+, MBPP+).

• Mejora constante: HARC mejoró consistentemente el rendimiento en todos los baselines de fusión probados. Por ejemplo, aplicado al sólido baseline WUDI-Merging, la puntuación general aumentó de 38,08 a 38,96.
• Mitigación de la sensibilidad al dominio: La fusión ingenua degradó severamente el rendimiento de la generación de código (por ejemplo, HumanEval+ cayó de 33,50 a 20,27 con Weight Averaging). HARC recuperó parcialmente esta pérdida, demostrando su capacidad para preservar las capacidades específicas de dominio corrigiendo la selección de expertos.
• Estudios de ablación:
  • Estructura del Hessiano: El uso del Hessiano completo superó a las aproximaciones diagonal o de identidad, confirmando que modelar la dinámica competitiva fuera de la diagonal es crucial.
  • Regularización: El factor de amortiguación $\alpha$ (que suprime los términos fuera de la diagonal) resultó ser más estable que el escalado diagonal $\gamma$.
  • Composición de datos: La calibración utilizando respuestas de origen filtradas como correctas (SrcCrt) produjo los mejores resultados. Sorprendentemente, el uso de solo prompts funcionó de manera comparable a las respuestas completas, lo que sugiere que las decisiones iniciales de enrutamiento son primordiales.
• Eficiencia de datos: HARC alcanzó un rendimiento cercano al óptimo utilizando solo el 40% de las muestras de calibración, lo que indica una convergencia rápida de las estadísticas de segundo orden.
• Análisis por capa: La visualización de la divergencia KL mostró que los errores de enrutamiento en los métodos base se acumulan exponencialmente en las capas más profundas. HARC suprimió con éxito esta tendencia, manteniendo la alineación en las capas profundas críticas.

5. Significado y afirmaciones

El artículo afirma que preservar la lógica de enrutamiento es tan esencial como alinear los pesos de los parámetros en arquitecturas dispersas. Al abordar el modo de fallo específico de la ruptura del enrutamiento, HARC ofrece un camino robusto y eficiente para consolidar modelos MoE sin reentrenamiento.

Los autores enfatizan que su método:

• No requiere reentrenamiento ni retropropagación de gradientes a través de los expertos.
• Se generaliza a través de diferentes estrategias de fusión (aritmética lineal, optimización) como un paso de post-procesamiento ligero.
• Reduce los costes computacionales y las barreras para desplegar modelos capaces al permitir el uso fiable de múltiples modelos de expertos.

El trabajo concluye que, si bien las hipótesis de fusión lineal existentes (Conectividad de Modo Lineal) se mantienen para las capas densas, son frágiles para los enrutadores MoE. HARC cierra esta brecha al considerar explícitamente la dinámica de compuerta no lineal inherente a las arquitecturas de Mezcla de Expertos.