Transport and Merge: Cross-Architecture Merging for Large Language Models

1. El Problema

Los Grandes Modelos de Lenguaje (LLMs) actuales logran capacidades excepcionales al escalar la capacidad del modelo y los datos de entrenamiento. Sin embargo, muchas implementaciones del mundo real dependen de modelos más pequeños entrenados o adaptados con datos de recursos limitados (ej. lenguas minoritarias como el malayo, cantonés o tailandés, o dominios especializados como medicina y finanzas).

El desafío principal es transferir conocimiento de modelos grandes y ricos en recursos (fuente) a modelos pequeños y con restricciones de recursos (objetivo). Aunque el fusionado de modelos (model merging) es una alternativa eficiente al ajuste fino basado en gradientes, la mayoría de los métodos existentes asumen que los modelos fuente y objetivo comparten la misma arquitectura. Esto impide la transferencia directa de parámetros entre arquitecturas heterogéneas (por ejemplo, fusionar un LLaMA-8B con un modelo de 1B de una arquitectura diferente), limitando su aplicabilidad en escenarios de recursos limitados.

2. Metodología

Los autores proponen un marco de trabajo basado en Transporte Óptimo (Optimal Transport - OT) para alinear modelos heterogéneos y fusionar sus pesos directamente sin necesidad de entrenamiento extensivo. El enfoque se basa en la observación de que, a pesar de las diferencias arquitectónicas, los modelos de lenguaje exhiben activaciones internas correlacionadas al procesar las mismas entradas.

El proceso se divide en las siguientes etapas clave:

• Extracción de Activaciones: Se utiliza un pequeño conjunto de datos de calibración ($D$) para extraer las activaciones intermedias de las capas del modelo fuente ($M_B$) y del modelo objetivo ($M_A$).
• Alineación mediante Transporte Óptimo (OT):
  • Nivel de Características (Features): Para cada par de capas (fuente $m$, objetivo $\ell$), se calcula una matriz de costos basada en la correlación de Pearson entre los canales de activación. Se resuelve un problema de Transporte Óptimo regularizado entrópicamente para obtener una matriz de transporte $Q^{\ell m}$, que define cómo las características del modelo fuente se combinan para igualar las del objetivo.
  • Nivel de Capas: Se agregan los costos de las características para inferir una correspondencia global entre capas ($P$), determinando qué capas del modelo fuente son más relevantes para cada capa del objetivo.
• Fusión de Pesos (Weight Fusion):
  • Las matrices de transporte ($Q$ y $P$) se utilizan para mapear los pesos del modelo fuente al espacio de neuronas del modelo objetivo.
  • Selección Selectiva (Top-k): Para evitar interferencias y mantener la robustez, la fusión no se aplica a todas las neuronas. Se seleccionan las $k$ neuronas más activas en el modelo objetivo (basado en el conjunto de datos $D$) y solo se fusionan los pesos correspondientes a estas neuronas.
  • La fórmula de fusión combina los pesos originales del objetivo con una versión transformada de los pesos del fuente, controlada por un coeficiente $\alpha$.
• Adaptación con Residuos Congelados (Residual-Frozen Adaptation):
  • Tras la fusión, se realiza un ajuste fino ligero. Se congelan los parámetros residuales transferidos (el conocimiento inyectado) y solo se actualizan los parámetros base originales del modelo objetivo. Esto asegura que el conocimiento transferido no se sobrescriba, permitiendo al modelo adaptarse a la distribución de la tarea específica.

3. Contribuciones Clave

• Marco de Fusión Cross-Arquitectura: Es uno de los primeros métodos que permite la fusión directa de parámetros entre modelos de LLM con arquitecturas y capacidades diferentes, sin depender de la destilación de conocimiento (distillation) previa.
• Correspondencia Estructurada Neuron a Neuron: Utiliza el Transporte Óptimo para inferir correspondencias estructuradas a nivel de neuronas basadas en estadísticas de activación, en lugar de alinear pesos directamente.
• Interpretación en el Espacio de Representación: Demuestran teóricamente que la fusión de pesos es equivalente a mapear las características del objetivo a las coordenadas del fuente, aplicar la transformación lineal del fuente y devolverlas al espacio del objetivo.
• Eficiencia de Datos: El método requiere únicamente un pequeño conjunto de datos de entrada (calibración) para inferir las correspondencias, eliminando la necesidad de grandes corpus de entrenamiento para la transferencia.

4. Resultados

Los experimentos se realizaron en múltiples lenguas de bajos recursos (Malayo, Indonesio, Cantonés, Tailandés) y dominios especializados (Finanzas, Medicina).

• Rendimiento Superior: El método propuesto ("Fused w/ Adaptation") superó consistentemente a los modelos base y a otras estrategias de transferencia como el Ajuste Fino Supervisado (SFT) y la Destilación de Conocimiento en la mayoría de las métricas.
  • Ejemplo: En el benchmark MalayMMLU, la fusión con adaptación mejoró la precisión en Humanidades en +6.51 puntos y en Ciencias Sociales en +5.99 puntos respecto al modelo base.
• Robustez Arquitectónica: El método funcionó bien al fusionar modelos de diferentes familias (ej. LLaMA con Qwen) y diferentes tamaños (1B, 8B, 32B). Se observó que modelos fuente más grandes proporcionan representaciones transferibles más ricas.
• Preservación de Capacidades Generales: La fusión no degradó significativamente las capacidades generales del modelo objetivo (razonamiento, conocimiento común), manteniendo un rendimiento comparable al modelo base original.
• Análisis de Correspondencia: El análisis de las matrices de transporte mostró que las correspondencias entre modelos heterogéneos son dispersas pero alineadas, donde un pequeño subconjunto de pares de neuronas explica la mayor parte del "masa de transporte", validando la estrategia de selección Top-k.

5. Significado e Impacto

Este trabajo aborda una brecha crítica en la implementación de IA: la necesidad de modelos pequeños y eficientes para lenguas y dominios específicos que carecen de datos masivos.

• Alternativa a la Destilación: Ofrece una alternativa práctica y teóricamente fundamentada a la destilación de conocimiento, que suele ser costosa computacionalmente y requiere grandes cantidades de datos de entrenamiento.
• Democratización de la IA: Facilita la creación de modelos de alto rendimiento para lenguas de bajos recursos y dominios especializados sin necesidad de entrenar modelos desde cero o realizar ajustes finos masivos.
• Eficiencia Computacional: Al evitar el entrenamiento de gradientes completo y utilizar solo un pequeño conjunto de datos para la alineación, reduce drásticamente los requisitos de computación y energía.

En resumen, "Transport and Merge" establece un nuevo paradigma para la transferencia de conocimiento en LLMs, demostrando que es posible fusionar modelos heterogéneos de manera efectiva mediante la alineación de sus espacios de activación, logrando mejoras consistentes en escenarios de recursos limitados.