Resumen Técnico: Enrutamiento de Índice Inverso Adaptativo para Mezclas de Expertos Granulares (AIR-MoE)

1. Enunciado del Problema

Los modelos de Mezcla de Expertos (MoE) se han convertido en un pilar fundamental de los Modelos de Lenguaje Grandes (LLM) escalables al activar solo un subconjunto de expertos por token. Las leyes de escalado empíricas recientes sugieren que el rendimiento mejora con expertos granulares: utilizar muchos expertos pequeños en lugar de unos pocos grandes. Sin embargo, este régimen introduce un cuello de botella computacional significativo: el costo de enrutamiento. A medida que aumenta el número de expertos ($E$), el enrutador debe calcular productos internos entre cada representación de token y cada centróide de experto para identificar los $K$ mejores expertos. Esta sobrecarga puede dominar el cómputo total, haciendo prohibitivo el entrenamiento de modelos altamente granulares.

Las soluciones existentes para reducir los costos de enrutamiento suelen imponer restricciones estructurales que limitan la expresividad del modelo:

• MoE Jerárquicos dividen los expertos en grupos fijos y no superpuestos, restringiendo el enrutamiento a dentro de esos grupos.
• Restricciones de Producto-Clave/Estructurales (por ejemplo, He, 2024) dividen los centroides de los expertos en componentes para permitir la recuperación exacta, pero restringen el espacio de las representaciones de expertos admisibles (por ejemplo, forzando una estructura de producto cartesiano).
• Los enfoques de Hash Fijo carecen de capacidad de aprendizaje y se ha demostrado que están fundamentalmente limitados.

El desafío consiste en aproximar el enrutamiento de los $K$ mejores de manera eficiente sin imponer restricciones estructurales a los parámetros de los expertos ni requerir la puntuación completa de todos los expertos.

2. Metodología: AIR-MoE

Los autores proponen el Enrutamiento de Índice Inverso Adaptativo para MoE (AIR-MoE), una arquitectura de enrutamiento inspirada en los índices de archivos inversos utilizados en la recuperación de información y la búsqueda aproximada de vecinos más cercanos. El método aproxima los verdaderos expertos de los $K$ mejores mediante un proceso de recuperación de dos etapas, de grueso a fino, evitando la puntuación completa de todos los expertos.

2.1 Recuperación de Dos Etapas

Para cada token de entrada $h$, el enrutamiento procede en paralelo:

1. Cribado Grueso: El token se mapea al código más cercano $c_g$ en un código aprendido $\mathcal{C} = \{c_1, \dots, c_G\}$ utilizando similitud coseno. Cada código $c_g$ está asociado a una lista corta precalculada $L_g$ de $M$ expertos candidatos. Estas listas cortas son esencialmente "listas de publicación" en un índice inverso, que almacenan los índices de los expertos con los productos internos más altos en relación con el centróide del código.
2. Puntuación Fina: El enrutador calcula productos internos exactos solo entre el token $h$ y los expertos dentro de la lista corta recuperada $L_g$. Los $K$ mejores expertos se seleccionan luego de este conjunto de candidatos reducido.

Esta jerarquía asegura que el conjunto de expertos activos $T$ sea un subconjunto de la lista corta $L$, que a su vez es un subconjunto del conjunto completo de expertos $[E]$.

2.2 Entrenamiento y Optimización

Un desafío central es que la asignación de tokens a códigos no es diferenciable. AIR-MoE aborda esto mediante una estrategia de optimización de dos niveles:

• Entrenamiento Basado en Gradientes: Las representaciones de tokens y los centroides de expertos se actualizan mediante retropropagación estándar desde la pérdida de modelado de lenguaje aguas abajo.
• Entrenamiento del Código Libre de Gradientes: El código $\mathcal{C}$ y sus listas cortas asociadas se actualizan de forma independiente utilizando un algoritmo de k-medias esférica adaptativa. Este algoritmo:
  • Utiliza un Promedio Móvil Exponencial (EMA) para rastrear la distribución de las representaciones de tokens.
  • Opera sobre la esfera unitaria para manejar eficazmente la geometría del coseno de alta dimensión.
  • Incluye un mecanismo de reinicialización para códigos "muertos" (aquellos con recuentos de asignación bajos) para evitar el colapso.
  • Actualiza las listas cortas $L_g$ solo después de los pasos del optimizador, asegurando que el índice permanezca consistente con el estado actual de los centroides de expertos.

Crucialmente, AIR-MoE no utiliza el truco de estimación directa; los gradientes no fluyen a través del paso de cuantización hacia el código.

2.3 Garantías Teóricas

El artículo deriva una cota inferior para la Recuperación de Masa (la masa de probabilidad total de los verdaderos expertos de los $K$ mejores retenida en la lista corta). La cota depende de:

1. El error de cuantización $\epsilon(h) = \|h - c(h)\|$: Un error menor conduce a una mejor retención.
2. La masa de la lista corta $\rho_M(c)$: La masa de probabilidad capturada por los $M$ mejores expertos en relación con el código.
   Teóricamente, AIR-MoE retiene la mayor parte de la masa de enrutamiento cuando los tokens están bien cuantizados y el tamaño de la lista corta $M$ es suficientemente grande.

3. Contribuciones Clave

1. Arquitectura AIR-MoE: Un mecanismo de enrutamiento inspirado en índices inversos que aproxima el enrutamiento de los $K$ mejores para MoE granulares sin imponer restricciones estructurales (por ejemplo, agrupación o restricciones paramétricas) a las representaciones de los expertos.
2. Optimización de Dos Niveles: Una estrategia de entrenamiento que desacopla la adaptación libre de gradientes del código (mediante k-medias esférica adaptativa) del entrenamiento basado en gradientes de los parámetros del modelo, evitando la necesidad de estimación directa.
3. Análisis Teórico: Una cota inferior derivada para la recuperación de masa, proporcionando intuición sobre las condiciones bajo las cuales la aproximación preserva la calidad del enrutamiento.
4. Rendimiento Empírico: Demostración de que AIR-MoE logra una compensación superior entre perplejidad y FLOPs en comparación con los métodos existentes de enrutamiento granular eficiente.

4. Resultados Experimentales

Los autores evaluaron AIR-MoE en WikiText-103, C5 y OpenWebText2 a través de escalas de modelos pequeñas, medianas y grandes (hasta ~0.45B de parámetros) con 65,536 expertos.

• Rendimiento vs. Líneas Base: AIR-MoE superó consistentemente a las líneas base granulares eficientes existentes, incluyendo PEER (He, 2024) y MoE Jerárquico.
  • Logró mejoras de perplejidad de hasta un 10% sobre la línea base PEER con costos de FLOP similares.
  • Mientras que el enrutamiento granular estándar (puntuando todos los expertos) logró la perplejidad más baja, incurrió en costos computacionales prohibitivos (por ejemplo, ~467.5P FLOPs frente a 324.1P para AIR-MoE en WikiText-103 Pequeño).
• Eficiencia: AIR-MoE redujo significativamente los FLOPs de entrenamiento en comparación con el enrutamiento granular exacto, manteniendo una perplejidad competitiva.
• Estudios de Ablación:
  • Eliminar la normalización de centroides o utilizar distancia euclidiana en lugar de k-medias esférica degradó ligeramente el rendimiento.
  • Utilizar un código estático (no adaptativo) resultó en el peor rendimiento, confirmando la necesidad del código adaptativo.
  • El enrutamiento de elección de experto condujo a una alta superposición pero a una severa inanición de expertos (78.6% de expertos muertos), lo que indica que la superposición por sí sola no es una métrica suficiente para la calidad del enrutamiento.
• Uso de Expertos: A diferencia de los enrutadores jerárquicos o PEER, que imponen estructuras rígidas que conducen a artefactos de enrutamiento, AIR-MoE permitió un uso flexible de expertos sin restricciones estructurales.

5. Significado y Afirmaciones

El artículo afirma que AIR-MoE proporciona un mecanismo efectivo para el enrutamiento en arquitecturas de mezcla de expertos dispersa, habilitando específicamente el uso de expertos granulares sin la penalización computacional del enrutamiento completo.

• Sustitución Directa: El método no requiere modificaciones en la arquitectura del modelo subyacente ni en la función de pérdida, sirviendo como un reemplazo directo para los enrutadores estándar.
• Flexibilidad: Al evitar restricciones estructurales en los parámetros de los expertos, AIR-MoE permite que el modelo aprenda dependencias token-experto de manera más flexible que los métodos que requieren agrupaciones fijas o estructuras de producto-clave.
• Escalabilidad: El enfoque escala favorablemente en regímenes de lotes efectivos grandes porque el costo de cribado grueso se amortiza y es independiente del número de tokens enrutados.

Los autores reconocen limitaciones, señalando que, aunque las reducciones de FLOP son significativas, las aceleraciones reales en tiempo de ejecución dependen del soporte de hardware y de implementaciones optimizadas para mitigar el tráfico de memoria y las sobrecargas de indexación inherentes a los MoE granulares. También señalan que la implementación actual no emplea reducción de dimensionalidad, dejando esto como una dirección ortogonal para trabajos futuros.