Por Canuto  

Un informe técnico de Moonshot AI sostiene que Muon puede igualar el rendimiento de AdamW con cerca de la mitad de los FLOPs, mientras su modelo Moonlight mejora resultados en código, matemáticas y comprensión de lenguaje.
***

  • Muon incorporó weight decay y ajustes en la escala de actualización para alcanzar estabilidad en modelos grandes.
  • Las pruebas de escalamiento indican que Muon necesita aproximadamente 52% de los FLOPs de AdamW para lograr un rendimiento comparable.
  • Moonlight, un modelo Mixture-of-Experts de 3B/16B parámetros, fue entrenado con 5,7 billones de tokens y superó a varios modelos comparables.

Muon busca reemplazar a AdamW en el entrenamiento de modelos grandes

El optimizador Muon, basado en la ortogonalización de matrices, logró escalar desde experimentos pequeños hasta el entrenamiento de modelos de lenguaje con miles de millones de parámetros. El avance aparece en el informe técnico Muon is Scalable for LLM Training, elaborado por Jingyuan Liu, Jianlin Su, Xingcheng Yao, Zhejun Jiang y otros investigadores de Moonshot AI, con participación de Xingcheng Yao desde UCLA.

AdamW se mantiene como una de las herramientas estándar para entrenar grandes modelos de lenguaje. Su popularidad se debe a su estabilidad, disponibilidad y capacidad para gestionar parámetros de arquitecturas complejas.

Muon propone una ruta diferente para actualizar las matrices de pesos. En lugar de tratar cada elemento de forma independiente, utiliza el momento del gradiente y aplica una ortogonalización aproximada mediante iteraciones de Newton-Schulz.

La ortogonalización busca evitar que el aprendizaje se concentre en unas pocas direcciones dominantes. Según los investigadores, este comportamiento puede producir actualizaciones más diversas y adecuadas para los pesos que funcionan como operadores sobre espacios de representación.

El informe concluye que Muon puede reemplazar a AdamW como optimizador principal en entrenamientos de gran escala. La afirmación se apoya en experimentos de leyes de escalamiento, pruebas de preentrenamiento y evaluaciones de ajuste supervisado.

Dos cambios permitieron llevar Muon a modelos de mayor tamaño

Los investigadores identificaron dos obstáculos principales cuando intentaron ampliar Muon. El primero fue el crecimiento excesivo de los pesos y de la raíz cuadrática media de las salidas de las capas durante entrenamientos prolongados.

Para corregir ese problema, incorporaron weight decay al optimizador. El mecanismo reduce gradualmente el tamaño de los pesos y evitó que algunas matrices superaran el rango de precisión de bfloat16.

Las pruebas con un modelo de 800 millones de parámetros y 100.000 millones de tokens mostraron una dinámica relevante. Muon sin weight decay avanzó más rápido al principio, pero perdió ventaja durante el sobreentrenamiento.

Muon con weight decay terminó con una pérdida de validación menor que la versión original y que AdamW. La observación llevó al equipo a incluir una tasa de weight decay de 0,1 en las etapas de entrenamiento de Moonlight.

El segundo cambio abordó la variación de la escala de actualización entre matrices de distintas dimensiones. El análisis del informe indica que el RMS teórico de una actualización completa de Muon equivale a la raíz de uno dividido por la dimensión mayor de la matriz.

Esta diferencia puede producir actualizaciones demasiado pequeñas en matrices grandes, como las de los bloques MLP. También puede generar actualizaciones excesivas en estructuras pequeñas, como ciertas cabezas de clave y valor.

Para compensar el efecto, el equipo escaló cada actualización con la raíz cuadrada de la dimensión mayor de la matriz. Además, multiplicó el resultado por 0,2 para aproximar el RMS habitual de AdamW.

El ajuste permitió reutilizar los hiperparámetros de aprendizaje y weight decay utilizados por AdamW. Los investigadores fijaron el momento en 0,95 y utilizaron cinco iteraciones de Newton-Schulz, ya que diez iteraciones ofrecían mayor precisión sin mejorar de forma consistente el rendimiento.

Las pruebas reportan una ventaja de eficiencia computacional

Moonshot AI comparó Muon y AdamW mediante una serie de modelos densos con arquitectura similar a Llama. El estudio utilizó distintos tamaños, desde 399 millones hasta 1.500 millones de parámetros sin contar los componentes de embedding y salida.

Los modelos se entrenaron con cantidades de tokens ajustadas al presupuesto computacional. Para AdamW, el equipo realizó una búsqueda de hiperparámetros con el objetivo de establecer una línea base competitiva bajo un régimen de entrenamiento óptimo.

La curva de escalamiento ajustada registró una pérdida de lenguaje de 2,506 multiplicada por el cómputo elevado a -0,052 para Muon. En AdamW, la relación fue de 2,608 multiplicada por el cómputo elevado a -0,054.

El resultado principal indica que Muon necesita aproximadamente 52% de los FLOPs de AdamW para alcanzar un rendimiento comparable. En términos prácticos, el informe presenta la diferencia como una eficiencia computacional cercana al doble.

La ventaja no depende únicamente de reducir la memoria del optimizador. Muon utiliza un solo búfer de momento, mientras AdamW emplea dos estados de momento, por lo que la alternativa requiere aproximadamente la mitad de memoria para ese componente.

El equipo también diseñó una implementación distribuida basada en ZeRO-1. El sistema particiona los estados del optimizador entre los dispositivos y reúne temporalmente los gradientes necesarios para calcular la actualización completa de cada matriz.

La implementación agrega operaciones de comunicación frente a AdamW, pero los investigadores estiman que la carga total de comunicación se ubica entre una y 1,25 veces la de su equivalente distribuido. En sus pruebas, el costo adicional se acercó al límite inferior cuando utilizaron múltiples grupos de paralelismo de datos.

El proceso de Newton-Schulz se ejecuta en bfloat16 para reducir el tráfico de datos. Además, técnicas como la superposición de comunicaciones y cálculos pueden disminuir el impacto de la latencia.

Moonshot AI liberó su implementación distribuida de Muon, descrita como eficiente en memoria y comunicación. También anunció la disponibilidad de modelos previamente entrenados, versiones ajustadas a instrucciones y checkpoints intermedios.

Moonlight muestra mejoras en código, matemáticas y lenguaje

Para validar Muon en una arquitectura moderna, el equipo entrenó Moonlight, un modelo Mixture-of-Experts inspirado en DeepSeek-V3-Small. El sistema tiene 2.240 millones de parámetros activos y 15.290 millones de parámetros totales, o cerca de 3.000 millones y 16.000 millones al incluir embedding.

Moonlight recibió entrenamiento con 5,7 billones de tokens y una longitud máxima de contexto de 8.000 tokens. El proceso tuvo una etapa inicial de aumento de la tasa de aprendizaje, una fase extensa de descenso coseno y un periodo final de enfriamiento.

Durante el enfriamiento, el modelo utilizó datos de mayor calidad centrados en matemáticas, código y razonamiento. Esta etapa llevó la tasa de aprendizaje a 0,0001 antes de iniciar un descenso lineal hasta cero durante los últimos 500.000 millones de tokens.

A 1,2 billones de tokens, Moonlight superó a su versión equivalente entrenada con AdamW en varias pruebas. En MMLU obtuvo 60,4 frente a 60,2, mientras que en MMLU-Pro consiguió 28,1 frente a 26,8.

La diferencia fue mayor en generación de código. Moonlight registró 37,2 en HumanEval frente a 29,3 de la versión con AdamW, y alcanzó 52,9 en MBPP frente a 49,2.

También obtuvo mejores resultados en razonamiento matemático. En GSM8K marcó 45,0 frente a 43,8, mientras que en MATH consiguió 19,8 frente a 16,1.

Después de completar los 5,7 billones de tokens, Moonlight logró 70,0 en MMLU y 42,4 en MMLU-Pro. En BBH alcanzó 65,2, mientras que en TriviaQA registró 66,3.

En código, el modelo obtuvo 48,1 en HumanEval y 63,8 en MBPP. En matemáticas registró 77,4 en GSM8K y 45,3 en MATH.

Moonlight también consiguió 81,1 en CMATH, 77,2 en C-Eval y 78,2 en CMMLU. El informe señala que el modelo superó a alternativas de arquitectura y presupuesto de entrenamiento similares.

Frente a Qwen2.5-3B, que fue entrenado con 18 billones de tokens, Moonlight logró una puntuación superior en MMLU, MMLU-Pro, BBH, TriviaQA, HumanEval y MBPP. Sin embargo, Qwen mantuvo ventaja en GSM8K, CMATH, C-Eval y CMMLU.

La ortogonalización abre preguntas sobre el futuro de la optimización

El análisis espectral aportó una explicación posible para el desempeño de Muon. Los investigadores midieron la entropía de los valores singulares en matrices de atención, expertos, expertos compartidos, routers y capas densas.

En todos los grupos, las matrices entrenadas con Muon mostraron una entropía SVD superior a las matrices optimizadas con AdamW. La diferencia fue especialmente marcada en los pesos del router que seleccionan expertos.

Más del 90% de las matrices examinadas presentó una entropía SVD mayor con Muon. Una distribución más plana de los valores singulares sugiere que el optimizador explora un conjunto más diverso de direcciones.

El trabajo también analizó el ajuste supervisado. Un modelo preentrenado y ajustado con Muon superó a las variantes que combinaron AdamW en una de las dos etapas.

En el modelo Moonlight de 1,2 billones de tokens, el ajuste con Muon produjo 55,7 en MMLU, 57,3 en HumanEval, 55,6 en MBPP y 68,0 en GSM8K. La versión con AdamW en ambas fases obtuvo 52,0, 53,1, 55,2 y 64,6, respectivamente.

Los resultados también revelaron que cambiar de optimizador entre el preentrenamiento y el ajuste puede reducir la ventaja. En particular, un modelo preentrenado con AdamW no obtuvo una mejora clara al recibir ajuste con Muon.

Una prueba adicional con Qwen2.5-7B mostró un desempeño similar entre ambos optimizadores durante el ajuste supervisado. AdamW alcanzó 71,4 en MMLU y 89,8 en GSM8K, frente a 70,8 y 85,8 de Muon.

El informe plantea que la integración de todos los parámetros dentro de Muon sigue pendiente. Actualmente, ciertos componentes, como RMSNorm, la cabeza de lenguaje y los embeddings, continúan utilizando AdamW.

Los autores también proponen extender el método a normas de Schatten diferentes de la norma espectral. Esa línea podría permitir nuevas formas de controlar las actualizaciones en modelos con arquitecturas más heterogéneas.

La investigación no elimina todos los desafíos técnicos, pero presenta evidencia de que la elección del optimizador puede cambiar de manera sustancial el costo de desarrollar modelos de IA. Con Moonlight, Moonshot AI busca demostrar que una menor factura computacional puede coexistir con mejores resultados en tareas exigentes.

Imagen original de DiarioBitcoin, creada con inteligencia artificial, de uso libre, licenciada bajo Dominio Público.

Este artículo fue escrito por un redactor de contenido de IA y revisado por un editor humano para garantizar calidad y precisión.


ADVERTENCIA: DiarioBitcoin ofrece contenido informativo y educativo sobre diversos temas, incluyendo criptomonedas, IA, tecnología y regulaciones. No brindamos asesoramiento financiero. Las inversiones en criptoactivos son de alto riesgo y pueden no ser adecuadas para todos. Investigue, consulte a un experto y verifique la legislación aplicable antes de invertir. Podría perder todo su capital.

Suscríbete a nuestro boletín

Artículos Relacionados