Sharp feature-learning transitions and Bayes-optimal neural scaling laws in extensive-width networks

Resumen Técnico: Transiciones Agudas de Aprendizaje de Características y Leyes de Escala Óptimas Bayes en Redes de Ancho Extensivo

Configuración del Problema
El artículo investiga los límites teóricos de la información para aprender una red neuronal de profesor de una sola capa oculta con características jerárquicas a partir de consultas ruidosas, específicamente en el contexto de la transferencia de conocimiento a un modelo estudiante más pequeño. El estudio opera en el régimen de ancho extensivo, donde el ancho del profesor $k$ escala linealmente con la dimensión de entrada $d$ ($k = \Theta(d)$). Este escenario se sitúa entre el régimen de múltiples índices bien estudiado ($k = \Theta(1)$) y el régimen del Núcleo Tangente Neuronal (NTK) ($k \gg d$), capturando el ancho "grande pero finito" de las redes prácticas.

La red del profesor genera etiquetas $y_\mu = v^\top \sigma(W_0 x_\mu)$ a partir de entradas $x_\mu \sim \mathcal{N}(0, I_d)$, con ruido gaussiano aditivo. El vector de lectura $v$ sigue una estructura jerárquica (por ejemplo, decaimiento de ley de potencia $v_i \propto i^{-\beta}$), lo que implica que las características tienen una importancia que decae geométricamente. Los autores analizan dos tipos de aprendices:

1. Estudiante Bayesiano: Conoce el mecanismo de generación de datos (activación $\sigma$, varianza del ruido $\Delta$, lectura $v$) pero no los pesos específicos del profesor $W_0$. Predice mediante la media posterior, logrando el error óptimo bayesiano ($\varepsilon_{BO}$).
2. Estudiante SGD: Una red de una sola capa oculta de ancho $k_s \leq k$ entrenada con ADAM.

Las preguntas principales abordadas son: (P1) Dado un conjunto de datos fijo, ¿qué características del profesor son estadísticamente recuperables? y (P2) ¿Cómo restringe el subconjunto recuperable de características la ley de escala de $\varepsilon_{BO}$ con respecto al tamaño de los datos $n$?

Metodología
Los autores emplean un argumento heurístico de desacoplamiento de dejar-uno-afuera (cavidad), una técnica estándar en física estadística, para derivar caracterizaciones asintóticamente precisas del proceso de aprendizaje.

• Desacoplamiento: El método aísla la contribución de una sola característica $w_i$ tratando las contribuciones de todas las demás características como un término de ruido efectivo.
• Suposición Gaussiana: Los autores asumen que los "campos de cavidad" (el ruido efectivo de otras características) se comportan como variables gaussianas independientes. Esto reduce el problema de inferencia multi-característica a una secuencia de modelos lineales generalizados (GLM) escalares independientes.
• Sistema de Punto Fijo: Al analizar la superposición entre la dirección de la característica verdadera y la estimación posterior, los autores derivan un sistema cerrado de ecuaciones de punto fijo. Estas ecuaciones relacionan la superposición $q_i$ de la característica $i$ con la relación señal-ruido (SNR) y el error de generalización global $\varepsilon$.
• Validación: Las predicciones teóricas se validan numéricamente utilizando simulaciones de Monte Carlo Hamiltoniano (HMC) en diversos regímenes de parámetros, mostrando un excelente acuerdo incluso en dimensiones moderadas ($d \sim 10^2$).

Resultados Clave

1. Escalera Aguda de Transiciones de Aprendizaje de Características:
   El análisis revela que la recuperabilidad de las características está gobernada por una transición de fase aguda. Una característica $i$ es recuperable (superposición $q_i > 0$) si y solo si su SNR efectiva, $snr_i = \frac{(n/d)v_i^2}{\Delta + \varepsilon_{BO}}$, supera un umbral crítico dependiente de la activación $\lambda_\sigma$.

   • A medida que aumenta el tamaño de los datos $n$, las características cruzan este umbral secuencialmente, comenzando por las más fuertes (mayor $v_i$).
   • Esto resulta en una "escalera" de saltos discontinuos en las superposiciones de características y caídas correspondientes en el error de generalización, en lugar de una transición suave.

2. Ancho Efectivo ($k_c$) y Descomposición del Error:
   El número de características recuperables con un presupuesto de datos dado se define como el ancho efectivo $k_c = |\{i : q_i > 0\}|$. El error óptimo bayesiano se descompone en dos términos:
   $$\varepsilon_{BO} = \Theta\left(\frac{k_c d}{n}\right) + \Theta\left(\sum_{j > k_c} v_j^2\right)$$
   El primer término representa el error de estimación en las características recuperables (escalando como $d/n$ por característica), mientras que el segundo término es el sesgo irreducible de las características no recuperables.

3. Leyes de Escala Unificadas:
   Para lecturas de ley de potencia con exponente $\beta > 1/2$, el artículo deriva dos regímenes de escala distintos que se unifican bajo la relación $\varepsilon_{BO} = \Theta(k_c d / n)$:

   • Régimen de Aprendizaje de Características ($n \ll k^{2\beta}d$): Aquí, $k_c \sim (n/d)^{1/(2\beta)}$. El error escala como $\varepsilon_{BO} \sim n^{1/(2\beta) - 1}$. En este régimen, el modelo está limitado por el número de características que puede aprender.
   • Régimen de Refinamiento ($n \gg k^{2\beta}d$): Todas las características son recuperables ($k_c = k$). El error escala como $\varepsilon_{BO} \sim n^{-1}$, recuperando la tasa paramétrica clásica.
     El punto de cruce $n^\star = \Theta(k^{2\beta}d)$ marca la transición donde la característica más débil se vuelve recuperable.

4. Implicaciones para Estudiantes Entrenados con SGD:
   Los resultados empíricos muestran que un estudiante entrenado con ADAM con ancho $k_s$ cercano al ancho efectivo teórico $k_c$ logra estas leyes de escala óptimas (hasta una pequeña brecha algorítmica).

   • Si $k_s < k_c$, el estudiante se ajusta insuficientemente, limitado por su capacidad.
   • Si $k_s \gg k_c$, el estudiante sobreajusta direcciones no recuperables, degradando el rendimiento.
   • El ancho óptimo del estudiante escala con el tamaño de los datos según la jerarquía de características, proporcionando una base teórica para la transferencia de hiperparámetros.

Importancia y Afirmaciones
El artículo afirma proporcionar la primera caracterización analítica del aprendizaje óptimo bayesiano en el régimen de ancho extensivo ($k = \Theta(d)$) para redes con características jerárquicas y coeficientes de Hermite segundos nulos ($\mu_2 = 0$).

• Unificación: Unifica dos regímenes de escala empíricos previamente distintos (aprendizaje de características y refinamiento) en un único marco impulsado por el ancho efectivo $k_c$.
• Mecanismo: Ofrece una explicación mecanicista de la dinámica de aprendizaje en "escalera" observada en la práctica, atribuyéndola a transiciones de fase secuenciales en la recuperabilidad de características.
• Contribución Metodológica: El trabajo extiende el método de cavidad más allá de lecturas densas hacia estructuras de lectura arbitrarias (incluyendo dispersas y de ley de potencia), recuperando resultados conocidos para lecturas densas mientras se generaliza a configuraciones dispersas donde las tasas de decaimiento del error difieren significativamente.
• Perspectiva Práctica: El ancho efectivo $k_c$ derivado proporciona una prescripción de primeros principios para dimensionar modelos estudiantes basados en presupuestos de datos, sugiriendo que la capacidad del modelo debe escalar con el número de características recuperables en lugar del número total de características del profesor.

Limitaciones
Los autores reconocen que sus derivaciones dependen de suposiciones no probadas, específicamente la gaussianidad de los campos de cavidad y la concentración de las superposiciones. Además, el análisis se restringe a activaciones con un coeficiente de Hermite segundo nulo ($\mu_2 = 0$), lo que incluye activaciones impares pero excluye otras como ReLU (aunque el artículo señala que ReLU puede analizarse mediante el régimen complementario $\mu_2 \neq 0$ en otros trabajos). La reducción de dejar-uno-afuera es heurística, aunque validada numéricamente.