Towards Understanding Gradient Flow Dynamics of Homogeneous Neural Networks Beyond the Origin

Resumen Técnico: Hacia la comprensión de la dinámica del flujo de gradiente en redes neuronales homogéneas más allá del origen

Planteamiento del problema
Las redes neuronales profundas modernas entrenadas mediante métodos basados en gradiente exhiben a menudo una fuerte generalización a pesar de la sobreparametrización, un fenómeno atribuido a la regularización implícita inducida por el algoritmo de entrenamiento. Si bien se ha logrado un progreso significativo en la comprensión de la dinámica de entrenamiento en el régimen del Kernel de Tangente Neuronal (NTK) (inicialización grande) y en las etapas tardías del entrenamiento, el régimen de "inicialización pequeña" sigue siendo menos comprendido. En este régimen, las redes operan en un modo de aprendizaje de características, pero la evolución altamente no lineal de los pesos plantea desafíos analíticos.

Trabajos previos (por ejemplo, Maennel et al., 2018; Kumar y Haupt, 2024, 2025) establecieron que para redes neuronales homogéneas con inicialización pequeña, los pesos permanecen cerca del origen durante un período prolongado. Durante esta fase, los pesos convergen en dirección hacia un punto de Karush-Kuhn-Tucker (KKT) de una Función de Correlación Neuronal (NCF), un fenómeno conocido como "convergencia direccional temprana". Además, estos puntos KKT de la etapa temprana suelen exhibir dispersión (sparsity). Sin embargo, el comportamiento del flujo de gradiente después de que los pesos escapan del origen y la preservación de la estructura de dispersión emergente durante esta transición no fueron caracterizados completamente. Este artículo aborda la brecha analizando la dinámica del flujo de gradiente de redes neuronales homogéneas inmediatamente después de que escapan del origen, hasta el primer punto de silla encontrado.

Metodología y configuración
El artículo estudia redes neuronales homogéneas $H(x; w)$ que son $L$-positivamente homogéneas ($L \ge 2$), lo que significa que $H(x; cw) = c^L H(x; w)$ para $c \ge 0$. El análisis asume que la red tiene gradientes localmente Lipschitz, lo que incluye redes lineales profundas y redes con activaciones ReLU polinómicas ($p \ge 2$), pero excluye explícitamente las redes ReLU estándar ($p=1$) debido a la no diferenciabilidad en el origen, aunque los resultados empíricos sugieren posibles extensiones.

El objetivo de entrenamiento es la minimización de una función de pérdida $\mathcal{L}(w)$ mediante el flujo de gradiente:
$$\dot{w} = -\nabla \mathcal{L}(w)$$
Los autores construyen su análisis sobre el concepto de la Función de Correlación Neuronal (NCF), definida como $N(u) = \bar{e}_y^\top H(X; u)$, donde $\bar{e}_y = -\ell'(0, y)$. Analizan la optimización restringida de $N(u)$ en la esfera unitaria.

La metodología central consiste en:

1. Construcción de la trayectoria límite: Definir una solución límite $p(t)$ que representa el camino del flujo de gradiente después de escapar del origen. Esto se logra mediante el desplazamiento temporal de la solución $\psi(t, \delta w_0)$ (donde $\delta$ es la escala de inicialización) por un término dependiente de $\ln(1/\delta)$ y tomando el límite cuando $\delta \to 0$.
2. Análisis de estabilidad: Utilizar las propiedades de los puntos KKT positivos de segundo orden $\Delta$ de la NCF restringida. Los autores asumen que la dirección inicial de los pesos se encuentra en el conjunto estable de tal punto KKT de la NCF.
3. Subconjuntos de preservación de ceros: Introducir el concepto de un "subconjunto de preservación de ceros" de los pesos —pesos inicializados en cero que permanecen en cero durante todo el entrenamiento bajo el flujo de gradiente. Esto se utiliza para analizar la preservación de la dispersión.

Contribuciones clave y resultados

1. Dinámica del flujo de gradiente post-escape (Teoremas 4 y 10)
El artículo establece que, para una inicialización suficientemente pequeña, la trayectoria del flujo de gradiente después de escapar del origen sigue de cerca un camino específico $p(t)$, independientemente de la escala de inicialización pequeña $\delta$, siempre que la dirección inicial se alinee con un conjunto estable de un punto KKT positivo de segundo orden de la NCF.

• Redes de dos homogeneidades ($L=2$): El Teorema 4 muestra que, tras un escalado de tiempo de $O(\ln(1/\delta))$, la trayectoria $\psi(t, \delta w_0)$ converge a $p(t)$ con un límite de error de $O(\delta^{\frac{\Delta}{\Delta + 8N(w^*)}})$.
• Redes homogéneas profundas ($L > 2$): El Teorema 10 extiende esto a redes $L$-homogéneas. El tiempo de escape escala como $O(\delta^{-(L-2)})$, y la trayectoria converge de manera similar a un camino límite $p(t)$ con un límite de error específico.
• Implicación: El comportamiento del flujo de gradiente inmediatamente después de escapar del origen está determinado principalmente por el punto KKT de la NCF al que los pesos se aproximaron mientras estaban cerca del origen.

2. Caracterización del primer punto de silla (Corolarios 5, 6, 11, 12)
Los autores demuestran que la trayectoria límite $p(t)$ converge a un punto estacionario (punto de silla) de la función de pérdida de entrenamiento. En consecuencia, para una inicialización suficientemente pequeña, la trayectoria real del flujo de gradiente $\psi(t, \delta w_0)$ se aproximará a este mismo punto de silla después de escapar del origen.

• El artículo distingue entre casos donde el punto de silla es finito (trayectoria acotada) y casos donde está en el infinito (trayectoria no acotada, común con la pérdida logística).
• Los resultados implican que el flujo de gradiente sigue una trayectoria de "silla a silla", caracterizando específicamente el primer punto de silla encontrado después del origen.

3. Preservación de la dispersión en redes de alimentación hacia adelante (Sección 4 y Teorema 18)
Una contribución crítica es la prueba de que la estructura de dispersión que emerge durante la fase de "cerca del origen" se preserva después de escapar del origen, hasta el primer punto de silla.

• Subconjuntos de preservación de ceros: Los autores demuestran que, para redes homogéneas de alimentación hacia adelante (feed-forward), las filas y columnas de las matrices de pesos correspondientes a las entradas de magnitud cero en un punto KKT de la NCF forman un "subconjunto de preservación de ceros". Si estos pesos se inicializan en cero, permanecen en cero durante todo el entrenamiento.
• Preservación de la dispersión: Dado que el flujo de gradiente con inicialización pequeña converge en dirección a un punto KKT con una dispersión específica (donde ciertos pesos son cero), y que estos pesos cero forman un subconjunto de preservación de ceros, la estructura de dispersión se mantiene mientras la red escapa del origen y se aproxima al primer punto de silla.
• Teorema 18: Formaliza que tanto para redes con $L=2$ como con $L>2$, la magnitud de los pesos en el subconjunto de preservación de ceros permanece insignificante ($O(\delta^{\text{exponente}}$) incluso después del tiempo de escape, preservando efectivamente el patrón de dispersión.

Validación experimental
El artículo proporciona experimentos numéricos (Figuras 1, 4, 5, 6) utilizando funciones de activación cuadráticas ($L=2$) y activaciones ReLU polinómicas.

• Activación cuadrática: Los experimentos confirman que los pesos se vuelven dispersos antes de escapar del origen y que esta dispersión se preserva en el siguiente punto de silla.
• Activación ReLU: Aunque el marco teórico excluye la ReLU estándar ($L=1$), los experimentos en redes ReLU muestran que, si bien los pesos no se vuelven perfectamente dispersos (hay más entradas no nulas que en el caso polinómico), la estructura de dispersión relativa (donde los pesos pequeños permanecen pequeños) se preserva. Esto sugiere que los hallazgos pueden extenderse a las redes ReLU empíricamente, a pesar de la falta de gradientes localmente Lipschitz en el origen.

Significado y afirmaciones
El artículo afirma proporcionar una caracterización rigurosa de un segmento específico de la dinámica de entrenamiento: la transición de la fase de "cerca del origen" al primer punto de silla.

• Mecanismo de regularización implícita: Vincula la regularización implícita observada en los regímenes de inicialización pequeña con la geometría de la NCF y los puntos KKT específicos a los que la red converge.
• Estabilidad de la dispersión: Ofrece una explicación teórica de por qué los patrones de dispersidad observados en las etapas tempranas del entrenamiento persisten, atribuyéndolo a la naturaleza de "preservación de ceros" de subconjuntos de pesos específicos en arquitecturas de alimentación hacia adelante.
• Limitaciones: Los autores declaran explícitamente la modestia de su alcance. Los resultados están limitados a:
  1. Redes homogéneas con gradientes localmente Lipschitz (excluyendo la ReLU estándar).
  2. La dinámica solo hasta el primer punto de silla encontrado después de escapar del origen. El comportamiento más allá de este punto (por ejemplo, escapar del punto de silla hacia un mínimo global) no está cubierto y se identifica como una dirección para trabajos futuros.
  3. La suposición de que la dirección inicial se encuentra en el conjunto estable de un punto KKT positivo de segundo orden.

En resumen, el artículo cierra la brecha entre la convergencia direccional de la etapa temprana y la fase de entrenamiento subsiguiente, demostrando que las propiedades estructurales (específicamente la dispersión) establecidas mientras se está cerca del origen dictan la trayectoria y la estabilidad de la red a medida que comienza a aprender características y se aproxima al primer punto de silla.