InfiGFusion: Graph-on-Logits Distillation via Efficient Gromov-Wasserstein for Model Fusion

Resumen Técnico: InfiGFusion

Enunciado del Problema
Los esfuerzos recientes para fusionar Modelos de Lenguaje Grandes (LLM) heterogéneos en sistemas unificados han dependido en gran medida de métodos de destilación basados en logits. Si bien estos enfoques mantienen la eficiencia de inferencia, típicamente tratan las dimensiones del vocabulario de forma independiente, alineando las distribuciones de tokens mediante objetivos como la divergencia KL o la distancia de Wasserstein. Esta alineación token a token pasa por alto las dependencias semánticas—las interacciones entre dimensiones y los patrones de co-activación que codifican los procesos de razonamiento interno de un modelo. En consecuencia, los métodos existentes luchan por alinear modelos con comportamientos de generación divergentes, particularmente en tareas que requieren razonamiento multi-paso, inferencia relacional y alineación semántica fina.

Metodología: InfiGFusion
Los autores proponen InfiGFusion, un marco de fusión consciente de la estructura que modela explícitamente estas dependencias semánticas. La innovación central es la función de pérdida de Destilación de Grafos sobre Logits (GLD), que opera de la siguiente manera:

1. Construcción de Grafos de Logits: En lugar de tratar los logits como vectores independientes, el método construye un grafo global de co-activación. Para cada salida, se retienen los logits superiores a $k$ para garantizar la dispersión. Los productos externos de estos logits a través de las posiciones de la secuencia se agregan para formar un grafo donde los nodos representan canales de vocabulario y las aristas cuantifican sus activaciones conjuntas (similitud).
2. Alineación Consciente de la Estructura: Para alinear los grafos de logits de los modelos fuente con un modelo pivote, el marco emplea la distancia de Gromov-Wasserstein (GW). GW se elige por su capacidad para igualar distribuciones basándose en la similitud estructural interna en lugar de la correspondencia directa de nodos.
3. Aproximación Eficiente: El cálculo estándar de la distancia GW tiene una complejidad prohibitiva de $O(n^4)$, lo que lo hace inviable para vocabularios grandes ($n > 10^5$). Los autores introducen una novedosa aproximación de forma cerrada basada en ordenamiento que reduce la complejidad a $O(n \log n)$.
   • Compresión de Características del Grafo: La matriz de similitud se comprime en características escalares a nivel de nodo (específicamente, medias de filas o características de tipo grado).
   • Emparejamiento Basado en Ordenamiento: La alineación se realiza ordenando estas características y emparejando la entrada $i$-ésima más grande de la fuente con la entrada $i$-ésima más grande del objetivo.
   • Garantías Teóricas: El artículo proporciona un límite de error de aproximación demostrable de $\frac{n-1}{n^2} + \frac{m-1}{m^2}$, que decae como $O(1/n)$. Además, los autores demuestran que la pérdida guiada por GW tiene una constante de Lipschitz significativamente menor en comparación con las pérdidas KL y Wasserstein estándar, lo que sugiere límites de generalización más ajustados y un entrenamiento más estable.
4. Objetivo Unificado: El objetivo final de entrenamiento combina la pérdida GLD consciente de la estructura con una pérdida de Destilación Universal de Logits (ULD) de grano grueso (para la alineación de distribuciones a nivel de token) y señales estándar de Ajuste Fino Supervisado (SFT).

Contribuciones Clave

• Primer Marco de Fusión Consciente de la Estructura: InfiGFusion es el primer método que modela dependencias semánticas mediante grafos de logits a nivel de características y los alinea utilizando una novedosa pérdida GLD, superando la alineación de tokens independiente.
• Aproximación GW Escalable: Los autores proponen una aproximación de forma cerrada para la distancia de Gromov-Wasserstein que reduce la complejidad computacional de $O(n^4)$ a $O(n \log n)$ con límites de error demostrables, haciendo viable la alineación basada en grafos para vocabularios a escala de LLM.
• Superioridad Empírica en Razonamiento: Extensos experimentos en 11 benchmarks demuestran que InfiGFusion supera consistentemente a las líneas base de fusión más avanzadas (por ejemplo, MiniLogit, FuseLLM, FuseChat) y a las líneas base de SFT.

Resultados
El método se evaluó en 11 benchmarks que abarcan razonamiento, codificación, matemáticas y seguimiento de instrucciones, fusionando múltiples modelos fuente (Qwen2.5, Mistral, Qwen2.5-Coder) en un modelo pivote Phi-4.

• Ganancias en Razonamiento Complejo: InfiGFusion mostró mejoras significativas en tareas que requieren inferencia multi-paso y relacional. Notablemente, logró una mejora de +35.6 puntos en Aritmética Multi-paso y una mejora de +37.06 puntos en Juicio Causal sobre las líneas base de SFT.
• Rendimiento General: El modelo superó al modelo pivote y a los modelos fuente en diversos dominios, alcanzando una puntuación promedio de 83.85 en comparación con los 79.54 del pivote.
• Eficiencia: A pesar de la alineación estructural, el método mantuvo una eficiencia competitiva, requiriendo solo 195 horas de GPU para el entrenamiento, significativamente menos que las tuberías de fusión multi-paso como FuseChat (650 horas de GPU).
• Estudios de Ablación: Los experimentos confirmaron que tanto los componentes ULD (alineación de grano grueso) como GLD (alineación estructural) son necesarios, siendo GLD específicamente el impulsor de las ganancias en tareas de razonamiento.

Significado
El artículo afirma que InfiGFusion aborda una brecha crítica en la fusión de modelos: la incapacidad de los objetivos a nivel de token para capturar las estructuras relacionales subyacentes al razonamiento del modelo. Al alinear explícitamente la "forma" del espacio de logits en lugar de solo las distribuciones marginales, el marco permite la integración de LLM heterogéneos sin aumentar los costos de inferencia. Los autores postulan que esta alineación que preserva la estructura es esencial para avanzar en las capacidades colectivas de los LLM, particularmente en escenarios de razonamiento complejo donde la consistencia lógica y la inferencia multi-salto son primordiales. El trabajo sugiere que el progreso futuro en la fusión de modelos depende de ir más allá del emparejamiento superficial de tokens para preservar las dependencias semánticas internas que impulsan un razonamiento coherente.