Resumen Técnico: Mejora del Entrenamiento de Redes Neuronales mediante el Desacoplamiento de la Magnitud y la Dirección de los Vectores de Peso

1. Planteamiento del Problema

El entrenamiento de las redes neuronales modernas depende de optimizadores como Adam y Muon que tratan las matrices de pesos como objetos únicos y monolíticos. Sin embargo, cada matriz de pesos $W$ conlleva inherentemente dos cantidades distintas: una magnitud ($\|W\|$) y una dirección ($\hat{W} = W/\|W\|$). En la optimización estándar, estas dos cantidades están acopladas, lo que conduce a una dinámica de entrenamiento inestable e impredecible:

• Dependencia del Cambio Direccional: El cambio angular inducido por una actualización depende inversamente de la magnitud actual. En consecuencia, la tasa de aprendizaje (LR) no controla directamente la tasa de cambio direccional, la cual puede variar entre capas y a lo largo del tiempo.
• Deriva de Magnitud No Intencionada: Las actualizaciones tienden a ser perpendiculares a los pesos actuales (debido a la invarianza de escala o al ruido), lo que provoca que la magnitud crezca incluso en ausencia de un gradiente radial. Este "crecimiento descontrolado" obliga a los profesionales a depender de mecanismos indirectos como el decaimiento de pesos (weight decay) para estabilizar el entrenamiento y el calentamiento (warmup) para prevenir la inestabilidad temprana.
• Limitaciones de Escalamiento: Los métodos actuales para escalar el entrenamiento (por ejemplo, $\mu$P) a menudo requieren el reajuste de hiperparámetros al cambiar la anchura o la profundidad del modelo. Además, los enfoques recientes que restringen los pesos a una esfera de norma fija carecen de parámetros de magnitud aprendibles, dejando el control de la escala únicamente en manos de las capas de normalización.

2. Metodología: Desacoplamiento Magnitud–Dirección (MD)

Los autores proponen el Desacoplamiento Magnitud–Dirección (MD), una modificación del optimizador que factoriza cada matriz de pesos en una dirección de norma fija y ganancias de magnitud aprendibles, actualizándolas con tasas de aprendizaje separadas.

Factorización Central

Cada matriz de pesos $W$ se reparametriza como:
$$W = \text{diag}(\gamma_{\text{row}}) \, \hat{W} \, \text{diag}(\gamma_{\text{col}})$$
Donde:

• $\hat{W}$ es la dirección, restringida a una hiperesfera de norma fija (por ejemplo, norma de Frobenius).
• $\gamma_{\text{row}}$ y $\gamma_{\text{col}}$ son ganancias de magnitud aprendibles (vectores que actúan sobre filas y columnas, respectivamente).

Procedimiento de Optimización

El método opera enteramente dentro del paso del optimizador, manteniendo sin cambios el pase hacia adelante (forward pass) del modelo (el modelo ve un único tensor fusionado $W$):

1. Recuperación: Recuperar $\hat{W}$ y las ganancias $\gamma$ a partir del peso fusionado $W$.
2. División de Gradientes: Dividir el gradiente $G = \partial L / \partial W$ en un gradiente de dirección ($\partial L / \partial \hat{W}$) y gradientes de ganancia (sumados sobre los respectivos ejes).
3. Actualización de Dirección: Actualizar $\hat{W}$ utilizando el optimizador base (por ejemplo, Adam, Muon) con una tasa de aprendizaje $\eta_W$, y luego proyectarlo de nuevo sobre la esfera para mantener la norma fija.
4. Actualización de Ganancia: Actualizar las ganancias $\gamma$ (opcionalmente mediante un mapa suave como softplus para asegurar la positividad) utilizando una tasa de aprendizaje separada $\eta_\gamma$.
5. Reensamblaje: Reconstruir $W$ para el siguiente pase hacia adelante.

Propiedades Clave

• Agnóstico al Optimizador: El método actúa como un envoltorio (wrapper) alrededor de cualquier optimizador base (Adam, Muon, Shampoo, etc.).
• Control Directo de la LR: Al fijar la dirección en una esfera, la tasa de aprendizaje establece directamente la actualización angular relativa, eliminando la dependencia de la magnitud actual del peso.
• Eliminación de Heurísticas: El método elimina la necesidad de decaimiento de pesos (ya que los pesos ya están restringidos en norma) y de calentamiento (ya que las actualizaciones tempranas grandes y desestabilizadoras se regulan desde el primer paso).
• Control de Escala de Grano Fino: A diferencia de los métodos que dependen únicamente de las ganancias de las capas de normalización, el Desacoplamiento MD proporciona un control explícito y aprendido sobre la escala de filas y columnas individuales, permitiendo que el modelo amplifique o atenúe características específicas.

3. Contribuciones Clave y Hallazgos

Resultados Empíricos

Los autores validan el Desacoplamiento MD en modelos densos (de 181M a 1.29B de parámetros) y en modelos grandes de Mezcla de Expertos (MoE) (hasta 6.7B de parámetros totales).

• Ganancias de Rendimiento: El Desacoplamiento MD mejora sobre las líneas base bien ajustadas (AdamW y Muon). En modelos MoE grandes, MuonMD (Muon con Desacoplamiento MD) alcanza la pérdida de AdamW con aproximadamente 2× menos cómputo.
• Transferencia de Tasa de Aprendizaje: El método permite la transferencia de tasa de aprendizaje zero-shot a través de la anchura del modelo. La LR de matriz óptima encontrada en un modelo pequeño sigue siendo óptima para modelos significativamente más grandes sin necesidad de reajuste. Esto se cumple tanto para el escalado de anchura como de profundidad.
• Entrenamiento sin Calentamiento: Los experimentos muestran que eliminar el calentamiento no solo mantiene la estabilidad, sino que mejora la pérdida final, ya que los pasos iniciales de entrenamiento ya no se desperdician con tasas de aprendizaje reducidas.
• Perspectivas de Ablación:
  • Eje de Normalización: Se encontró que restringir la dirección mediante la norma de Frobenius (toda la matriz) es lo más flexible y rinde de manera comparable a las restricciones de fila/columna.
  • Parametrización de Ganancia: Una ganancia combinada por fila y columna ($\gamma_{\text{row}} + \gamma_{\text{col}}$) supera a las ganancias escalares o de un solo eje. La parametrización softplus para las ganancias proporciona una ventaja ligera pero constante sobre las actualizaciones directas o exponenciales.
  • Embeddings: Restringir los vectores de embedding a una norma L2 unitaria funciona ligeramente mejor que dejarlos sin restricciones.

Comparación con el Trabajo Relacionado

El artículo sitúa el Desacoplamiento MD dentro de un cuerpo creciente de trabajo sobre entrenamiento esférico y reparametrización:

• A diferencia de nGPT, que agrupa las restricciones esféricas con cambios arquitectónicos específicos (por ejemplo, normalización L2, residuales interpolados), el Desacoplamiento MD es una modificación del lado del optimizador que puede aplicarse a arquitecturas estándar. Cuando se aplica a la arquitectura de nGPT, el Desacoplamiento MD mejora aún más el rendimiento.
• A diferencia de Muown o HyperP, que pueden fijar normas sin magnitudes aprendibles o ligar el método a optimizadores específicos, el Desacoplamiento MD aprende explícitamente las magnitudes y es agnóstico al optimizador.
• Aborda la "interferencia magnitud-dirección" identificada en trabajos previos (por ejemplo, Kosson et al., 2024a) al separar la dinámica de ambas cantidades.

4. Significado y Reivindicaciones

El artículo afirma que tratar la magnitud y la dirección como cantidades controladas por separado produce dinámicas de entrenamiento más predecibles y una mejora simple y ampliamente aplicable a los optimizadores modernos.

• Simplicidad: El método no requiere cambios arquitectónicos y añade una sobrecarga computacional mínima (operaciones elemento a elemento durante el paso del optimizador).
• Robustez: Simplifica la receta de entrenamiento al eliminar la necesidad de decaimiento de pesos y calentamiento, reduciendo el espacio de búsqueda de hiperparámetros.
• Escalabilidad: Proporciona un mecanismo fundamentado para la transferencia de la tasa de aprendizaje a través de la anchura del modelo, abordando un cuello de botella clave en el escalado de los modelos de lenguaje de gran tamaño.

Los autores concluyen que, si bien la lógica adicional en el paso del optimizador introduce una sobrecarga modesta, los beneficios en la estabilidad del entrenamiento, la eficiencia computacional y la transferencia de hiperparámetros hacen que el Desacoplamiento MD sea un avance significativo para el entrenamiento de redes neuronales a gran escala.