Resumen Técnico: Inferencia Bayesiana con MLPs Profundos No Lineales con Forma (Shaped)

Planteamiento del Problema

La teoría del aprendizaje profundo busca caracterizar cómo las redes neuronales realizan predicciones cuando los parámetros del modelo (anchura $N$, profundidad $L$) y el tamaño del conjunto de datos ($P$) son simultáneamente grandes. Un desafío fundamental surge porque los límites de $N, L, P \to \infty$ no necesariamente conmutan. Los marcos existentes suelen depender del límite de anchura infinita con profundidad fija (que produce Núcleos de Tangente Neuronal o Procesos Gaussianos de Redes Neuronales) o de regímenes de redes lineales específicos. Sin embargo, estos enfoques a menudo no logran capturar el papel de la profundidad en redes no lineales o la interacción entre el aprendizaje de características y los métodos de kernel en regímenes de escalado proporcional donde $L \propto N$ y $P \propto N$.

El problema específico abordado es la caracterización riguroza de la inferencia bayesiana en Perceptrones Multicapa (MLP) profundos y no lineales bajo un límite de escalado conjunto donde el número de muestras de entrenamiento ($P$), la dimensión de entrada ($N_0$), la anchura de la capa oculta ($N$) y el número de capas ocultas ($L$) divergen todos. Los autores buscan determinar:

1. Cómo se comporta el posterior predictivo en este régimen.
2. Bajo qué condiciones la profundidad proporciona un beneficio (aumenta la evidencia del modelo bayesiano).
3. Si las redes no lineales pueden aproximarse mediante métodos de kernel dependientes de los datos y de qué manera.

Metodología

1. Funciones de Activación con Forma (Shaped Activation Functions)

Para asegurar un comportamiento no degenerado a medida que $L, N \to \infty$, los autores emplean funciones de activación con forma $\phi_s$. Para una no linealidad base $\phi_{base}$, la activación se escala como $\phi_s(x) = s \phi_{base}(x/s)$. Este mecanismo asegura que el efecto acumulativo de las no linealidades a través de muchas capas permanezca controlable y comparable a los efectos lineales. El factor de forma $s$ se elige de tal manera que la razón $t = L/N$ controla de forma única la distribución de inicialización.

2. Ecuación Diferencial Estocástica de Covarianza Neuronal (NSDE)

La herramienta técnica central es la Ecuación Diferencial Estocástica de Covarianza Neuronal (NSDE). Los autores analizan la evolución del kernel conjugado (matriz de Gram de las características ocultas) $\Phi^{(\ell)}$ a través de las capas.

• En el límite $N, L \to \infty$ con $L/N \to t$, la evolución discreta del kernel converge a un proceso de difusión $\Phi_t$.
• El proceso está gobernado por:
  $$d\Phi_t = \alpha_{act} \cdot b(\Phi_t) dt + \Phi_t^{1/2} dB_t \Phi_t^{1/2}$$
  donde $b(\Phi_t)$ es un término de deriva dependiente de la no linealidad base, y $B_t$ es una matriz de movimiento browniano.
• Para manejar un $P$ grande, se introduce un cambio de tiempo $\tau = P t$, definiendo una "profundidad-tiempo" $\tau^* = LP/N$. La dinámica se reescala para equilibrar los términos de difusión y deriva.

3. Función de Partición Bayesiana vía SDE

Los autores reformulan la función de partición bayesiana $Z_\beta(x_0, \kappa)$ (que codifica el posterior predictivo) como una esperanza sobre la distribución a priori del kernel conjugado.

• Utilizando una identidad de integral gaussiana, la función de partición se expresa como:
  $$Z_\beta(x_0, \kappa) \propto \int_{\mathbb{R}^P} \exp\left(-\frac{1}{2\beta}\|p\|^2 + ip^\top Y\right) \mathbb{E}_{prior}\left[\exp\left(-\frac{1}{2}v^\top \Phi v\right)\right] dp$$
  donde $v = [p^\top, \kappa]^\top$ y $\Phi$ es el kernel conjugado en la penúltima capa.
• Mediante el Teorema 1, en el límite conjunto, la esperanza sobre el prior del MLP converge a la esperanza sobre la solución de la NSDE. Esto reduce el problema de la inferencia bayesiana en MLPs profundos a resolver la SDE para $\Phi_\tau$ y calcular la esperanza resultante.

4. Análisis Perturbativo

Dado que resolver la SDE para un $\tau^*$ arbitrario es inabordable, los autores realizan una expansión perturbativa para una profundidad efectiva pequeña $\tau = LP/N$. Calculan la derivada de primer orden de la función de partición logarítmica con respecto a $\tau$ en $\tau=0$. Esto les permite separar las contribuciones del componente de la red lineal (solo difusión) y el componente del kernel no lineal (solo deriva).

Contribuciones Clave y Resultados

1. Marco Unificador Riguroso

El artículo proporciona el primer marco matemáticamente riguroso para analizar la inferencia bayesiana en MLPs profundos no lineales bajo el límite proporcional $LP/N \in \Theta(1)$. Une la brecha entre los métodos de kernel de anchura infinita y las redes lineales profundas, cubriendo tanto funciones de activación suaves como ReLU a temperaturas arbitrarias ($\beta$).

2. Recuperación y Generalización de Resultados Previos

• Redes Lineales Profundas: El marco recupera resultados exactos para redes lineales profundas (donde la función de partición involucra integrales Meijer-G) como un caso especial donde el término de deriva desaparece.
• Límite de Anchura Infinita: Recupera el resultado de que, en el límite $N \gg LP$, la inferencia bayesiana se reduce a un método de kernel agnóstico a los datos con un mapa de características determinista derivado de la Neural ODE.
• Redes Débilmente No Lineales: Los resultados generalizan los hallazgos perturbativos de la literatura de física (por ejemplo, [HZ24]) sobre redes débilmente no lineales, pero con total rigor matemático y sin depender de simplificaciones no probadas.

3. Expansión Perturbativa de la Evidencia Bayesiana

Los autores derivan un criterio simple para determinar cuándo aumentar la profundidad (aumentar $\tau = LP/N$) mejora la evidencia del modelo bayesiano (verosimilitud marginal).

• Conclusión 2.1: El cambio de primer orden en la log-evidencia es la suma de un término de red lineal y un término de kernel no lineal.
• Conclusiones 3.1 y 3.2: Para distribuciones de datos específicas (normalizadas, entradas aproximadamente ortogonales), los autores derivan desigualdades explícitas que determinan cuándo la profundidad es beneficiosa.
  • Para activaciones suaves, la profundidad es beneficiosa si las estadísticas de los datos y las derivadas de la activación satisfacen una desigualdad específica que involucra la relación señal-ruido y la curvatura de la función de activación.
  • Para activaciones ReLU, la profundidad es beneficiosa si y solo si la varianza de las etiquetas excede un umbral relacionado con la asimetría de la ReLU ($c_+ - c_-$).
• Transición de Fase: Los resultados sugieren una transición de fase donde la profundidad solo es beneficiosa si la no linealidad no es "demasiado fuerte" en relación con la estructura de los datos.

4. Equivalencia con Métodos de Kernel Dependientes de los Datos

• Conclusión 2.2: Al primer orden en $\tau$, el posterior predictivo bayesiano para datos de prueba es equivalente a un método de kernel, pero con un mapa de características dependiente de los datos.
• A diferencia del límite de anchura infinita donde el mapa de características es fijo, el mapa de características efectivo aquí incluye un término de corrección proporcional a las etiquetas de los datos de entrenamiento y al componente ortogonal de la entrada de prueba. Esto proporciona una justificación teórica para el "aprendizaje de características" en el sentido bayesiano: la red aprende un mapa de características que se adapta ligeramente al conjunto de entrenamiento, distinto del mapa de características prior.

Significado y Reivindicaciones

Los autores afirman que su trabajo ofrece una alternativa matemáticamente rigurosa a las derivaciones inspiradas en la física que a menudo dependen de supuestos heurísticos. Al utilizar la NSDE y las expansiones perturbativas, ellos:

1. Clarifican el papel de la profundidad: Demuestran que la profundidad no es universalmente beneficiosa; su utilidad depende de la interacción entre el proceso generador de datos, la función de activación y la temperatura.
2. Unifican el Aprendizaje de Kernel y de Características: Muestran que incluso en el régimen profundo y no lineal, el posterior predictivo puede verse como un método de kernel, pero con un mapa de características ajustado dinámicamente por los datos de entrenamiento (a primer orden).
3. Proporcionan Herramientas Analíticas: El marco permite el cálculo de estadísticas bayesianas (evidencia, media posterior, varianza) en regímenes previamente considerados inabordables, específicamente para redes no lineales con anchura y profundidad finitas pero grandes.

El artículo establece explícitamente que no pretende resolver el caso general para un $\tau$ arbitrario ni proporcionar validación experimental en conjuntos de datos del mundo real. En su lugar, establece una base teórica para comprender los mecanismos de la inferencia bayesiana en redes profundas a través de la lente de las ecuaciones diferenciales estocásticas y los límites de escalado proporcional.