Resumen Técnico: Agudeza de la Geometría de Fisher y el Sesgo Implícito de SGD hacia Mínimos Planos

1. Planteamiento del Problema

La intuición predominante en el aprendizaje profundo postula que el Descenso de Gradiente Estocástico (SGD) favorece implícitamente los "mínimos planos" en el paisaje de pérdida, y que estos mínimos planos generalizan mejor que los "agudos". Sin embargo, la base teórica de esta narrativa ha sido críticamente socavada por Dinh et al. [2017], quienes demostraron que las medidas euclidianas estándar de agudeza (por ejemplo, la traza o el autovalor máximo de la Hessiana de la pérdida) no son invariantes ante la reparametrización.

Específicamente, para una red neuronal, se pueden reescalar los pesos (por ejemplo, $W_1 \to \alpha W_1, W_2 \to \alpha^{-1} W_2$) para preservar exactamente la función de entrada-salida de la red mientras se altera arbitrariamente la agudeza euclidiana. Dado que la generalización es una propiedad de la función, no de la parametrización específica, una medida de agudeza no invariante no puede explicar de manera significativa la generalización. Aunque se han tomado pasos parciales hacia medidas invariantes (por ejemplo, Jang et al. [2022]; Kristiadi et al. [2023]), ha permanecido ausente un marco unificado que conecte la agudeza invariante, el sesgo implícito de SGD y los límites de generalización.

2. Metodología

2.1 Agudeza Riemanniana ($S_R$)

Los autores reemplazan la métrica euclidiana en el espacio de parámetros con la Matriz de Información de Fisher (FIM), $G(\theta)$, que induce una estructura Riemanniana intrínseca a la distribución predictiva del modelo. Definen la agudeza Riemanniana en un mínimo $\theta^*$ como:
$$S_R(\theta^*) = \text{tr}\left( G(\theta^*)^{-1} \nabla^2 L(\theta^*) \right)$$
donde $\nabla^2 L$ es la Hessiana de la pérdida de entrenamiento. A diferencia de la agudeza euclidiana, $S_R$ mide la curvatura relativa a la geometría natural de la variedad estadística.

2.2 Marco Teórico: SGD como una Ecuación Diferencial Estocástica (SDE)

Los autores formalizan el SGD de mini-lotes como una SDE de tiempo continuo. Introducen la Asunción 3, postulando que la covarianza del ruido del gradiente es proporcional a la FIM:
$$\text{Cov}[\nabla \tilde{L}(\theta^*)] = \frac{\eta}{B} G(\theta^*)$$
donde $\eta$ es la tasa de aprendizaje y $B$ es el tamaño del lote. Bajo esta asunción, derivan la distribución estacionaria $\pi^*$ del proceso SGD.

2.3 Límites de Generalización

Los autores conectan el sesgo geomético con el rendimiento de prueba utilizando un marco PAC-Bayes. Construyen una distribución posterior centrada en el mínimo entrenado y derivan un límite de generalización controlado explícitamente por $S_R$.

2.4 Estimación Empírica

Dado que computar la FIM completa es intratable, los autores utilizan un estimador de FIM empírico diagonal (muestreando predicciones de la propia distribución del modelo) y el estimador de traza estocástica de Hutchinson para la Hessiana. Reconocen que, si bien la verdadera FIM garantiza la invariancia, el estimador diagonal empírico utilizado en los experimentos hereda esta invariancia solo aproximadamente.

3. Contribuciones Clave

1. Definición de Agudeza Invariante a la Reparametrización: El artículo define $S_R$ y demuestra (Lema 1) que es invariante bajo reparametrizaciones suaves que preservan la función, abordando directamente la crítica de Dinh et al. [2017].
2. Análisis de la Distribución Estacionaria: Los autores derivan la distribución estacionaria del SGD de mini-lotes (Teorema 1) y demuestran que la masa de probabilidad está exponencialmente concentrada en mínimos con baja agudeza Riemanniana (Corolario 1). Esto establece un sesgo implícito hacia los mínimos planos Riemannianos.
3. Límite de Generalización PAC-Bayes: Se establece un límite de generalización (Corolario 2) donde la brecha de generalización está controlada explícitamente por $S_R$, proporcionando un vínculo teórico entre la planitud Riemanniana y el rendimiento de prueba.
4. Validación Empírica: Experimentos en MNIST y CIFAR-10 confirman que $S_R$ rastrea la generalización a través de diferentes optimizadores (SGD, Adam), tamaños de lote y tasas de aprendizaje, de formas en que la agudeza euclidiana ($S_E$) no lo hace. El escalamiento de $S_R$ con $\eta/B$ coincide con las predicciones teóricas.

4. Resultados

4.1 Invariancia y Alineación

• Invariancia: Si bien la verdadera FIM es estrictamente invariante, el estimador diagonal empírico no es perfectamente invariante bajo un reescalado arbitrario. Sin embargo, sigue siendo sustancialmente más estable que la agudeza euclidiana bajo el reescalado de pesos.
• Estructura del Ruido: Las pruebas empíricas (Sección 5) muestran que la covarianza del ruido del gradiente y la FIM están direccionalmente alineadas (similitud de coseno $\approx 1.0$) a través de las configuraciones. Sin embargo, la constante de proporcionalidad escalar no es exactamente $\eta/B$ en todos los regímenes, particularmente en entornos de lotes pequeños y alto ruido. Los autores refinan su asunción para permitir un factor escalar $\alpha(\theta)$, demostrando que las principales conclusiones teóricas (el sesgo hacia la planitud) siguen siendo robustas incluso si la afirmación escalar es imperfecta.

4.2 Leyes de Escalamiento y Generalización

• Tamaño de Lote ($B$): En MNIST, aumentar $B$ conduce a un aumento significativo de $S_R$ y una disminución correspondiente en la precisión de prueba, consistente con la predicción $S_R \propto B/\eta$. En CIFAR-10, la relación es más compleja debido a las dificultades de optimización en regímenes de lotes pequeños, pero se mantiene la tendencia de que un mayor $B$ conduce a un mayor $S_R$.
• Tasa de Aprendizaje ($\eta$): Aumentar $\eta$ generalmente disminuye $S_R$ y mejora la generalización (hasta un punto donde el entrenamiento se vuelve inestable).
• Optimizadores:
  • SGD: Encuentra consistentemente mínimos con menor $S_R$ y mayor precisión en configuraciones de lotes pequeños.
  • Adam: Logra un bajo $S_R$ y alta precisión en MNIST. En CIFAR-10, Adam encuentra mínimos con un $S_R$ menor que SGD pero con una precisión ligeramente inferior. Los autores sugieren que esto se debe a que el precondicionamiento adaptativo de Adam cambia la métrica efectiva; el sesgo implícito de los optimizadores precondicionados es hacia la planitud relativa a su propia métrica tensorial.

4.3 Dinámicas de Entrenamiento

Las trayectorias de entrenamiento muestran que el SGD de lotes pequeños escapa progresivamente de los pozos agudos y se establece en regiones más planas (disminuyendo $S_R$) a medida que el entrenamiento continúa. El SGD de lotes grandes se estabiliza en valores de $S_R$ más altos, consistente con un ruido insuficiente para escapar de los mínimos agudos.

5. Significación y Reivindicaciones

El artículo afirma proporcionar la primera cuenta rigurosa e invariante a la reparametrización de por qué los mínimos planos generalizan. Al fundamentar la agudeza en la geometría Riemanniana de la variedad estadística (vía la FIM), los autores resuelven la objeción fundamental de que la planitud es meramente un artefacto de la parametrización.

La significación del trabajo radica en:

• Resolución Teórica: Une la intuición de "planitud-generalización" con la "crítica de la reparametrización" definiendo una medida de agudeza que es intrínseca a la función del modelo.
• Mecanismo de Sesgo Implícito: Proporciona una derivación matemática que muestra cómo la estructura de ruido de SGD (proporcional a la FIM) impulsa naturalmente al optimizador hacia mínimos con baja agudeza Riemanniana.
• Métrica Práctica: Demuestra que $S_R$ es un predictor más fiable de la generalización que la agudeza euclidiana, particularmente al comparar modelos entrenados con diferentes hiperparámetros u optimizadores.

Limitaciones Reconocidas por los Autores:

• El estimador diagonal empírico de la FIM no es perfectamente invariante a la reparametrización; la invariancia exacta requeriría estimadores estructurados como K-FAC.
• La afirmación escalar de la covarianza del ruido ($\alpha = \eta/B$) es una aproximación que se mantiene direccionalmente pero varía en magnitud dependiendo del conjunto de datos y la arquitectura.
• La relación entre $S_R$ y la generalización es más informativa cuando el optimizador ha convergido; en regímenes inestables (por ejemplo, tasas de aprendizaje muy altas), un bajo $S_R$ no garantiza una buena generalización.