Resumen Técnico: Conectividad de Modos Inducida por Optimizadores: De AdamW a Muon

Enunciado del Problema

La conectividad de modos, el fenómeno por el cual modelos de redes neuronales entrenados independientemente pueden conectarse mediante una trayectoria de baja pérdida, ha sido estudiada extensamente. Sin embargo, la literatura existente trata en gran medida al conjunto de soluciones como una entidad monolítica, asumiendo a menudo una conectividad independiente de las dinámicas de optimización específicas utilizadas para alcanzar dichas soluciones. Este trabajo postula que diferentes optimizadores imponen sesgos implícitos distintos (regularización), restringiendo efectivamente las soluciones a diferentes regiones del paisaje de pérdida. La pregunta central abordada es: ¿Cómo se comporta la conectividad de modos cuando se restringe a soluciones consistentes con el sesgo implícito de un optimizador específico? Específicamente, los autores investigan si las soluciones encontradas por el mismo optimizador forman un conjunto conectado (conectividad intra-optimizador) y si las soluciones encontradas por diferentes optimizadores pueden conectarse sin incurrir en una barrera de pérdida (conectividad inter-optimizador).

Metodología

Marco Teórico

Los autores formalizan el concepto de conjuntos de soluciones regularizados por optimizador. Basándose en caracterizaciones teóricas recientes de optimizadores adaptativos (por ejemplo, AdamW y la familia Lion-K), definen los puntos límite de estos optimizadores como soluciones a problemas de optimización con restricciones:

• AdamW: Convierte a soluciones que satisfacen una restricción de norma $\ell_\infty$ ($\|\theta\|_{\max} \leq 1/\lambda$).
• Muon (y Lion-K con norma nuclear): Convierte a soluciones que satisfacen una restricción de norma de operador ($\|\theta\|_{\text{op}} \leq 1/\lambda$).

El trabajo introduce dos definiciones:

1. Conectividad intra-optimizador: El conjunto de soluciones $O_A(\lambda)$ para un optimizador específico $A$ es conexo.
2. Conectividad inter-optimizador: La unión de conjuntos de soluciones $O_{A_1}(\lambda_1) \cup O_{A_2}(\lambda_2)$ es conexa.

El análisis teórico se realiza sobre redes ReLU de dos capas en el régimen de interpolación (donde es posible lograr una pérdida de entrenamiento cero). Los autores analizan la geometría de estos conjuntos regularizados bajo diferentes anchos de red ($m$).

Configuración Experimental

La validación empírica se realiza en tareas de preentrenamiento de GPT-2 en cinco conjuntos de datos (enwik8, WikiText-103, Stories, BookCorpus, One Billion Word) utilizando dos escalas de modelo (~30M y ~50M parámetros).

• Optimizadores: AdamW y Muon.
• Construcción de Trayectorias: Los autores emplean un algoritmo modificado de interpolación policadena combinado con alineación simétrica (permutación y alineación ortogonal de cabezas de atención y neuronas MLP). Este método construye una trayectoria lineal por tramos con un punto de inflexión aprendible para minimizar las barreras de pérdida entre modelos.
• Métricas: El estudio rastrea las barreras de pérdida de validación y la evolución de los espectros de valores singulares (específicamente el rango estable) a lo largo de las trayectorias de interpolación.

Contribuciones Clave

1. Demostraciones Teóricas de Conectividad Intra-Optimizador

El trabajo demuestra que, para redes ReLU de dos capas suficientemente anchas, los conjuntos de soluciones inducidos por optimizadores específicos son conexos:

• Teorema 3.1: La conectividad intra-optimizador se cumple para Muon (y Descenso de Gradiente Normalizado con momento) cuando el ancho de la red $m \geq 4P$, donde $P$ es el número de patrones de disposición de hiperplanos distintos.
• Teorema 3.2: La conectividad intra-optimizador se cumple para AdamW (y Signum) cuando $m \geq m^*$, donde $m^*$ es un ancho crítico determinado por los datos y los parámetros de regularización.
• Mecanismo: Las demostraciones implican reducir soluciones arbitrarias a formas canónicas "esparcidas" o "igualadas" dentro del conjunto regularizado y demostrar que estas formas pueden conectarse mediante trayectorias continuas que preservan las restricciones de norma implícitas.

2. Caracterización de la Conectividad Inter-Optimizador

Los autores establecen que la conectividad entre regiones inducidas por diferentes optimizadores depende críticamente del ancho de la red y de la fuerza de regularización:

• Régimen de Gran Ancho: La unión de dos conjuntos regularizados es conexa si y solo si los conjuntos se superponen (es decir, el conjunto de pérdida cero pasa a través de la intersección de sus respectivas bolas de norma). El número de componentes conexos puede transicionar de uno a dos a medida que aumenta la fuerza de regularización.
• Construcción de Ancho Finito: El trabajo construye un conjunto de datos y un ancho específicos ($m=2$) donde el conjunto de soluciones de pérdida cero se descompone en múltiples componentes desconectados. Crucialmente, demuestran que AdamW y Muon pueden forzarse a componentes diferentes eligiendo parámetros de regularización apropiados. En este escenario, no existe una trayectoria de pérdida cero entre una solución de AdamW y una de Muon; cualquier trayectoria de conexión debe incurrir en una barrera de pérdida demostrable (al menos $1/2$ para pérdida cuadrática).

3. Transiciones Espectrales Empíricas

Los experimentos en GPT-2 revelan comportamientos espectrales distintos:

• Trayectorias del Mismo Optimizador: La interpolación entre dos modelos entrenados con el mismo optimizador (por ejemplo, AdamW-AdamW) preserva el espectro de valores singulares (por ejemplo, AdamW mantiene valores atípicos, Muon permanece isotrópico) y mantiene una pérdida baja.
• Trayectorias Cruzadas entre Optimizadores: La interpolación entre diferentes optimizadores (por ejemplo, AdamW-Muon) resulta en una transición de fase espectral suave. El espectro evoluciona desde la firma de AdamW hasta la firma de Muon a lo largo de la trayectoria. Aunque la pérdida se mantiene relativamente baja, los modelos intermedios exhiben perfiles espectrales distintos a cualquiera de los extremos.

4. Generalización Fuera de Distribución (OOD)

El trabajo demuestra que las trayectorias de interpolación cruzada entre optimizadores pueden mejorar la generalización OOD. Los modelos muestreados a lo largo de la trayectoria AdamW-Muon exhiben barreras de pérdida negativas en conjuntos de datos OOD (rendimiento superior al de la interpolación lineal de los extremos), lo que sugiere que recorrer el espacio entre diferentes sesgos de optimizador produce modelos estructuralmente diversos con propiedades de generalización superiores.

Resultados y Significancia

Resumen de Hallazgos

• Estructura Dependiente del Optimizador: El paisaje de pérdida no es un único bloque conectado de soluciones de baja pérdida; más bien, contiene "cuencas" o regiones específicas del optimizador definidas por la regularización implícita.
• La Conectividad es Condicional: Mientras que las soluciones del mismo optimizador están conectadas en redes anchas, diferentes optimizadores pueden converger a componentes desconectados en configuraciones de ancho finito, separados por barreras de pérdida.
• Huella Digital Espectral: El espectro de valores singulares sirve como un proxy robusto para la región inducida por el optimizador. Las trayectorias del mismo optimizador preservan esta huella; las trayectorias cruzadas entre optimizadores atraviesan una transición.

Significancia Afirmada

El trabajo afirma revelar una estructura dependiente del optimizador que es invisible para la literatura clásica sobre conectividad de modos. Al refinar la definición de conectividad de modos para tener en cuenta el sesgo implícito, los autores muestran que:

1. La elección del optimizador moldea fundamentalmente la topología del espacio de soluciones alcanzable.
2. La interpolación cruzada entre optimizadores es un mecanismo viable para la fusión de modelos que accede a regiones diversas del espacio de soluciones, mejorando potencialmente propiedades posteriores como la cuantización y la generalización.
3. Las garantías teóricas de conectividad deben considerar ahora el regularizador específico inducido por el optimizador, y no solo la arquitectura de la red.

El trabajo cierra la brecha entre el análisis teórico del sesgo implícito y la fusión práctica de modelos, sugiriendo que la "mejor" trayectoria entre modelos puede depender de recorrer los límites de diferentes regularizaciones inducidas por optimizadores.