Resumen Técnico: Revisitando el Kernel de Tangente Neuronal (NTK)

Planteamiento del Problema

El marco del Kernel de Tangente Neuronal (NTK) describe la dinámica de aprendizaje de las redes neuronales sobreparametrizadas entrenadas con descenso de gradiente, postulando que en el límite de ancho infinito, la red se comporta como un método de kernel con un kernel fijo. Este régimen suele denominarse "entrenamiento perezoso" (lazy training) porque las activaciones de las neuronas ocultas individuales permanecen arbitrariamente cerca de sus valores de inicialización. Sin embargo, esta caracterización presenta dos desafíos teóricos significativos:

1. El Eje del Ancho: El "régimen perezoso" implica que las características (activaciones de neuronas ocultas) son estáticas y aleatorias, lo que no logra explicar el "aprendizaje de características" observado en las redes profundas prácticas. Trabajos previos sugieren que, si bien el NTK captura dependencias complejas, lo hace sin extraer características significativas, creando una paradoja entre su éxito de generalización y su falta de evolución de características.
2. El Eje de la Profundidad: A medida que la profundidad de la red aumenta en el límite de ancho infinito, el NTK converge a un predictor independiente de los datos (una matriz constante), lo que impide que el modelo distinga entre puntos de datos que no están presentes en el conjunto de entrenamiento. Esto sugiere un callejón sin salida teórico para las redes profundas bajo la escala estándar de NTK, donde el modelo pierde expresividad y capacidades de generalización.

Este artículo desafía las conclusiones de que el NTK es puramente "perezoso" respecto al aprendizaje de características y que los NTK profundos inevitablemente degeneran en predictores triviales.

Metodología

Los autores emplean una combinación de análisis asintótico teórico y validación empírica en redes ReLU totalmente conectadas con inicialización Gaussiana estándar y pérdida de Error Cuadrático Medio (MSE).

Marco Teórico

• Notación y Definiciones: El artículo utiliza las definiciones recursivas del NTK ($\Theta^{(L)}_\infty$) y la covarianza de las neuronas ($\Sigma^{(L)}$) establecidas por Jacot et al. (2018a) y Arora et al. (2019b). Se centra en el límite de ancho infinito donde el NTK converge a un kernel determinista.
• Análisis de Características (Ancho): Los autores analizan el comportamiento agregado de las capas ocultas en lugar de las neuronas individuales. Derivan una solución de forma cerrada para el producto interno de las activaciones de la capa oculta a lo largo del tiempo, $\langle h_t(x), h_t(x') \rangle$, para demostrar que, si bien las desviaciones de las neuronas individuales desaparecen, la desviación agregada de la norma no lo hace.
• Análisis de Profundidad: El artículo investiga el límite cuando la profundidad $L \to \infty$. Introduce un kernel normalizado $\bar{\Theta}^{(L)}_\infty$ y analiza sus propiedades de convergencia. Los autores derivan límites asintóticos sobre los autovalores de la matriz del kernel y la función de predicción $f_t(x)$.
• Estrategia de Escalamiento: Para abordar la degeneración de la profundidad, el artículo propone un régimen de escalamiento específico donde la profundidad $L$, el tiempo de parada $t$ y el tamaño del conjunto de datos $n$ están acoplados. Específicamente, evita la suposición de $t \to \infty$ o $L \to \infty$ de forma aislada, requiriendo en su lugar que $L$ escale apropiadamente con $n$ y $t$ para permanecer dentro de la "fase ordenada" de la dinámica del kernel.

Validación Empírica

• Modelos de Juguete y Conjuntos de Datos: Los experimentos se realizan con datos sintéticos, el conjunto de datos text8 (usando codificaciones one-hot) y MNIST.
• Métricas: Los autores miden la desviación de la última capa oculta de su inicialización ($\|h_t - h_0\|_2$) y las tasas de convergencia de los valores y autovalores del kernel a medida que la profundidad aumenta.

Contribuciones Clave

1. Reevaluación del "Régimen Perezoso" (Eje del Ancho)

El artículo refuta la distinción binaria entre regímenes de "aprendizaje perezoso" (características estáticas) y "aprendizaje de características" (características dinámicas).

• Dinámica Agregada: Aunque las activaciones de las neuronas individuales convergen a sus valores de inicialización a medida que el ancho aumenta, la norma agregada de las desviaciones en la última capa oculta no desaparece.
• Prueba Teórica (Proposición 4): Los autores demuestran que para redes ReLU con al menos dos capas ocultas, el producto interno de las activaciones de la capa oculta para distintas entradas, $\langle h_t(x), h_t(x') \rangle$, es una función estrictamente no constante del paso de entrenamiento $t$. Esto demuestra que las correlaciones estadísticas entre puntos de datos evolucionan durante el entrenamiento, incluso en el límite de ancho infinito.
• Implicación: El "régimen perezoso" no excluye el aprendizaje de características; más bien, las características evolucionan de una manera que es capturada por la dinámica de la capa agregada, no por el seguimiento de neuronas individuales.

2. Resolución de la Expresividad en Redes Profundas (Eje de la Profundidad)

El artículo aborda el colapso del NTK hacia un predictor independiente de los datos a medida que la profundidad aumenta.

• Convergencia a una Constante: Los autores confirám que el kernel normalizado $\bar{\Theta}^{(L)}_\infty(x, x')$ converge a $1/4$ para $x \neq x'$ cuando $L \to \infty$ (Teorema 2).
• Estabilización mediante Escalamiento (Teorema 3): El artículo demuestra que la función de predicción $f_t$ puede permanecer estable y expresiva si la profundidad $L$ se escala en relación con el tamaño del conjunto de datos $n$ y el tiempo de parada $t$ se elige apropiadamente (parada temprana).
  • Al establecer $t \in o(1/\lambda_n)$ (donde $\lambda_n$ es el autovalor más pequeño) y escalar $L$ de tal manera que sea lo suficientemente grande para estabilizar el kernel pero no tanto como para inducir la independencia de los datos, el modelo retiene la capacidad de distinguir puntos de datos.
  • Los autores proporcionan un límite sobre $L$ dependiente de la separación de los datos: $L \in O\left(\frac{1}{\sqrt{1-\rho^{(2)}(x_i^\top x_j)}}\right)$.
• Generalización: Estos resultados se mantienen para conjuntos de datos arbitrarios en la esfera sin asumir que los datos sean generados por una distribución de probabilidad con una función de densidad, relajando las restricciones encontradas en la literatura previa (ej., Xiao et al., 2020).

Resultados

• Evolución de Características: Los resultados empíricos en text8 muestran que, si bien las desviaciones de las neuronas individuales son pequeñas (consistentes con el comportamiento "perezoso"), la desviación agregada de la última capa oculta crece con el entrenamiento, confirmando la predicción teórica de características no estáticas.
• Convergencia del Kernel: Los experimentos en datos sintéticos y MNIST confirman que los valores del kernel convergen a $1/4$ a una tasa de $O(1/L)$.
• Estabilidad: Al aplicar las leyes de escalamiento derivadas (Teorema 3), los autores muestran que la norma del operador de proyección en la solución de forma cerrada permanece acotada ($O(n)$) y el modelo mantiene la expresividad incluso para profundidades grandes, siempre que la profundidad no sea excesivamente grande en relación con la separación de los datos.

Significado y Reivindicaciones

El artículo afirma proporcionar una comprensión teórica más matizada del escalamiento en redes neuronales:

1. Matiz al "Régimen Perezoso": Argumenta que el descriptor de "régimen perezoso" es incompleto. Si bien las neuronas individuales son estáticas, la dinámica colectiva de la capa permite el aprendizaje de características, reconciliando el éxito de generalización del NTK con el concepto de extracción de características.
2. Viabilidad de los NTK Profundos: Desafía la noción de que los NTK de profundidad infinita son callejones sin salida teóricos. Al escalar adecuadamente la profundidad y el tiempo de parada, el NTK puede producir salidas bien comportadas y no triviales para grandes conjuntos de datos.
3. Alineación Práctica: Los hallazgos sugieren que la zona de "punto medio" (Goldilocks zone) para el análisis del NTK reside en regímenes donde la profundidad es grande pero finita, y se emplea la parada temprana, reflejando entornos experimentales prácticos donde $L \ll n$.
4. Generalización: Los autores proporcionan un conjunto de propiedades clave (dominancia diagonal, definición positiva a profundidad suficiente) que permiten generalizar estos resultados a otros kernels y arquitecturas, como las CNN, aunque explícitamente señalan que su análisis no cubre las fluctuaciones estocásticas encontradas en regímenes donde la profundidad excede vastamente al ancho (Hanin y Nica, 2020).

El artículo concluye que el NTK sigue siendo una herramienta valiosa para analizar redes sobreparametrizadas, siempre que se vaya más allá de la suposición de características estáticas y tiempo/profundidad de entrenamiento infinito sin escalamiento.