El optimizador Muon se ha consolidado como una herramienta clave para el entrenamiento de modelos de lenguaje de última generación como Kimi K2 o GLM-5, ya que requiere menos pasos que AdamW para alcanzar una misma pérdida. Sin embargo, cada paso de Muon es más costoso debido a su procedimiento de ortogonalización Newton-Schulz, una operación con matrices de tiempo cúbico. Ante este cuello de botella, investigadores han desarrollado Gram Newton-Schulz, una reformulación matemáticamente idéntica al algoritmo original, pero que opera principalmente sobre la matriz de Gram, pequeña y simétrica, en lugar de sobre la matriz rectangular de entrada. Esto reduce drásticamente las multiplicaciones matriciales rectangulares, que son las más costosas en términos de tiempo de cálculo y uso de memoria.
El nuevo algoritmo consigue reducir el tiempo del paso de ortogonalización entre un 40 % y un 50 % en modelos MoE de billones de parámetros, como Kimi K2. Para lograrlo, se implementaron kernels personalizados en CuTeDSL para las arquitecturas Hopper y Blackwell de NVIDIA, maximizando el rendimiento de las multiplicaciones simétricas. Además, se identifica y corrige una posible inestabilidad numérica en precisión media mediante una estrategia de reinicio de la matriz de Gram, dando lugar a la versión estabilizada denominada Stabilized Gram Newton-Schulz.
El resultado es GramMuon, la versión del optimizador que incorpora Gram Newton-Schulz. Las pruebas confirman que el entrenamiento con GramMuon es estable y conserva la calidad de optimización del Muon estándar, con la ventaja de un tiempo de cómputo significativamente menor. Los autores han liberado implementaciones de código abierto para facilitar la adopción. Dado que la reformulación es matemáticamente equivalente, es compatible con la mayoría de variantes de Muon existentes, como Scion, Dion o SOAP.