Correlation flow governs learning at criticality

Resumen Técnico: El flujo de correlación gobierna el aprendizaje en la criticidad

Planteamiento del problema
La inicialización de las redes neuronales profundas es crítica para determinar si la información y los gradientes pueden propagarse a través de la profundidad. Si bien el éxito empírico ha impulsado el escalado del tamaño de las redes, una teoría unificada que conecte las propiedades de la inicialización con la dinámica de aprendizaje sigue siendo elusiva. Específicamente, existe una falta de comprensión sobre cómo la interacción entre la propagación de la correlación, el espectro del Jacobiano de extremo a extremo y el Kernel de Tangente Neural (NTK) está condicionada por las elecciones de inicialización (Gaussiana vs. ortogonal) en el límite de redes infinitamente anchas e infinitamente profundas. Trabajos previos han establecido que la "criticidad" (un punto específico en el plano de varianza de pesos-sesgo) permite una decaimiento algebraico en lugar de exponencial de las correlaciones, pero la relación precisa entre este flujo de correlación y la dinámica de aprendizaje resultante (gobernada por el NTK) no ha sido caracterizada completamente. Además, el papel de la inicialización ortogonal para controlar las correcciones de tamaño finito en este límite requiere un fundamento teórico riguroso.

Metodología
Los autores emplean una combinación de teoría de campo medio y teoría de matrices aleatorias para analizar redes neuronales totalmente conectadas en el límite secuencial, donde el ancho $N$ se lleva al infinito antes que la profundidad $L$ ($N \to \infty$ seguido de $L \to \infty$).

1. Marco Teórico:

   • Régimen de Campo Medio: Los autores analizan las relaciones de recursión deterministas para la covarianza y la correlación de las pre-activaciones a medida que aumenta la profundidad. Identifican tres regímenes: ordenado (correlación $\to 1$), caótico (correlación $\to$ constante $<1$), y un fase de varianza evanescente recientemente identificada (donde la varianza $\to 0$).
   • Punto Crítico: Localizan el punto crítico en la intersección de la frontera de la fase de varianza evanescente y la línea del "borde del caos", definida por $\sigma_w = 1/\phi_1$ y $\sigma_b = 0$ (donde $\phi_1$ es el primer coeficiente de Taylor de la función de activación).
   • Análisis del Jacobiano: Utilizando la teoría de probabilidad libre, analizan el espectro de valores singulares de los Jacobianos por capa y de extremo a extremo. Derivan leyes de escala para la media y la varianza del espectro del Jacobiano bajo inicialización tanto Gaussiana como ortogonal.
   • Derivación del NTK: Expresan el NTK como una suma de kernels de covarianza arrastrados por Jacobianos parciales. Al analizar el comportamiento asintótico de estos términos en la criticidad, derivan la relación entre el NTño y la correlación de punto fijo.

2. Análisis de Tamaño Finito:

   • Los autores derivan las correcciones de orden principal de tamaño finito ($O(L/N)$) para el kernel empírico. Descomponen estas fluctuaciones en "defectos de Wick" (derivados de la no-gaussianidad de las pre-activaciones) y efectos de "no-ortogonalidad de los pesos".
   • Comparan cómo se acumulan estos errores bajo la inicialización Gaussiana frente a la ortogonal.

3. Validación Numérica:

   • Las predicciones teóricas se validan utilizando redes totalmente conectadas con funciones de activación seno (y algunas de tanh). Los experimentos rastrean la propagación de la variedad (manifold), el flujo de correlación, los espectros del Jacobiano y el condicionamiento del NTK a través de diversas profundidades y anchuras.

Contribuciones Clave y Resultados

1. La criticidad es necesaria para la propagación de la correlación:
   Los autores demuestran que la propagación de la correlación no trivial hacia una profundidad infinita es posible solo en el punto crítico. En todos los demás regímenes, las correlaciones colapsan hacia 1 (ordenado) o hacia una constante independiente de la entrada (caótico), o la varianza desaparece exponencialmente, causando un colapso numérico. En la criticidad, la correlación converge algebraicamente, preservando la dependencia de la entrada.

2. Desvanecimiento algebraico del Jacobiano de extremo a extremo:
   Contrario a la esperanza de que la criticidad pudiera permitir la "isometría dinámica" (donde el espectro del Jacobiano se concentra cerca de la unidad), los autores demuestran que en la criticidad, el Jacobiano de extremo a extremo se desvanece algebraicamente con la profundidad ($\Xi \propto 1/L$). Esto ocurre incluso con la inicialización ortogonal. En consecuencia, la isometría dinámica está ausente en el límite secuencial en la criticidad.

3. Equivalencia entre el NTK y la Correlación de Salida:
   El resultado teórico central del artículo es la demostración de que, en el límite de profundidad infinita en la criticidad, el NTK se vuelve exactamente proporcional a la matriz de correlación de punto fijo (específicamente, $\Theta_L \to \frac{1}{A\sigma_w^2} \rho_\infty$). Esto establece un vínculo directo entre la propagación de la información (correlación) y la dinámica de aprendizaje (espectro del NTK). Esta equivalencia se mantiene independientemente del desvanecimiento del Jacobiano, ya que el decaimiento es compensado en los Jacobianos de traslape reescalados utilizados en la definición del NTK.

4. Papel de la Inicialización Ortogonal:
   Aunque la inicialización ortogonal no puede restaurar la isometría dinámica, juega un papel crucial en el control de los efectos de tamaño finito.

   • Bajo la inicialización Gaussiana, las fluctuaciones de tamaño finito en la correlación y el NTK escalan como $O(\sqrt{L/N})$, acumulándose con la profundidad y provocando espectros mal condicionados.
   • Bajo la inicialización ortogonal, se suprimen las fluctuaciones de orden principal derivadas de la no-ortogonalidad de los pesos. Las fluctuaciones restantes (defectos de Wick) escalan como $O(1/\sqrt{N})$ y son independientes de la profundidad. Esto permite que las predicciones del límite secuencial se mantengan para redes mucho más profundas.

5. Restricciones de la Función de Activación:
   Los autores muestran que para que la información se propague sin colapsar a un punto fijo trivial, la función de activación debe satisfacer $\phi_2 = 0$ (es decir, debe ser una función impar en su orden más bajo). Si $\phi_2 \neq 0$, las correlaciones colapsan independientemente de la inicialización.

Significancia y Reivindicaciones
El artículo afirma proporcionar un marco unificado que conecta la propagación de la información, el control del gradiente y la dinámica de aprendizaje en redes profundas. Su principal significancia radica en:

• Unificar Correlación y Aprendizaje: Demuestra que el espectro del NTK, que dicta la velocidad y dirección del aprendizaje, está determinado asintóticamente por la estructura de correlación de las entradas en la criticidad. Esto explica por qué preservar la similitud de la entrada a través de la red es esencial para el entrenamiento.
• Justificar la Inicialización Crítica: Proporciona una justificación teórica rigurosa para inicializar en el punto crítico, mostrando que es la única elección que preserva las correlaciones dependientes de la entrada mientras mantiene una estructura de Jacobiano bien definida (aunque evanescente).
• Clarificar la Inicialización Ortogonal: Resuelve la tensión entre los regímenes "congelado" (lineal) y de "aprendizaje de características" al mostrar que la inicialización ortogonal es la condición necesaria para suprimir el crecimiento incontrolado de las fluctuaciones de tamaño finito que de otro modo destruirían la predictibilidad del NTK en redes profundas.
• Implicaciones Prácticas: Los resultados sugieren que para redes muy profundas (por ejemplo, SIRENs, Representaciones Neuronales Implícitas), la inicialización crítica combinada con pesos ortogonales es el enfoque basado en principios para asegurar dinámicas de entrenamiento estables y un sesgo espectral controlado, donde la red aprende a lo largo de los modos principales de la estructura de correlación del conjunto de datos de entrada.

Los autores señalan que estos resultados dependen del límite secuencial y de supuestos específicos de la función de activación (acotada, diferenciable, $\phi_2=0$), y reconocen que extender estos hallazgos al límite conjunto ($L/N \to c$) o a activaciones no acotadas (como ReLU) sigue siendo un área abierta para trabajos futuros.