Resumen Técnico: Enrutamiento Pegajoso para la Inferencia de MoE con Eficiencia de Memoria

1. Planteamiento del Problema: El Cuello de Botella del Intercambio de Expertos

Los modelos de Mezcla de Expertos (MoE) están diseñados para escalar los Modelos de Lenguaje Extensos (LLMs) activando solo un subconjunto disperso de expertos por cada token. Teóricamente, esto permite recuentos masivos de parámetros con una computación por token baja, lo que los hace atractivos para el despliegue en el borde (edge) donde solo unos pocos expertos necesitan residir en la memoria rápida (VRAM/SRAM) en cualquier momento dado.

Sin embargo, el artículo identifica un modo de fallo crítico en la práctica: la inconsistencia de enrutamiento temporal. Los enrutadores MoE estándar toman decisiones independientes para cada token basándose únicamente en la representación oculta de ese token. En consecuencia, los tokens consecutivos activan frecuentemente expertos completamente diferentes. En hardware con restricciones de memoria, cada cambio requiere la expulsión de los pesos de los expertos en caché y la carga de nuevos pesos desde un almacenamiento lento (RAM de CPU/Flash). Dado que el ancho de banda de PCIe/NVMe es órdenes de magnitud inferior al ancho de banda de la memoria GPU, estos frecuentes "intercambios de expertos" dominan la latencia de inferencia, haciendo que los modelos MoE sean a menudo más lentos que los modelos densos de tamaño comparable.

Las soluciones existentes se consideran insuficientes:

• Los enfoques a nivel de sistema (por ejemplo, heurísticas de caché, prelectores/prefetchers) son reactivos; aprovechan cualquier localidad existente, pero no pueden corregir enrutadores intrínsecamente inconsistentes.
• Los enfoques post-hoc (por ejemplo, el ajuste fino solo del enrutador) fallan porque las representaciones de los expertos fueron moldeadas por un régimen de entrenamiento que nunca incentivó la localidad. El enrutador no puede realinear fácilmente las decisiones en un espacio de representación que no fue diseñado para ello.
• Los rediseños arquitectónicos (por ejemplo, Oracle-MoE) suelen requerir cambios significativos en la estructura o en el flujo de entrenamiento del modelo.

2. Metodología: StickyMoE

Los autores proponen StickyMoE, un enfoque de tiempo de entrenamiento que trata la localidad temporal de enrutamiento como un objetivo explícito. La innovación central es una pérdida de consistencia de enrutamiento diferenciable ($L_{cons}$) que penaliza los cambios abruptos en la asignación de expertos entre tokens adyacentes.

La Pérdida de Consistencia de Enrutamiento

La función de pérdida minimiza la distancia $\ell_2$ entre los vectores de probabilidad de la puerta softmax ($g_t$) de tokens consecutivos:
$$L_{cons} = \frac{1}{T-1} \sum_{t=2}^{T} \| g_t - g_{t-1} \|_2^2$$
Donde $g_t \in \mathbb{R}^N$ es la salida del enrutador para el token $t$.

Esta pérdida se integra en el objetivo de entrenamiento total junto con las pérdidas estándar de entropía cruzada ($L_{CE}$) y de equilibrio de carga ($L_{bal}$):
$$L = L_{CE} + \lambda L_{cons} + \mu L_{bal}$$

Mecanismos Clave

• Coadaptación: A diferencia de los métodos post-hoc, StickyMoE entrena simultáneamente el enrutador y las representaciones de los expertos. Minimizar $L_{cons}$ actualiza los pesos del enrutador ($W_r$) y realiza la retropropagación a través de los estados ocultos ($h_t$), fomentando que el modelo aprenda representaciones que sean naturalmente consistentes a través de extensiones semánticamente coherentes.
• Variante Suave-Dura (Soft-Hard): Para abordar la "deriva de largo alcance" (donde pequeños cambios paso a paso se acumulan en un cambio distante), los autores introducen una restricción de anclaje. La secuencia se divide en ventanas; el primer token de una ventana establece un ancla, y los tokens subsiguientes son penalizados por desviarse de esta ancla, con la penalización aumentando linealmente con la distancia desde el inicio de la ventana.
• Compatibilidad: El método no requiere cambios arquitectónicos y añade un gasto computacional insignificante. Es compatible con cualquier MoE top-k estándar y con las técnicas de equilibrio de carga existentes.

3. Contribuciones Clave

1. Pérdida de Consistencia de Enrutamiento: Un objetivo de entrenamiento simple y diferenciable que penaliza directamente el cambio de experto entre tokens adyacentes, requiriendo solo un hiperparámetro ($\lambda$).
2. Complementariedad con el Equilibrio de Carga: La evidencia empírica muestra que la pérdida de consistencia y la pérdida de equilibrio de carga son complementarias. La pérdida de consistencia moldea cuándo se usan los expertos sin causar un "colapso de expertos" (la entropía de utilización permanece alta, cerca de su máximo teórico).
3. Variante Suave-Dura: Una combinación de penalizaciones suaves por paso y restricciones de anclaje por segmento que garantiza una fuerte localidad sin costes adicionales de calidad.
4. Evaluación Exhaustiva: Experimentos controlados comparando StickyMoE contra MoE vanilla, ajuste fino de enrutador post-hoc (ReMoE) y bases arquitectónicas (Oracle-MoE) en modelos pequeños y medianos.

4. Resultados Experimentales

Los experimentos se realizaron en WikiText-2 utilizando modelos MoE pequeños (8.8M de parámetros) y medianos (22M de parámetros).

• Tasa de Cambio de Experto (SR): StickyMoE reduce la tasa de cambio de expertos hasta en un 59% (por ejemplo, de 0.71 a 0.29 en el modelo mediano).
• Tasa de Acierto de Caché (CHR): La reducción en los cambios se traduce en una reducción de 3.92× en los fallos de caché (el CHR aumentó de ~0.54 a ~0.88) en el modelo mediano con un caché de 2 ranuras de expertos.
• Perplejidad (Calidad):
  • Con valores moderados de $\lambda$ (por ejemplo, 0.05), StickyMoE mejora la perplejidad en comparación con el baseline (por ejemplo, -4.1% en el modelo mediano), lo que sugiere que la pérdida de consistencia actúa como un regularizador beneficioso.
  • Los valores agresivos de $\lambda$ introducen un compromiso, pero la frontera de Pareto de StickyMoE domina estrictamente a otros métodos.
• Fracaso del Ajuste Fino Post-Hoc: El baseline ReMoE (ajuste fino solo del enrutador en un modelo preentrenado) no logró reducir significativamente la tasa de cambio (<0.5% de cambio), confirmando que la localidad no puede ser retroajustada una vez que las representaciones están fijas.
• Análisis por Capas: La pérdida de consistencia es más efectiva en las capas más profundas (L1–L3). La capa de entrada (L0) se resiste a las mejoras de localidad, probablemente porque los embeddings de los tokens carecen de mezcla contextual.

5. Significado y Reivindicaciones

El artículo argumenta que el enrutamiento de la localidad temporal es un objetivo de entrenamiento legítimo que debe instilarse durante el preentrenamiento en lugar de tratarse como una consideración secundaria a nivel de sistema.

• Cambio de Paradigma: Así como el entrenamiento consciente de la cuantización se ha convertido en un estándar, los autores postulan que el "entrenamiento consciente de la localidad" debería ser el predeterminado para los modelos MoE destinados al despliegue en el borde.
• Resolución de la Causa Raíz: StickyMoE aborda la causa raíz (decisiones de enrutamiento inconsistentes) en lugar de los síntomas (fallos de caché). Al co-entrenar las representaciones y las decisiones de enrutamiento, evita el desajuste de representación que afecta a los métodos post-hoc.
• Practicidad: El método es agnóstico a la arquitectura y no requiere cambios en el flujo de inferencia. El enrutador entrenado produce naturalmente distribuciones de puerta consistentes, lo que permite que las estrategias de caché estándar (LRU, LFU) funcionen significativamente más eficientemente sin ingeniería adicional.

Los autores concluyen que, si bien las optimizaciones a nivel de sistema y los ajustes post-hoc siguen siendo complementos útiles, no pueden sustituir a un modelo que nunca fue entrenado para ser amigable con la caché en primer lugar. StickyMoE proporciona un camino para realizar la eficiencia de memoria teórica de las arquitecturas MoE en dispositivos de borde.