The Implicit Bias of Adam and Muon on Smooth Homogeneous Neural Networks

Este artículo demuestra que optimizadores basados en momentum como Adam y Muon, al aplicarse a redes neuronales homogéneas suaves, exhiben un sesgo implícito hacia la maximización de márgenes bajo normas específicas, comportándose como trayectorias de descenso de gradiente más pronunciado bajo un esquema de tasa de aprendizaje decreciente.

Lo esencial

¡Claro que sí! Imagina que entrenar una red neuronal es como enseñar a un grupo de estudiantes a resolver un problema matemático muy difícil. El objetivo es que todos lleguen a la respuesta correcta, pero lo más interesante es cómo llegan allí, porque el camino que eligen determina qué tan bien entenderán el problema en el futuro (esto se llama "generalización").

Los autores de este paper, Eitan Gronich y Gal Vardi, se preguntaron: ¿Qué "sesgo" o tendencia oculta tienen diferentes métodos de aprendizaje (optimizadores) al guiar a estos estudiantes?

Aquí tienes la explicación con analogías sencillas:

1. El Problema: ¿Quién es el mejor profesor?

En el mundo de la Inteligencia Artificial, hay muchos "profesores" (algoritmos) que pueden enseñar a la red neuronal.

• Descenso de Gradiente (GD): Es el profesor clásico y serio. Camina paso a paso, siempre bajando por la pendiente más empinada.
• Adam: Es el profesor moderno y popular. Usa un "momento" (inercia) para no detenerse en pequeños baches y ajusta su paso dinámicamente.
• Muon: Es un profesor nuevo y muy potente, especialmente para redes grandes (como las que usan las IAs generadoras de texto).

El paper descubre que, aunque todos intentan minimizar el error, cada uno tiene una "brújula" interna diferente que los empuja hacia un tipo específico de solución.

2. La Analogía del "Margen de Seguridad"

Imagina que tienes que dibujar una línea para separar dos grupos de personas en una habitación: los que usan zapatos rojos y los que usan zapatos azules.

• El objetivo: Dibujar una línea que los separe perfectamente.
• El "Margen": Es la distancia entre la línea y la persona más cercana de cada grupo.
• La teoría: Los investigadores dicen que los algoritmos no solo quieren separar a la gente; quieren maximizar esa distancia de seguridad. Cuanto más lejos esté la línea de la gente, más seguro es que no se equivocarán si entra alguien nuevo.

3. La Gran Revelación: Cada optimizador mide la distancia de forma distinta

Aquí es donde entra la magia del paper. Depende de qué "regla de medida" use el optimizador, la línea final se verá diferente:

• El Descenso de Gradiente (GD) y Muon (en ciertos casos):

  • Analogía: Usan una regla métrica estándar. Miden la distancia en línea recta (como si fueras en un avión sobre la ciudad).
  • Resultado: Buscan la solución que maximiza la distancia "euclidiana" (la más común). Es como buscar el camino más directo y equilibrado.

• Adam (y Signum):

  • Analogía: Usan una regla de "peor caso". Imagina que te preocupa tanto no chocar con nadie que mides la distancia solo por el obstáculo más cercano en cualquier dirección individual.
  • Resultado: Maximizan el margen en la dirección más "estrecha" (norma L-infinito). Es como si el profesor Adam dijera: "No me importa si estoy lejos de todos, solo me importa no chocar con el estudiante que está justo al lado de mi nariz".

• Muon (en redes profundas):

  • Analogía: Usa una regla de "fuerza de equipo". En lugar de mirar a cada estudiante individualmente, mira a las filas y columnas de la clase como bloques enteros.
  • Resultado: Maximiza un margen basado en la estructura de las matrices de pesos (norma espectral). Es como si el profesor Muon organizara a los estudiantes en filas y asegurara que toda la fila esté segura, no solo individuos sueltos.

4. El Concepto Clave: "Inercia" (Momentum)

El paper explica que algoritmos como Adam y Muon tienen "momento".

• Analogía: Imagina que empujas un carrito de compras. Si el carrito tiene mucho peso (momento), no gira instantáneamente cuando ves un obstáculo; sigue un poco de su trayectoria anterior.
• Los autores demuestran que, si el "profesor" (el algoritmo) reduce su velocidad (tasa de aprendizaje) muy lentamente con el tiempo, esa inercia hace que el carrito termine siguiendo exactamente el mismo camino que el "Descenso de Gradiente más empinado", pero adaptado a su propia regla de medida.

5. ¿Por qué importa esto?

Antes, pensábamos que Adam y Muon eran simplemente "más rápidos" o "más eficientes". Este paper nos dice que son diferentes en su esencia.

• Si eliges Adam, estás forzando a tu red neuronal a buscar una solución que sea muy robusta contra errores individuales (como un escudo contra ataques pequeños).
• Si eliges Muon, estás buscando una solución que respete la estructura global de los datos.

En resumen:
El paper es como un manual de instrucciones que revela que cada herramienta de construcción (optimizador) tiene su propia "huella digital". No todos construyen el mismo edificio, aunque todos intenten resolver el mismo problema. Entender esto nos ayuda a elegir la herramienta correcta para el trabajo específico que queremos hacer, asegurando que nuestra Inteligencia Artificial no solo aprenda, sino que aprenda de la manera más segura y robusta posible.

¡Es como saber que un martillo es genial para clavos, pero si intentas usarlo para atornillar, el resultado será diferente!

Resumen técnico

Resumen Técnico: Sesgo Implícito de Adam y Muon en Redes Neuronales Homogéneas Suaves

1. Planteamiento del Problema

Las redes neuronales profundas, a pesar de estar sobreparametrizadas y entrenadas sin regularización explícita, muestran un rendimiento de generalización notable. La literatura atribuye esto al sesgo implícito de los algoritmos de optimización basados en gradientes, que tienden a converger hacia soluciones que maximizan el margen de los puntos de entrenamiento.

Mientras que trabajos anteriores han estudiado exhaustivamente el sesgo implícito del Descenso de Gradiente (GD) y sus variantes en modelos lineales y redes homogéneas, existe una brecha de conocimiento sobre optimizadores modernos y populares como Adam y Muon. Específicamente:

• Se desconoce qué tipo de margen maximizan estos optimizadores en modelos no lineales (homogéneos).
• La mayoría de los análisis previos se limitan a modelos lineales o asumen la presencia de constantes de estabilidad en Adam que no reflejan su comportamiento en la práctica (donde la constante suele ser despreciable).
• No se ha caracterizado el comportamiento de optimizadores "híbridos" como Muon-Adam o Muon-Signum.

El objetivo de este trabajo es extender la teoría del sesgo implícito a modelos homogéneos suaves (que incluyen redes con activaciones como ReLU cuadrada o funciones suaves) para optimizadores basados en momentum como Adam, Muon y sus variantes.

2. Metodología y Marco Teórico

Los autores desarrollan un marco unificado basado en la noción de Descenso de Gradiente Aproximado (Approximate Steepest Descent).

• Modelos Homogéneos: Se asume que la función del modelo $f(x; \theta)$ es $L$-homogénea, es decir, $f(x; \alpha\theta) = \alpha^L f(x; \theta)$. Esto incluye redes lineales profundas y redes con activaciones suaves homogéneas (ej. $z \mapsto \max(0, z)^q$ con $q>1$).

• Pérdidas: Se consideran pérdidas log-cóncavas con cola exponencial (como la pérdida exponencial y logística).

• Definición de Descenso de Gradiente Aproximado:
  Un algoritmo se considera un "Descenso de Gradiente Aproximado" si, asintóticamente, la dirección de actualización se alinea con la dirección de descenso más pronunciado (steepest descent) bajo una norma específica, incluso si no lo es exactamente en cada paso.

  • Para Muon, se demuestra que es un descenso de gradiente momentum normalizado bajo la norma del espectro máximo ($\|\cdot\|_{msp}$), definida como el máximo de las normas espectrales de las matrices de pesos.
  • Para Adam (sin constante de estabilidad), se demuestra que se comporta asintóticamente como un descenso de gradiente de signo (Signum), lo que corresponde a la maximización del margen bajo la norma $L_\infty$.
  • Para Muon-Adam y Muon-Signum, se demuestra que son casos de descenso de gradiente momentum normalizado bajo una norma híbrida (máximo de la norma espectral de las matrices y la norma $L_\infty$ de los parámetros no matriciales).

• Condiciones Clave:

  • Tasa de Aprendizaje Decreciente: Se requiere una programación de tasa de aprendizaje $\eta(t)$ que decaiga pero cuya integral diverja ($\int_0^\infty \eta(t) dt = \infty$).
  • Convergencia Direccional: Se asume que la dirección de los parámetros $\frac{\theta_t}{\|\theta_t\|}$ converge a un límite $\bar{\theta}$.
  • Análisis de Momentum: Utilizan propiedades asintóticas de los estimadores de momentum (EMA) para demostrar que, bajo una tasa de aprendizaje decreciente, el momentum se alinea con los gradientes significativos, permitiendo tratar estos optimizadores como aproximaciones del descenso de gradiente más pronunciado.

3. Contribuciones Clave

1. Extensión a Descenso de Gradiente Normalizado: Generalizan los resultados de Tsilivis et al. (2025) para el descenso de gradiente normalizado con una programación de tasa de aprendizaje arbitraria, demostrando que cualquier punto límite de la dirección de los parámetros es un punto KKT (Karush-Kuhn-Tucker) del problema de maximización de margen correspondiente.
2. Caracterización del Sesgo de Muon: Demuestran que Muon (y variantes como Muon-Signum) tiene un sesgo implícito hacia la maximización del margen bajo la norma definida por las normas espectrales de las matrices de pesos ($\|\cdot\|_{msp}$).
3. Caracterización del Sesgo de Adam: Demuestran que Adam (sin constante de estabilidad) bajo una tasa de aprendizaje decreciente tiene un sesgo hacia la maximización del margen $L_\infty$. Esto corrige y extiende trabajos previos que solo analizaban Adam en modelos lineales.
4. Análisis de Optimizadores Híbridos: Proporcionan el primer análisis teórico para combinaciones como Muon-Adam, mostrando que maximizan un margen definido por una norma compuesta: $\max(\alpha \|\mathbf{W}\|_{msp}, \|\mathbf{u}\|_\infty)$.
5. Marco Unificado de "Descenso Aproximado": Introducen una definición flexible de "Descenso de Gradiente Aproximado" que permite unificar el análisis de algoritmos adaptativos (como Adam) y de momentum bajo un mismo marco teórico, superando la necesidad de que el algoritmo sea un descenso de gradiente exacto en cada paso.

4. Resultados Principales

• Teorema de Convergencia a KKT: Bajo las suposiciones de homogeneidad, suavidad (o condiciones de estratificación para no suaves) y convergencia direccional, cualquier punto límite de la trayectoria de Adam, Muon o sus variantes es un punto KKT del problema de minimización de norma sujeto a restricciones de margen.
• Identidad de la Norma: La norma específica que se minimiza (y por tanto, el margen que se maximiza) depende estrictamente del optimizador:
  • GD / Momentum GD: Maximiza el margen $L_2$.
  • Signum / Adam (sin $\epsilon$): Maximiza el margen $L_\infty$.
  • Muon: Maximiza el margen bajo la norma espectral máxima ($L_{sp}$ o $L_{msp}$).
  • Muon-Adam: Maximiza un margen híbrido.
• Validación Empírica: Los experimentos en redes de dos capas (MNIST binario) confirman la teoría. Se observa que:
  • Adam y Signum convergen a soluciones con márgenes $L_\infty$ más altos.
  • Muon converge a soluciones con márgenes $L_{msp}$ más altos.
  • La convergencia direccional se mantiene en la práctica (alta similitud de coseno con el último iterado).
  • Los resultados se mantienen tanto para activaciones ReLU como ReLU cuadrada (suave), aunque las redes suaves muestran un comportamiento más alineado con la teoría.

5. Significado e Impacto

Este trabajo es fundamental por varias razones:

• Puente Teórico: Cierra la brecha entre la teoría del sesgo implícito en modelos lineales y la realidad de los modelos no lineales modernos utilizados en LLMs y Vision Transformers.
• Guía para la Elección de Optimizadores: Proporciona una base teórica para entender cómo la elección del optimizador afecta la geometría de la solución final. Por ejemplo, si se desea una solución con propiedades específicas de robustez o generalización asociadas a la norma $L_\infty$ o espectral, el optimizador debe elegirse en consecuencia.
• Validación de Prácticas Reales: Al analizar Adam sin la constante de estabilidad ($\epsilon \to 0$), el estudio refleja con mayor precisión el comportamiento real de los optimizadores en entrenamiento a gran escala, donde los gradientes suelen ser mucho mayores que $\epsilon$.
• Nuevas Direcciones de Investigación: Abre preguntas sobre la generalización de estos resultados a modelos no homogéneos y la posibilidad de probar la convergencia direccional sin asumir la condición (T2) para optimizadores basados en momentum.

En conclusión, el artículo establece que el sesgo implícito no es universal, sino que está intrínsecamente ligado a la norma geométrica subyacente del algoritmo de optimización, y demuestra que optimizadores modernos como Adam y Muon poseen sesgos bien definidos hacia la maximización de márgenes bajo normas específicas ($L_\infty$ y espectral, respectivamente) en redes neuronales homogéneas.