Resumen Técnico: LoRA-Muon: Descenso de Gradiente Estepado Espectral en la Variedad de Bajo Rango

1. Planteamiento del Problema

La Adaptación de Bajo Rango (LoRA) es ampliamente adoptada para el ajuste fino de grandes modelos de aprendizaje profundo debido a su eficiencia en la reducción de los FLOPs de entrenamiento y los requisitos de memoria. Sin embargo, los optimizadores estándar como AdamW, cuando se aplican de forma independiente a los factores de LoRA ($A$ y $B$), introducen dificultades de ajuste significativas en comparación con el entrenamiento denso. Estos problemas incluyen:

• Sensibilidad a la Inicialización: El rendimiento depende altamente de las elecciones de inicialización.
• Mala Transferencia de Hiperparámetros: Las tasas de aprendizaje óptimas no se transfieren bien entre diferentes rangos, anchuras, profundidades o convenciones de escalado de factores.
• Dinámicas Divergentes: Incluso partiendo de pesos compuestos idénticos ($W = AB^T$), las dinámicas de entrenamiento suelen divergir si las escalas de los factores se desvían.
• Rendimiento Subóptimo: El preentrenamiento nativo de bajo rango suele tener dificultades para igualar las líneas base densas sin un extra de estabilización.

El desafío central es diseñar un optimizador para la variedad de bajo rango que mantenga las propiedades geométricas de los optimizadores de rango completo (específicamente Muon y la familia Shampoo) y que asegure que las tasas de aprendizaje óptimas se transfieran a través de rango, anchura, profundidad y el reescalado de factores.

2. Metodología

2.1 Fundamento Geométrico

El artículo reformula el diseño de optimizadores como un problema geométrico: identificar la norma bajo la cual un optimizador realiza el descenso más escarpado (steepest descent).

• Optimizador Muon: Realiza el descenso más escarpado bajo la norma espectral ($\|\cdot\|_{2\to2}$) en el espacio de matrices ambiente $\mathbb{R}^{m \times n}$. La dirección de actualización viene dada por la función de signo de la matriz, $-\text{msign}(G)$.
• Variedad de Bajo Rango: LoRA restringe los pesos a una variedad $\mathcal{M}_r = \{W \in \mathbb{R}^{m \times n} : \text{rank}(W) = r\}$. Los autores proponen resolver el problema del descenso más escarpado directamente en el espacio tangente de esta variedad, $T_W \mathcal{M}_r$, en lugar de optimizar los factores $A$ y $B$ como coordenadas independientes.

2.2 Derivación de LoRA-Muon

Los autores derivan LoRA-Muon resolviendo el problema de minimización lineal restringida en el espacio tangente de bajo rango bajo la norma espectral.

1. Descomposición del Espacio Tangente: Las actualizaciones en $W$ se descomponen en $\Delta W = \Delta A B^T + A \Delta B^T$.
2. División de la Región de Confianza: Para hacer el problema tratable, el presupuesto de la región de confianza $\eta$ se divide equitativamente entre los dos componentes tangentes. Esto se justifica por la observación empírica de que los dos componentes se alinean rápidamente durante el entrenamiento, haciendo que la desigualdad triangular sea casi ajustada.
3. Formulación del Proyector: Utilizando la invariancia unitaria, el problema se reduce a encontrar Oráculos de Minimización Lineal (LMOs) sobre los gradientes proyectados. La actualización ideal es:
   $$\Delta W = \frac{\eta}{2} \text{LMO}(\nabla W f \cdot P_B) + \frac{\eta}{2} \text{LMO}(P_A \cdot \nabla W f)$$
   donde $P_A$ y $P_B$ son proyectores sobre los espacios columna de $A$ y $B$.
4. Especialización Espectral: Para la norma espectral, el LMO es $-\text{msign}(\cdot)$. Las actualizaciones de los factores resultantes (Ecuaciones 14–15) implican el blanqueo (whitening) de los gradientes de los factores mediante las matrices de Gram actuales ($S_A = A^TA, S_B = B^TB$) y la aplicación de la función de signo espectral:
   $$\Delta A^* = -\frac{\eta}{2} \text{msign}(\nabla_A f \cdot S_B^{-1/2}) S_B^{-1/2}$$
   $$\Delta B^* = -\frac{\eta}{2} \text{msign}(\nabla_B f \cdot S_A^{-1/2}) S_A^{-1/2}$$

2.3 Decaimiento de Peso Dividido (Split Weight Decay)

El decaimiento de peso estándar por factor provoca un desajuste entre las dinámicas de actualización de rango completo y bajo rango. Los autores derivan una regla de decaimiento de peso dividido donde el factor de decaimiento $s = \sqrt{1 - \lambda \eta}$ se aplica de forma asimétrica:
$$A_{t+1} = s A_t + \frac{1}{s} \Delta A^*_t, \quad B_{t+1} = s B_t + \frac{1}{s} \Delta B^*_t$$
Esto asegura que la actualización del peso compuesto coincida con las dinámicas deseadas de rango completo, preservando la transferibilidad de los hiperparámetros.

2.4 Invarianza de Calibre (Gauge Invariance)

Una propiedad crítica de LoRA-Muon es la invarianza de calibre. La actualización sobre el peso compuesto $W$ permanece invariante bajo transformaciones invertibles arbitrarias $(A, B) \to (AR, BR^{-T})$. Esto se logra porque la actualización depende solo de los proyectores $P_A$ y $P_B$, que son, en sí mismos, invariantes de calibre.

3. Contribuciones Clave

1. Derivación de LoRA-Muon: Una nueva familia de optimizadores derivados al resolver el problema del descenso más escarpado en la variedad de bajo rango bajo normas invariantes de la unidad, especializándose en la norma espectral.
2. Regla de Decaimiento de Peso Dividido: Una regla de actualización específica que alinea las dinámicas del entrenamiento de bajo rango con las del de rango completo, permitiendo la transferencia de hiperparámetros.
3. Prueba de Invarianza de Calibre: Prueba teórica de que la actualización de LoRA-Muon es invariante bajo la acción del grupo de calibre completo, asegurando que el optimizador dependa de la actualización de peso inducida y no de las coordenadas arbitrarias de los factores.
4. Comparación con Spectron: Demostración de que Spectron no es invariante de calibre y es sensible a los reescalados arbitrarios de los factores, a diferencia de LoRA-Muon.
5. Relación con LoRA-RITE: Prueba de que el núcleo de coordenadas QR simplificado de LoRA-RITE es algebraicamente idéntico a la actualización espectral de LoRA-Muon, aunque LoRA-Muon evita la sobrecarga computacional de las descomposiciones QR y el almacenamiento del segundo momento.
6. Validación Empírica: Evidencia de que LoRA-Muon sirve como un proxy efectivo para la búsqueda de la tasa de aprendizaje de Muon denso.

4. Resultados Experimentales

Los experimentos se realizaron utilizando un modelo de lenguaje a nivel de caracteres con el conjunto de datos TinyShakespeare utilizando líneas base densas y de bajo rango con cómputo emparejado.

• Transferencia de Tasa de Aprendizaje: La tasa de aprendizaje óptima para Muon denso ($lr = 0.1$) se transfiere perfectamente a LoRA-Muon a través de barridos de rango (2, 4, 8, 16, 32), anchura, profundidad y reescalado del factor escalar.
  • En Rango 2, LoRA-Muon recupera la mejor tasa de aprendizaje probada para el modelo denso.
  • En Rango 32, LoRA-Muon logra una pérdida de validación media inferior ($1.776 \pm 0.002$) que la línea base densa ($1.789 \pm 0.002$) en el barrido de promedio de semillas.
• Reequilibrio de Calibre: Cuando los factores se inicializan con un gran desequilibrio escalar (escala 9), aplicar un paso de reequilibrio tiene un efecto insignificante en LoRA-Muon (cambio de pérdida $\approx -5.1 \times 10^{-5}$), confirmando la invarianza de calibre. En contraste, el rendimiento de Spectron cambia significativamente ($2.0 \times 10^{-2}$), confirmando su sensibilidad al escalado de factores.
• Eficiencia: Un proxy de rango-2 utiliza solo ~4.3% de los parámetros entrenables del modelo denso mientras soporta un presupuesto de paso del optimizador ~9.2 veces mayor bajo las mismas restricciones de cómputo.

5. Significado y Reivindicaciones

El artículo afirma que LoRA-Muon proporciona una solución teóricamente fundamentada a las dificultades de ajuste de LoRA. Al tratar a LoRA como un problema de optimización en la variedad de bajo rango en lugar de un sistema de coordenadas factorizado, los autores logran:

• Robustez: El optimizador es invariante a la elección arbitraria de la factorización (simetría de calibre).
• Transferibilidad: Las tasas de aprendizaje óptimas (y, por extensión, otros hiperparámetros vía leyes de escala) se transfieren sin problemas a través de rango, anchura, profundidad y escalas de factores.
• Eficiencia: Computa la actualización espectral sin descomposiciones QR costosas ni almacenamiento de segundo momento, lo que lo hace apto para aceleradores.

Los autores posicionan a LoRA-Muon como un "buen proxy de bajo rango" para el Muon de rango completo y los optimizadores de la familia Shampoo. Enfatizan que la parametrización de bajo rango no debe dictar la geometría del optimizador; en su lugar, la actualización debe definirse por la geometría de la propia variedad de matrices de bajo rango.

Limitaciones: Los autores señalan que las afirmaciones empíricas se establecen actualmente en modelos de la escala de TinyShakespeare y aún no se han validado en el preentrenamiento a gran escala o en benchmarks de ajuste fino de tareas específicas. El estudio se centra en la transferencia de la tasa de aprendizaje más que en un barrido completo de hiperparámetros.