🤖 machine learning

Approximate Muon with low-rank adapters

Este artículo presenta sMuon, un método eficiente que aproxima el optimizador Muon para el ajuste fino eficiente de parámetros de bajo rango mediante la linealización del objetivo y su resolución vía mínimos cuadrados, permitiendo así los beneficios de rendimiento de Muon en entornos como SFT y el preentrenamiento de ReLoRA donde anteriormente era incompatible con LoRA.

Autores originales: Ben Anson, Conor Houghton, Edward Milsom

Publicado 2026-08-17

📖 4 min de lectura☕ Lectura para el café

CC0 1.0

Autores originales: Ben Anson, Conor Houghton, Edward Milsom

Artículo original dedicado al dominio público bajo CC0 1.0 (http://creativecommons.org/publicdomain/zero/1.0/). ✨ Esta es una explicación generada por IA del artículo a continuación. No ha sido escrita ni avalada por los autores. Para mayor precisión técnica, consulte el artículo original. Leer descargo de responsabilidad completo

Imagina que estás intentando enseñarle a un robot gigante y superinteligente una nueva habilidad, como escribir poesía o resolver problemas matemáticos. No quieres reentrenar todo el cerebro del robot desde cero porque eso toma una eternidad y cuesta una fortuna. En su lugar, quieres acoplarle un pequeño y ligero "módulo de entrenamiento". Esto se llama Ajuste Fino de Parámetros Eficientes (PEFT). Piensa en ello como ponerle un auricular especializado a un robot de propósito general; el auricular le enseña el nuevo truco sin cambiar el hardware central del robot.

Una forma popular de construir estos auriculares es LoRA (Adaptación de Bajo Rango). Es como construir el auricular a partir de dos hojas de papel finas y planas que, al apilarse, crean la forma de la nueva habilidad. Para enseñar al robot, necesitamos un optimizador: un entrenador matemático que le dice al auricular cómo ajustar su forma para mejorar. El entrenador más común es AdamW, pero hay un entrenador más nuevo y potente llamado Muon. Muon es famoso por ser increíblemente eficiente cuando se realiza el preentrenamiento (enseñando al robot sus primeras lecciones desde cero). Funciona asegurándose de que las actualizaciones del cerebro del robot ocurran de una manera perfectamente equilibrada y "ortogonal", como un bailarín girando perfectamente sobre un pie sin tambalearse.

Sin embargo, hay un problema. Cuando intentas usar Muon para entrenar el pequeño auricular LoRA, las matemáticas fallan. Muon quiere equilibrar toda la capa del cerebro del robot, pero el auricular LoRA es solo una pequeña sección de bajo rango de esa capa. Es como intentar usar la batuta de un director para dirigir a un solo violinista de una manera que tenga sentido para toda la orquesta; la geometría simplemente no encaja. El método de solución estándar ha sido ignorar la orquesta y simplemente decirle al violinista que gire por su cuenta, pero eso pierde la magia de Muon. Este artículo se pregunta: ¿Podemos encontrar una manera de hacer que Muon funcione para estos pequeños auriculares sin romper las matemáticas o ralentizar todo?

Los autores, Ben Anson, Conor Houghton y Edward Milsom, dicen que sí. Introducen un nuevo método llamado sMuon (small Muon). En lugar de intentar forzar lo imposible, utilizan un truco matemático ingenioso: "linealizan" el problema. Imagina que estás intentando caminar por la cuerda floja mientras sostienes una sombrilla gigante y tambaleante. En lugar de intentar equilibrar toda la sombrilla a la vez, pretendes que la sombrilla es solo un palo recto por un breve segundo, calculas el paso perfecto y luego te ajustas. Al hacer esto, encuentran la mejor manera posible de actualizar el auricular LoRA que aproxima la danza perfecta de Muon.

El artículo encuentra que este nuevo método sMuon funciona sorprendentemente bien. Cuando lo probaron en varios modelos de lenguaje, sMuon a menudo funcionó mejor que el entrenador estándar AdamW y otros intentos de Muon. Fue particularmente impresionante cuando se usó en un robot que ya había sido preentrenado con Muon (llamado Moonlight), donde sMuon logró las puntuaciones más altas en seis de once tareas diferentes. En un experimento de preentrenamiento, sMuon empató por los mejores resultados en la reducción de la tasa de error, pero lo hizo aproximadamente un 30% más rápido que su rival más cercano, un método llamado Riemannion.

Crucialmente, los autores demuestran que su método no es solo una idea teórica; también es práctico. Mientras que otros métodos avanzados requieren cálculos pesados y lentos (como desarmar y reensamblar matrices complejas), sMuon se basa casi por completo en multiplicaciones de matrices simples, el tipo de matemáticas para las que los chips de computadora modernos están construidos para trabajar a la velocidad del rayo. Esto significa que sMuon es lo suficientemente rápido como para ser utilizado en el entrenamiento del mundo real sin ralentizar el proceso. El artículo sugiere que, al respetar la geometría de toda la capa en lugar de tratar las partes de LoRA como piezas separadas y no relacionadas, podemos obtener un aumento significativo de rendimiento con muy poco costo adicional. Es un recordatorio de que, a veces, la mejor manera de enseñarle un nuevo truco a un robot gigante es darle un pequeño, inteligente y matemáticamente elegante empujón.

=== RESUMEN ===

Resumen Técnico: Aproximación de Muon con Adaptadores de Bajo Rango (sMuon)

Planteamiento del Problema

El optimizador Muon ha demostrado un rendimiento superior en el preentrenamiento de grandes modelos de lenguaje (LLMs) al ortogonalizar las actualizaciones de peso en cada capa, igualando eficazmente los valores singulares. Este comportamiento corresponde al descenso más pronunciado bajo la norma espectral, controlando el cambio máximo de una capa como un operador lineal en lugar de un vector aplanado. Sin embargo, Muon rara vez se utiliza para el Ajuste Fino Eficiente en Parámetros (PEFT), particularmente con la Adaptación de Bajo Rango (LoRA).

La obstrucción matemática central es que Muon requiere ortogonalizar el paso de descenso para la matriz de peso completa de la capa. En el entorno LoRA, donde la actualización de peso está restringida a una parametrización de bajo rango $\Delta W = BA$ (con $B \in \mathbb{R}^{d_1 \times r}$ y $A \in \mathbb{R}^{r \times d_2}$ ), es matemáticamente imposible encontrar generalmente actualizaciones $\delta B$ y $\delta A$ tales que la $\delta \Delta W$ resultante coincida perfectamente con el gradiente ortogonalizado de rango completo requerido por Muon.

Las soluciones alternativas existentes incluyen:

Muon por factor: Ortogonalizar $B$ y $A$ por separado, lo que ignora la geometría de la capa completa.
LoRA-Muon: Relajar el objetivo utilizando la desigualdad triangular.
Riemannion: Formular el problema como una optimización de Riemann en una variedad de bajo rango, lo que requiere Descomposiciones de Valores Singulares (SVD) truncadas costosas durante el paso de retracción, lo que genera una sobrecarga computacional significativa.

Metodología

Los autores proponen sMuon (small Muon), un método que aproxima la solución a un objetivo de Muon relajado dentro del entorno de bajo rango. El enfoque se basa en la linealización y la optimización de mínimos cuadrados para encontrar el paso de descenso óptimo.

1. Linealización y Formulación de Mínimos Cuadrados

El objetivo es encontrar actualizaciones $\delta B, \delta A$ tales que la actualización implícita $\delta \Delta W \approx \delta B A + B \delta A$ aproxime el paso ideal de Muon $\delta \Delta W^* = -\eta \cdot \text{msign}(G_{\Delta W})$ , donde $\text{msign}(X) = UV^\top$ es el factor polar de la SVD.

Dado que el paso ideal es generalmente de rango completo y no puede realizarse exactamente mediante factores de bajo rango, los autores linealizan el problema descartando los términos de segundo orden ( $\delta B \delta A$ ) y resolviendo un problema de mínimos cuadrados:
$(\delta B^*, \delta A^*) = \arg \min_{\delta B, \delta A} \| \delta \Delta W^* - (\delta B A + B \delta A) \|_F^2$
Esto produce una solución de forma cerrada que involucra inversas de Moore-Penrose ( $B^\dagger, A^\dagger$ ). Crucialmente, los autores demuestran que la actualización óptima no es única en términos de los factores $(\delta B, \delta A)$ , pero el cambio de primer orden resultante en la capa completa $\delta \Delta W$ sí es único.

2. Computación Eficiente Basada Solo en Multiplicación de Matrices (Matmul-Only)

Una implementación directa de la solución de forma cerrada requiere ortogonalizar una matriz grande de $d_1 \times d_2$ , lo cual es computacionalmente prohibitivo. Para abordar esto, los autores derivan una formulación eficiente (Algoritmo 2) que evita la ortogonalización explícita de matrices grandes:

Gradiente de Bajo Rango: En lugar de utilizar el gradiente completo $G_{\Delta W}$ , utilizan una proyección de bajo rango $H = P_B G_{\Delta W} + (I-P_B)G_{\Delta W} P_A$ , que es equivalente de primer orden para el objetivo de optimización.
Reducción de Dimensionalidad: Explotan la identidad $\text{msign}(H) = U \text{msign}(U^\top H V) V^\top$ , donde $U$ y $V$ son bases ortonormales que abarcan los espacios columna y fila de $H$ . Dado que $H$ tiene un rango de como máximo $2r$ , esto reduce la costosa ortogonalización de $d_1 \times d_2$ a una ortogonalización de $2r \times 2r$ .
Implementación: Todo el algoritmo se basa únicamente en multiplicaciones de matrices, raíces cuadradas inversas de matrices y descomposiciones polares de matrices pequeñas de $r \times r$ o $2r \times 2r$ . Esto evita la sobrecarga de las descomposiciones QR y SVD requeridas por Riemannion.

3. Detalles Algorítmicos

El optimizador sMuon incorpora:

Decaimiento de Peso Dividido (Split Weight Decay): Para igualar la dinámica del entrenamiento de parámetros completos, el decaimiento de peso se aplica mediante un factor de escala $s = \sqrt{1-\lambda\eta}$ a los factores $A$ y $B$ antes de la actualización del gradiente.
Transporte de Momentum: Para asegurar la invarianza de la parametrización (es decir, que el optimizador se comporte de manera consistente independientemente de la factorización específica $B, A$ ), los búferes de momentum se proyectan sobre los espacios de los factores actuales en cada paso.

Contribuciones Clave

Derivación de sMuon: Un método de forma cerrada para entrenar adaptadores LoRA con Muon mediante la toma de una aproximación de primer orden del objetivo de descenso más pronunciado regularizado por la norma espectral.
Implementación Eficiente: Un algoritmo de "solo multiplicación de matrices" que es computacionalmente eficiente para el hardware moderno, evitando descomposiciones costosas. Ofrece una velocidad comparable a Per-factor Muon y LoRA-Muon, siendo significamente más rápido que Riemannion, especialmente para rangos LoRA más grandes.
Evaluación Empírica: Experimentos extensos a través de cuatro modelos base (incluyendo uno preentrenado con Muon y tres preentrenados con AdamW) y once benchmarks (sentido común y programación).

Resultados

El artículo reporta los siguientes hallazgos:

Rendimiento: sMuon se desempeña favorablemente en los benchmarks de SFT. Aunque los resultados dependen del modelo, sMuon logró la mejor precisión en seis de once tareas al ajustar el modelo Moonlight-16B-A3B (el cual fue preentrenado con Muon). Generalmente supera a AdamW y Per-factor Muon.
Preentrenamiento ReLoRA: En un entorno de preentrenamiento usando ReLoRA (fusión y refresco periódico de adaptadores), sMoun alcanzó la pérdida de validación final más baja, empatando con Riemannion pero siendo aproximadamente un 30% más rápido.
Propiedades Espectrales: El análisis empírico de los espectros de actualización de peso muestra que sMuon retiene un espectro más uniforme (más cercano al comportamiento ideal de Muon) en comparación con Per-factor Muon y LoRA-Muon.
Eficiencia: La Tabla 4 demuestra que, si bien sMuon es ligeramente más lento que AdamW, es significativamente más rápido que Riemannion. Para un rango de $r=64$ , Riemannion añade un ~33% de sobrecarga al tiempo de paso, mientras que sMuon añade solo un ~0.8%.

Significancia y Reivindicaciones

El artículo afirma que sMuon proporciona una forma fundamentada, eficiente en memoria y computacionalmente tratable de aplicar la optimización Muon a adaptadores de bajo rango. Al respetar la geometría de la capa completa en lugar de tratar los factores de LoRA como matrices independientes, sMuon logra mejoras de rendimiento moderadas, particularmente cuando el modelo base también fue preentrenado con Muon. Los autores enfatizan que su método evita las "costosas restricciones de variedad" de la optimización de Riemann sin sacrificar el rendimiento, haciendo que el ajuste fino basado en Muon sea más accesible para aplicaciones prácticas.

Los autores mantienen la modestia, señalando que los resultados "dependen del modelo y la evaluación" y que la contribución principal es una aproximación limpia y eficiente que cierra la brecha entre los beneficios teóricos de Muon y las restricciones prácticas de PEFT.