Optimal Sparsity of Mixture-of-Experts Language Models for Reasoning Tasks

Aquí presento un resumen técnico detallado del artículo "Optimal Sparsity of Mixture-of-Experts Language Models for Reasoning Tasks", publicado como ponencia en ICLR 2026.

1. Problema y Motivación

Las leyes de escalado empíricas han impulsado la evolución de los Modelos de Lenguaje de Gran Escala (LLM), vinculando la pérdida de entrenamiento con el tamaño del modelo, el tamaño del conjunto de datos y el presupuesto de cómputo. Sin embargo, estos coeficientes no son universales y cambian cuando se altera la arquitectura o el pipeline de datos.

El paradigma Mixture-of-Experts (MoE) se ha convertido en un estándar para la escalabilidad económica, permitiendo alta capacidad con un costo de cómputo fijo mediante el enrutamiento disperso de tokens. A pesar de su popularidad, la comunidad científica carece de una comprensión clara sobre cómo la dispersión (sparsity) de los MoE afecta diferencialmente a dos capacidades críticas:

1. Habilidades de memorización: Tareas como preguntas de conocimiento general (TriviaQA) o comprensión de texto (HellaSwag).
2. Habilidades de razonamiento: Tareas matemáticas y lógicas complejas (GSM8K, GSM-Plus).

El problema central es que las leyes de escalado tradicionales asumen que reducir la pérdida de pre-entrenamiento siempre mejora el rendimiento downstream. Este trabajo cuestiona si esta relación se mantiene para tareas de razonamiento en arquitecturas MoE, especialmente bajo presupuestos de cómputo fijos, y si técnicas posteriores como el aprendizaje por refuerzo (GRPO) o el cómputo en tiempo de prueba (TTC) pueden mitigar los efectos negativos de una dispersión subóptima.

2. Metodología

Los autores realizaron un estudio a gran escala utilizando una arquitectura basada en Mixtral (Transformers con MoE, SwiGLU, RMSNorm y enrutamiento top-k sin caída de tokens).

• Configuración Experimental:

  • Se entrenaron familias de modelos variando sistemáticamente tres hiperparámetros arquitectónicos:
    • Ancho del modelo ($d \in \{512, 1024, 2048\}$).
    • Número de expertos por capa ($E \in \{8, 16, 32, 64, 128, 256\}$).
    • Número de expertos activos por token ($k \in \{2, 4, 8, 16\}$).
  • Presupuesto Fijo: Se mantuvo constante el número total de tokens de entrenamiento (125B) y el presupuesto de FLOPs por token (en experimentos de IsoFLOP).
  • Dispersión: Definida como $1 - (k/E)$.

• Evaluación:

  • Se midió la pérdida de pre-entrenamiento, la pérdida de tarea (task loss) y la precisión en benchmarks específicos.
  • Benchmarks de Memorización: TriviaQA, HellaSwag.
  • Benchmarks de Razonamiento: GSM8K, GSM-Plus.
  • Benchmarks de Código: HumanEval, MBPP.
  • Post-entrenamiento y TTC: Se aplicó GRPO (Group Relative Policy Optimization) y Test-Time Compute (Self-Consistency) para verificar si podían recuperar el rendimiento en modelos con dispersión subóptima.

3. Contribuciones Clave y Hallazgos Principales

El estudio revela dos principios fundamentales que desafían la visión clásica de las leyes de escalado para modelos MoE:

A. La FLOPs Activos son Determinantes para el Razonamiento

• Hallazgo: Dos modelos con la misma pérdida de entrenamiento pueden tener rendimientos de razonamiento muy diferentes si tienen diferentes números de FLOPs activos (determinados por $k$).
• Evidencia: Incluso manteniendo constante la pérdida de entrenamiento, los modelos con un valor de top-k mayor (más expertos activos) superan consistentemente a aquellos con top-k menor en tareas de razonamiento.
• Implicación: La calidad del razonamiento no depende solo de la pérdida de pre-entrenamiento, sino de la cantidad de cómputo activo disponible tanto en entrenamiento como en inferencia.

B. Tokens por Parámetro (TPP) y la Curva en "U" Invertida

• Memorización (Hambre de Parámetros): Las tareas de memorización mejoran monótonamente a medida que aumenta el número total de parámetros (es decir, a medida que aumenta la dispersión y disminuye el TPP).
• Razonamiento (Hambre de Datos): Las tareas de razonamiento muestran un comportamiento no monótono. La precisión alcanza un pico óptimo alrededor de 20 tokens por parámetro (TPP).
  • Si el TPP es demasiado bajo (demasiados parámetros para pocos datos), el modelo sufre de "hambre de datos" y el rendimiento decae.
  • Si el TPP es demasiado alto (pocos parámetros para muchos datos), el modelo está sub-parametrizado.
• Resultado Sorprendente: Aumentar el número de expertos (aumentando la dispersión) más allá de un umbral óptimo, manteniendo el presupuesto de cómputo fijo, deteriora el rendimiento en matemáticas, a pesar de que la pérdida de entrenamiento sigue disminuyendo. Esto crea una relación en forma de "U" invertida entre la pérdida de entrenamiento y la precisión en razonamiento.

C. Ineficacia de Post-entrenamiento y TTC

• Se investigó si técnicas avanzadas como GRPO (Refuerzo) o Test-Time Compute (mayor tiempo de inferencia) podían corregir el deterioro del razonamiento causado por una dispersión excesiva.
• Resultado: Ni el GRPO ni el aumento del cómputo en tiempo de prueba eliminan la brecha de rendimiento. La brecha entre habilidades de memorización y razonamiento causada por la dispersión subóptima persiste.
• Conclusión: La dispersión óptima debe determinarse durante la fase de pre-entrenamiento; no es posible "arreglar" una arquitectura subóptima posteriormente con estos métodos.

4. Resultados Detallados

• Pérdida de Tarea vs. Pérdida de Entrenamiento: Mientras que en TriviaQA y HellaSwag la pérdida de tarea disminuye monótonamente con la pérdida de entrenamiento, en GSM8K y GSM-Plus, una vez que la pérdida de entrenamiento cae por debajo de cierto umbral, la pérdida de tarea aumenta (el modelo se vuelve peor en la tarea específica).
• Análisis de IsoFLOP: Bajo un presupuesto de FLOPs fijo, para tareas de razonamiento, los modelos más densos (menos dispersos) superan a los modelos muy dispersos cuando el presupuesto de cómputo es alto. La dispersión óptima depende del nivel de cómputo disponible.
• Tareas de Código: Los resultados en HumanEval y MBPP replican los hallazgos de las tareas matemáticas, confirmando que el razonamiento procedimental también requiere un equilibrio específico entre densidad y datos, favoreciendo configuraciones más densas en regímenes de alto cómputo.

5. Significancia e Impacto

Este trabajo redefine la práctica de escalado para modelos LLM basados en MoE:

1. Revisión de las Leyes de Escalado: Sugiere que la ley de Chinchilla (que optimiza la relación parámetros/datos) debe ser reevaluada para MoE, diferenciando entre tareas de "hambre de parámetros" (memorización) y "hambre de datos" (razonamiento).
2. Guía de Diseño Arquitectónico: Para entrenar modelos de razonamiento bajo un presupuesto fijo, no basta con maximizar el número total de parámetros. Es crucial encontrar el equilibrio correcto entre:
   • FLOPs Activos: Asegurar que el número de expertos activos ($k$) sea suficiente para el cómputo requerido.
   • TPP Óptimo: Mantener una relación de aproximadamente 20 tokens por parámetro para tareas de razonamiento, evitando la sobre-dispersión que lleva al modelo a ser "hambriento de datos".
3. Limitación de Técnicas Post-entrenamiento: Demuestra que el pre-entrenamiento es la fase crítica para establecer la capacidad de razonamiento; las técnicas de ajuste fino o inferencia avanzada no pueden compensar una arquitectura mal configurada en términos de dispersión.

En resumen, el artículo argumenta que la dispersión óptima de MoE no es un valor fijo, sino que debe determinarse conjuntamente por los FLOPs activos y la relación Tokens/Parámetro, requiriendo configuraciones más densas para el razonamiento en regímenes de alto cómputo, a diferencia de la intuición previa de que mayor dispersión siempre es mejor para la eficiencia.