DLLG: Dynamic Logit-Level Gating of LLM Experts

Resumen Técnico: DLLG: Gating de Nivel de Logit Dinámico para Expertos de LLM

Planteamiento del Problema
La proliferación de Modelos de Lenguaje de Gran Escala (LLM) especializados y optimizados para dominios estrechos (por ejemplo, razonamiento matemático, generación de código) presenta un desafío: cómo integrar expertos entrenados de forma independiente como "cajas negras" en un sistema unificado que aproveche sus fortalezas complementarias sin costos prohibitivos de reentrenamiento. Los enfoques existentes sufren una tensión persistente entre adaptabilidad, robustez y practicidad:

• Los métodos basados en enrutamiento (Routing-based methods) dependen de una selección de expertos dura y temprana. Esto conduce a un compromiso prematuro; si el experto elegido falla en una subtarea específica dentro de una sola respuesta, el error es irrecuperable. Además, muchos enrutadores requieren una calibración exhaustiva específica de cada benchmark.
• El ensamble heurístico (Heuristic Ensembling) (por ejemplo, ponderación por perplejidad, votación de tokens) utiliza señales proxy en el tiempo de inferencia. Estas métricas suelen ser empíricas y no están alineadas con la corrección real de la tarea, lo que conduce a salidas subóptimas.
• La fusión de parámetros (Parameter Merging) (por ejemplo, model souping, aritmética de tareas) fusiona los pesos de forma estática. Esto sacrifica la modularidad y a menudo induce una interferencia destructiva al combinar expertos con direcciones de gradiente conflictivas, lo que resulta en una degradación del rendimiento.
• Las arquitecturas de Mezcla de Expertos (Mixture-of-Experts, MoE) típicamente requieren el entrenamiento conjunto de extremo a extremo de las compuertas (gates) y los expertos, lo que limita su capacidad para utilizar especialistas preentrenados listos para usar.

Metodología: DLLG
Los autores proponen DLLG (Dynamic Logit-Level Gating), un marco que realiza una fusión de logits dinámica a nivel de token utilizando un modelo de gating ligero entrenado con supervisión dispersa a nivel de respuesta.

1. Arquitectura:

   • Expertos Congelados: $K$ LLMs especializados permanecen congelados. En cada paso de decodificación $t$, dado un prompt $x$ y un prefijo $y_{<t}$, cada experto $E_i$ produce un estado oculto $h^{(i)}_t$ y los logits del siguiente token $\ell^{(i)}_t$.
   • Modelo de Gating: Un modelo ligero ($g_\phi$) procesa los estados ocultos concatenados de todos los expertos. Proyecta estas representaciones en un espacio de incrustación compartido (usando proyección de bajo rango) y las procesa a través de un tallo autorregresivo (por ejemplo, Qwen2.5-0.5B-Instruct con adaptadores LoRA).
   • Predicción de Pesos: Una cabeza de predicción genera pesos de fusión por token $w_{t,i}$ mediante una transformación sigmoide. Estos pesos se aplican a los logits de los expertos: $\ell_t = \sum_{i=1}^K w_{t,i} \ell^{(i)}_t$. Los pesos no están restringidos a sumar uno, lo que permite contribuciones suaves y flexibles.

2. Entrenamiento (Supervisión a Nivel de Respuesta):

   • El marco aborda la falta de anotaciones a nivel de token utilizando señales de corrección a nivel de respuesta ($s_i \in \{0, 1\}$) derivadas de verificadores automáticos (por ejemplo, pruebas unitarias para código, verificación de respuestas para matemáticas).
   • Teacher Forcing: Durante el entrenamiento, los expertos generan estados ocultos condicionados por el prefijo de verdad fundamental $y_{<t}$. El modelo de gating predice los pesos $w_{t,i}$ para cada posición de token.
   • Objetivo: El modelo se entrena para minimizar el Error Cuadrático Medio (MSE) entre los pesos predichos y la señal de corrección de respuesta transmitida: $\mathcal{L} = \frac{1}{T} \sum_{t=1}^T \sum_{i=1}^K (w_{t,i} - s_i)^2$. Esto enseña a la compuerta (gate) a alinear la fusión a nivel de token con el éxito general de la trayectoria.

3. Inferencia:

   • DLLG opera de forma autorregresiva. En cada paso, el modelo de gating predice los pesos de fusión basados en los estados ocultos actuales de los expertos, los logits se combinan y se selecciona el siguiente token. Esto permite una utilización de expertos recuperable; si un experto comete un error temprano, el modelo de gating puede cambiar diná el peso hacia un experto más competente en los pasos subsiguientes.

Contribuciones Clave

1. Propuesta de Marco: Introducción de DLLG, un marco de ensamble de nivel de logit dinámico que une la generación a nivel de token con la supervisión dispersa a nivel de respuesta, permitiendo una utilización de expertos de grano fino sin necesidad de anotación a nivel de token o aprendizaje por refuerzo en línea.
2. Análisis de Mecanismo: Demostración de que el mecanismo de gating aprendido adapta dinámicamente las contribuciones de los expertos durante el despliegue (rollout), ajustando los pesos según el contexto evolutivo (por ejemplo, pasando de razonamiento matemático a generación de código dentro de una misma respuesta).
3. Validación Empírica: Evaluación a través de diversos benchmarks de razonamiento y código (GSM8K, MATH, HumanEval, etc.) y escalas de modelos (0.5B, 1.5B), mostrando ganancias consistentes sobre las líneas base de enrutamiento, ensamble heurístico y fusión de parámetros.

Resultados

• Rendimiento: A través de ambas escalas de modelos (0.5B y 1.5B), DLLG alcanza consistentemente el mayor rendimiento promedio en comparación con las fuertes líneas base, incluyendo RouterDC, EmbedLLM, varios métodos de ensamble de logits (GaC, UniTe) y técnicas de fusión de parámetros (Linear, SLERP, Task Arithmetic).
• Robustez: En entornos de baja capacidad (0.5B) donde los expertos individuales son débiles, DLLG supera a los métodos de enrutamiento frágiles y heurísticos, demostrando su capacidad para mitigar errores tempranos mediante la fusión suave.
• Generalización: El método muestra un fuerte rendimiento en benchmarks fuera de dominio (por ejemplo, HumanEval, MBPP, BBH), indicando que las reglas de fusión aprendidas se generalizan más allá de la distribución específica de entrenamiento.
• Ablación: Los experimentos confirman que una arquitectura de gating suficientemente expresiva (que incorpora un LLM auxiliar y proyección de bajo rango) es necesaria para una fusión efectiva a nivel de token, superando a las cabezas MLP simples.

Significancia y Reivindicaciones
El artículo afirma que DLLG ofrece un paradigma robusto y escalable para integrar expertos especializados. Al aprender la fusión a nivel de logit a partir de supervisión débil, logra la adaptabilidad de un MoE sin el costo de reentrenar los expertos o el riesgo de interferencia de parámetros. Los autores enfatizan que DLLG preserva la modularidad de los expertos y permite una utilización de expertos de grano fino y recuperable, superando las limitaciones del compromiso prematuro del enrutamiento y la naturaleza estática de la fusión de parámetros. El trabajo sugiere que aprender la utilización de expertos de grano fino a partir de la supervisión débil es un camino práctico para aprovechar los modelos de lenguaje cada vez más especializados.