Hierarchical Copula-Gumbel-Top-\texorpdfstring{$K$}{K} Routing: Two-Sided Dependence Control for Frozen Mixture-of-Experts at Fixed Per-Token Routing Laws

Resumen Técnico: Enrutamiento Jerárquico de Copula-Gumbel-Top-K

Planteamiento del Problema
En los modelos de Mezcla de Expertos (MoE) dispersos, el mecanismo de enrutamiento estocástico estándar (Gumbel-Top-K) define una "ley de enrutamiento" para cada token: una distribución de probabilidad específica sobre listas ordenadas de expertos y pesos de mezcla, condicionada a los logits del token. Mientras que casi todo el trabajo existente modifica esta ley por token (por ejemplo, mediante el ajuste fino de los logits o pérdidas auxiliares), la distribución conjunta de las elecciones de enrutamiento entre diferentes tokens permanece prácticamente inexplorada. El artículo plantea la siguiente pregunta: dado que la ley de enrutamiento de cada token individual se mantiene exactamente fija (como en un modelo congelado), ¿qué distribuciones conjuntas sobre las elecciones de enrutamiento de diferentes tokens son alcanzables? Específicamente, ¿es posible controlar la dependencia entre tokens para mejorar la coherencia de los expertos o el equilibrio de carga sin alterar el comportamiento marginal de ningún token individual?

Metodología: Hierarchical Copula-Gumbel-Top-K (H-CGA)
El autor propone H-CGA, una capa de dependencia en el lado del enrutamiento que opera sobre el vector de ruido Gumbel antes de la selección Top-K. El método se basa en el teorema de Sklar para separar las distribuciones marginales de las estructuras de dependencia. La construcción procede en cuatro etapas:

1. Agrupación Pre-Enrutamiento: Los tokens se particionan en grupos disjuntos (por ejemplo, ventanas fijas de tokens adyacentes). Un controlador lee las características pre-enrutamiento congeladas (por ejemplo, estados ocultos, entropía de la compuerta) para determinar las fuerzas de acoplamiento.
2. Acoplamiento Positivo Intra-Grupo: Para los tokens dentro de un grupo $g$, las perturbaciones Gumbel para cada coordenada de experto se generan mediante una cópula gaussiana intercambiable. Los tokens comparten una variable latente $\zeta_{ge}$, lo que induce una correlación positiva en sus perturbaciones aleatorias. Esto aumenta la probabilidad de que los tokens relacionados seleccionen los mismos expertos (coherencia).
3. Acoplamiento Antitético Inter-Grupo: Se emparejan pares disjuntos de grupos $(g, g')$. Se introduce una fuerza de oposición ajustable $\alpha_{g,g'} \in [0, 1]$. El latente compartido del segundo grupo se construye como una combinación lineal del latente del primer grupo y un término de ruido independiente. Cuando $\alpha_{g,g'} = 1$, los grupos reciben impulsos compartidos totalmente opuestos; cuando $\alpha_{g,g'} = 0$, son independientes. Esto introduce una dependencia negativa entre grupos para contrarrestar la ráfaga de carga (load burstiness).
4. Preservación de las Marginales: La transformación de los latentes compartidos a las variates Gumbel asegura que, para cualquier token $t$ individual, el vector de perturbaciones $(\gamma_{t1}, \dots, \gamma_{tE})$ siga siendo una distribución Gumbel i.i.d., idéntica al enrutamiento independiente.

Contribuciones Clave

1. Arquitectura H-CGA: Una capa de cópula jerárquica que controla la dependencia entre tokens en ambas direcciones (acoplamiento positivo intra-grupo, oposición inter-grupo ajustable) sin modificar los logits del enrutador, los pesos de los expertos ni ninguna ley de enrutamiento por token.
2. Invarianza Teórica: El autor demuestra (Teorema 1) que toda la jerarquía preserva la distribución condicional de la lista Top-K ordenada y los pesos de mezcla de cada token. En consecuencia, los conteos esperados condicionales de inclusión de expertos (Corolario 1) son idénticos a los del enrutamiento independiente.
3. Caracterización del Compromiso Coherencia–Dispersión: El artículo caracteriza el compromiso de la varianza (Proposición 1):
   • El acoplamiento positivo intra-grupo aumenta la varianza de las cargas de expertos realizadas en relación con el enrutamiento independiente (un costo de la coherencia local).
   • La oposición inter-grupo no negativa reduce esta varianza en relación con el acoplamiento "plano" (grupos independientes) al mismo nivel de fuerza intra-grupo, aunque no necesariamente reduce la varianza por debajo de la línea base del enrutamiento independiente.
4. Adaptación Solo de Enrutamiento: El autor formula un método de adaptación donde un pequeño controlador establece los diales de dependencia ($\rho_g, \alpha_{g,g'}$) basándose en características congeladas. Este controlador se entrena utilizando un estimador de función de puntuación (REINFORCE), requiriendo solo pasadas hacia adelante a través del modelo base congelado y gradientes confinados al controlador.

Resultados y Estudio Piloto
El artículo presenta un pequeño estudio piloto en un decodificador de 15.8M de parámetros entrenado con TinyStories, donde el modelo base permanece congelado durante la evaluación.

• Verificaciones de Invarianza: Pruebas sintéticas confirmaron que la frecuencia empírica de las listas Top-K ordenadas y las probabilidades de inclusión de expertos bajo H-CGA coincidieron con el enrutamiento independiente dentro de los márgenes de error de Monte Carlo, validando la invarianza teórica.
• Efectos Estadísticos:
  • El aumento del acoplamiento intra-grupo ($\rho$) aumentó significativamente la similitud de Jaccard de los conjuntos de expertos dentro de las ventanas y redujo el número de expertos distintos utilizados por ventana.
  • El aumento de la oposición inter-grupo ($\alpha$) redujo significativamente la similitud de Jaccard en los límites de los grupos, confirmando la capacidad de desacoplar las demandas de los grupos.
• Aprendizaje: Un controlador escalar aprendido (3 parámetros) se entrenó de forma estable pero no mostró mejoras en la entropía cruzada de validación en comparación con la línea base independiente. Las fuerzas de acoplamiento aprendidas fueron bajas, lo cual es consistente con la observación del artículo de que la preservación de las marginales limita la señal disponible para la optimización cuando solo se utiliza una pérdida por token.
• Varianza de Carga: El piloto observó cambios mínimos en el coeficiente de variación (CV) de la carga agregada, señalando que esta métrica de resumen no prueba directamente los límites de la varianza condicional propuestos en la Proposición 1.

Significancia y Reivindicaciones
El artículo afirma que H-CGA expone un nuevo grado de libertad en el enrutamiento de MoE: la estructura de dependencia conjunta del ruido de enrutamiento. Su principal significancia es teórica y mecánica:

• Seguridad: Proporciona una "garantía de seguridad" de que la modificación de la dependencia entre tokens no altera el comportamiento de enrutamiento marginal de ningún token, asegurando que las preferencias aprendidas del modelo base no sean tocadas.
• Control: Ofrece dos "diales" complementarios para intercambiar la coherencia del conjunto de expertos local frente a la dispersión de la carga global.
• Eficiencia: Permite la adaptación solo de enrutamiento donde el modelo base se evalúa solo en la dirección hacia adelante, evitando el costo computacional de actualizar los pesos de los expertos.

Limitaciones y Modestia
El autor es explícito sobre el alcance y las limitaciones:

• Garantía Local de la Capa: Los resultados de invarianza son condicionales a los estados ocultos y los logits que entran en una capa de enrutamiento específica. No garantizan la invarianza de extremo a extremo en una red de múltiples capas, ya que el acoplamiento puede alterar los estados ocultos posteriores.
• Sin Ganancias a Nivel de Tarea: El piloto no establece que la adaptación solo de enrutamiento mejore el rendimiento en tareas de descenso. El autor señala que la preservación de las marginales implica que los objetivos que dependen únicamente de las distribuciones de un solo token no pueden identificar al controlador; las señales útiles deben surgir de efectos conjuntos (por ejemplo, interacciones de atención), que pueden ser débiles.
• Varianza vs. Desbordamiento: Aunque el método caracteriza teóricamente la varianza de la carga, el piloto no midió directamente la varianza de la carga condicional o los desbordamientos de capacidad. El artículo no pretende resolver el desbordamiento de capacidad, sino que sugiere que el dial de oposición es un mecanismo para abordarlo.
• Sobrecarga: El método requiere generar y coordinar perturbaciones Gumbel para todos los expertos, lo que puede introducir una sobrecarga computacional en comparación con el enrutamiento estándar.

En resumen, el artículo establece un marco matemáticamente riguroso para controlar la dependencia de enrutamiento entre tokens en MoEs congelados, demostando que la coherencia y la dispersión de la carga pueden manipularse sin alterar las leyes de enrutamiento por token, manteniendo la cautela sobre los beneficios prácticos inmediatos para el rendimiento de las tareas.