Resumen Técnico: MobileMoE: Escalado de Mezcla de Expertos en Dispositivo

Enunciado del Problema

Si bien las arquitecturas de Mezcla de Expertos (MoE) se han convertido en el estándar para escalar modelos de lenguaje grandes (LLM) a cientos de miles de millones de parámetros en servidores en la nube, su potencial a escalas inferiores a mil millones para su despliegue en dispositivos sigue siendo en gran medida inexplorado. Los LLM existentes en dispositivos (por ejemplo, MobileLLM, MobileLLM-Pro) son abrumadoramente densos, requiriendo un aumento en la cantidad de parámetros activos y en la computación de inferencia para mejorar la calidad. Esto crea un cuello de botella para la IA en el borde, donde los dispositivos enfrentan restricciones estrictas en memoria (DRAM) y computación (duración de la batería, latencia). Aunque las capacidades de DRAM de los teléfonos inteligentes recientes han aumentado (por ejemplo, 12–16 GB), no se ha establecido una metodología fundamentada para escalar arquitecturas MoE bajo estas restricciones móviles específicas, equilibrando el tamaño total del modelo, los parámetros activos y el costo de inferencia.

Metodología

1. Ley de Escalado de MoE en Dispositivo

Los autores formulan una ley de escalado generalizada de MoE en dispositivo para optimizar conjuntamente la arquitectura bajo restricciones de memoria y computación. La función de pérdida $L$ se modela como una función de los parámetros activos ($N_{act}$), el tamaño de los datos de entrenamiento ($D$), una transformación monótona del número de expertos ($\hat{E}$) y las decisiones arquitectónicas ($x$, que incluyen la granularidad del experto $g$ y los expertos compartidos $s$).

El objetivo de optimización minimiza la pérdida sujeto a:

• Restricciones de Computación: FLOPs de entrenamiento ($F_{train} = 6N_{act}D$) y FLOPs de inferencia por token ($F_{inf} = 2N_{act}$).
• Restricciones de Memoria: Huella de memoria total $M(N_{total}, T)$, que incluye pesos cuantizados ($b_w$-bit) y caché KV ($b_{kv}$-bit) a una longitud de contexto $T$.

2. Diseño Arquitectónico mediante Divide y Vencerás

Utilizando la ley de escalado, los autores identifican un "punto óptimo" ideal para MoE en dispositivo a través de tres estudios de ablación controlados:

• Número de Expertos ($E$): La dispersión moderada es óptima. Aumentar $E$ reduce la pérdida pero con rendimientos decrecientes más allá de $E=8$. Se selecciona $E=8$ como el punto óptimo práctico para dispositivos con restricciones de memoria.
• Granularidad del Experto ($g$): Los expertos de granularidad fina (dividiendo cada experto en $g$ sub-expertos) permiten rutas de enrutamiento más diversas. El estudio encuentra que $g=8$ es la granularidad óptima en términos de computación, ofreciendo una reducción significativa de la pérdida frente a expertos de granularidad gruesa sin incurrir en una sobrecarga excesiva.
• Experto Compartido ($s$): Incorporar un experto compartido y siempre activo (una vía densa) complementa a los especialistas enrutados. La configuración con un experto compartido ($s=\checkmark$) logra una pérdida menor que su contraparte solo enrutada bajo computación fija.

Arquitectura Resultante: La familia MobileMoE derivada consta de tres escalas (S, M, L) con parámetros activos de 0.3B, 0.5B y 0.9B, y parámetros totales de 1.3B, 2.8B y 5.3B, respectivamente. El diseño utiliza $E=8$, $g=8$ (64 expertos enrutados de granularidad fina) y un experto compartido con enrutamiento top-4.

3. Receta de Entrenamiento de Cuatro Etapas

Para materializar las ventajas arquitectónicas, se emplea una tubería de entrenamiento integral en conjuntos de datos de código abierto:

1. Pre-entrenamiento: ~6T de tokens con una mezcla pesada en la web para garantizar una cobertura lingüística amplia y especialización de expertos. Las técnicas incluyen equilibrio sin pérdida auxiliar, pérdida z del enrutador y núcleos MLP agrupados (GMM) para eficiencia.
2. Entrenamiento Intermedio: ~500B de tokens centrados en datos de alta calidad específicos de dominio (matemáticas, código, conocimiento) para extender el contexto a 8k y afinar la especialización de expertos.
3. Ajuste Fino de Instrucciones (SFT): ~80M de muestras en diversos dominios utilizando despacho de tokens sin pérdida para preservar las señales de instrucción estructuradas.
4. Entrenamiento Consciente de la Cuantización (QAT): Cuantización de pesos INT4 (tamaño de grupo 32) con activaciones dinámicas INT8. Crucialmente, los pesos del enrutador se mantienen en FP32 para mantener la estabilidad del enrutamiento, añadiendo una sobrecarga de memoria negligible.

4. Inferencia Eficiente en Dispositivo

Para cerrar la brecha hacia el despliegue, los autores implementan un núcleo fusionado personalizado de MoE en ExecuTorch. Este núcleo:

• Convierte el despacho de expertos disperso en GEMM agrupados densos reordenando los tokens.
• Fusiona todas las sub-operaciones (enrutamiento, despacho, proyecciones, activación, dispersión) en un único operador para amortizar el lanzamiento del núcleo y la sobrecarga de cuantización.
• Habilita una inferencia eficiente en CPUs de teléfonos inteligentes comerciales (Snapdragon 8 Elite, Apple A18 Pro) y GPUs.

Resultados Clave

Rendimiento en Puntos de Referencia

En 14 puntos de referencia fundamentales (sentido común, conocimiento, ciencia, lectura, razonamiento), MobileMoE establece una nueva frontera de Pareto:

• Eficiencia vs. Modelos Densos: MobileMoE-S/M igualan o superan a los principales LLM densos en dispositivo (por ejemplo, MobileLLM-Pro, Llama 3.2 1B) utilizando 2–4× menos FLOPs de inferencia con huellas de memoria comparables.
• Eficiencia vs. Modelos MoE: MobileMoE-M iguala la precisión del estado del arte OLMoE-1B-7B con ~60% menos parámetros activos y totales. MobileMoE-L supera a OLMoE-1B-7B con 30% menos parámetros activos y una huella de memoria un 23% menor.
• Eficiencia de Tokens: MobileMoE-L logra un rendimiento superior con solo ~6T de tokens de pre-entrenamiento, superando a las líneas base densas entrenadas con 1.5–2× más tokens (por ejemplo, Llama 3.2 1B en 9T).

Perfilado de Tiempo de Ejecución en Dispositivo

Desplegado en Samsung Galaxy S25 e iPhone 16 Pro con pesos INT4:

• Latencia: MobileMoE-S logra una prefill 1.8–3.8× más rápida y una decodificación 2.2–3.4× más rápida en comparación con la línea base densa MobileLLM-Pro con memoria de pesos INT4 comparable.
• Memoria: A 8k de contexto, MobileMoE-S consume hasta un 22% menos de RAM pico que la línea base densa. MobileMoE-L cabe dentro de 5 GB de DRAM a 8k de contexto, confirmando la viabilidad para teléfonos inteligentes modernos.
• Precisión: MobileMoE-L ofrece una precisión sustancialmente mayor (Promedio 57.8 vs. 45.5 para MobileLLM-Pro) con solo un costo de tiempo de ejecución moderado, que disminuye a medida que aumenta la longitud del contexto.

Significado y Afirmaciones

El artículo afirma presentar la primera familia MoE con menos de mil millones de parámetros activos optimizada específicamente para el despliegue en el borde, derivada de una nueva ley de escalado en dispositivo. Su importancia radica en tres contribuciones principales:

1. Fundamento Teórico: Proporciona un marco fundamentado para el diseño de MoE bajo restricciones conjuntas de memoria y computación, identificando la dispersión moderada ($E=8$), los expertos de granularidad fina ($g=8$) y los expertos compartidos como la configuración óptima para dispositivos móviles.
2. Eficiencia de Entrenamiento: Demuestra que las arquitecturas MoE pueden lograr un rendimiento en dispositivo de vanguardia con significativamente menos tokens de entrenamiento (~6T) en comparación con sus contrapartes densas, validando la eficiencia de aprendizaje de MoE a escalas inferiores a mil millones.
3. Despliegue Práctico: Cierra el "último kilómetro" hacia el despliegue móvil proporcionando la primera inferencia eficiente de MoE en teléfonos inteligentes comerciales. Los resultados confirman que MoE no es solo una arquitectura a escala de nube, sino una vía práctica y eficiente para la próxima generación de IA en el borde, habilitando aplicaciones de baja latencia, que preservan la privacidad y son rentables en teléfonos inteligentes y dispositivos portátiles.

Los autores concluyen que, si bien los modelos existentes en dispositivo siguen siendo predominantemente densos, MobileMoE demuestra que MoE ofrece una alternativa superior para el borde, estableciendo una nueva frontera de Pareto que equilibra precisión, memoria y latencia.