Resumen Técnico: Inferencia Variacional de Posterior Basada en Puntuación a Gran Escala para Redes Neuronales Profundas Bayesianas

Enunciado del Problema

Las Redes Neuronales Profundas Bayesianas (BNN) ofrecen ventajas significativas sobre el aprendizaje profundo de estimación puntual, incluyendo una cuantificación robusta de la incertidumbre, resistencia al ruido y mitigación del sobreajuste. Sin embargo, la inferencia de posterior exacta en BNN es generalmente intratable. La Inferencia Variacional (IV) es un método aproximado dominante que formula la inferencia como un problema de optimización al aproximar el posterior objetivo $\pi(\theta)$ con una distribución tratable $q_\lambda(\theta)$.

El enfoque dominante actual, la IV basada en ELBO (minimizando la divergencia KL inversa), sufre de un conocido problema de "colapso de modos", donde la aproximación falla al capturar todos los modos de la densidad objetivo. Si bien los métodos de IV basados en puntuación (que minimizan la diferencia entre la puntuación objetivo $\nabla_\theta \log \pi(\theta)$ y la puntuación del modelo $\nabla_\theta \log q_\lambda(\theta)$) pueden teóricamente abordar el colapso de modos, los enfoques basados en puntuación existentes no son escalables a BNN a gran escala. Específicamente:

1. Intratabilidad Computacional: Métodos como la Coincidencia de Puntuación Gaussiana (GSM) y Lote y Coincidencia (BaM) dependen de inversiones o multiplicaciones de matrices que escalan cúbica o cuadráticamente con la dimensión del estado, haciéndolos prohibitivos para los parámetros de redes neuronales de alta dimensión.
2. Incompatibilidad con Gradientes Estocásticos: Los métodos basados en puntuación existentes típicamente requieren puntuaciones exactas o no pueden manejar puntuaciones de mini-lotes ruidosas y no sesgadas derivadas de gradientes estocásticos, las cuales son esenciales para el entrenamiento en grandes conjuntos de datos.
3. Restricciones de Reparametrización: Muchos métodos dependen del muestreo reparametrizado, lo cual puede ser computacionalmente costoso o inestable en altas dimensiones.

Metodología

Los autores proponen un nuevo algoritmo de IV Basada en Puntuación con Gradiente Estocástico Proximal diseñado para escalar a BNN a gran escala, incluidos los Transformadores de Visión (ViT) y ResNets.

Objetivo de Optimización Central

El método actualiza iterativamente la distribución variacional $q_t(\theta)$ a $q_{t+1}(\theta)$ resolviendo el siguiente problema de optimización en cada iteración $t$:

$$q_{t+1} = \arg \min_{q \in \mathcal{Q}} \mathbb{E}_{\theta \sim q_t(\theta)} \left[ \alpha_t \|\nabla_\theta \log q(\theta) - \nabla_\theta \log q_t(\theta)\|_2^2 + \|\nabla_\theta \log q(\theta) - s(\theta)\|_2^2 \right]$$

Donde:

• $s(\theta) = \nabla_\theta \log \pi(\theta)$ es la puntuación objetivo.
• $\alpha_t$ es un parámetro de impacto proximal adaptativo (programado linealmente de 0 a una constante $\le 1$).
• El primer término es una penalización proximal que obliga a la nueva iteración a permanecer cerca de la iteración actual en términos de puntuaciones.
• El segundo término es la pérdida de coincidencia de puntuación.

Innovaciones Técnicas Clave

1. Manejo de Puntuaciones de Mini-lotes Ruidosas: A diferencia de los métodos anteriores basados en puntuación que requieren puntuaciones exactas, este enfoque utiliza puntuaciones de mini-lotes ruidosas y no sesgadas $\hat{s}(\theta)$ obtenidas mediante gradientes estocásticos. La linealidad del término de puntuación en la función de pérdida asegura que el gradiente de la pérdida con respecto a los parámetros variacionales $\lambda$ permanezca una estimación no sesgada, incluso con puntuaciones ruidosas.
2. Evitación de Reparametrización: El término proximal permite que la esperanza se tome con respecto a la distribución actual fija $q_t$ en lugar de la variable $q$. Esto elude la necesidad del truco de reparametrización (utilizado en ADVI) y evita la diferenciación doble con respecto a $\theta$ (requerida por los métodos tradicionales de Fisher/Puntuación), reduciendo significativamente la sobrecarga computacional y la huella de memoria.
3. Escalabilidad: Al utilizar normas euclidianas en lugar de normas inducidas por covarianza (que requieren inversión de matrices) y soportar gradientes estocásticos, el método escala a redes con cientos de millones de parámetros (por ejemplo, ViT-L-32 con ~305M de parámetros).

Convergencia Teórica

Los autores proporcionan un teorema de convergencia (Teorema 3.1) que establece que, bajo la suposición de que la familia variacional es lo suficientemente rica para representar campos vectoriales suaves, la puntuación de la distribución variacional converge a la puntuación objetivo a medida que aumentan las iteraciones. Aunque la suposición de riqueza representacional es ideal, los autores notan que se cumple exactamente si tanto el posterior verdadero como la familia variacional son Gaussianos.

Contribuciones Clave

1. Algoritmo Novel: Introducción de un algoritmo de IV de coincidencia de puntuación con gradiente estocástico proximal que combina la pérdida de coincidencia de puntuación con una penalización proximal.
2. Escalabilidad a BNN a Gran Escala: El método es el primer enfoque de IV basado en puntuación demostrado para funcionar efectivamente en redes profundas a gran escala (ResNets, ViTs) y grandes conjuntos de datos, superando los cuellos de botella computacionales de GSM y BaM.
3. Robustez al Ruido: El algoritmo maneja efectivamente puntuaciones de mini-lotes ruidosas y no sesgadas, permitiendo el uso de gradientes estocásticos sin requerir cálculos de puntuación de lote completo.
4. Eliminación de Reparametrización: El método evita el truco de reparametrización, simplificando el paisaje de optimización y mejorando la eficiencia en comparación con la IV basada en ELBO (ADVI).

Resultados Experimentales

Los autores evaluaron el método en datos sintéticos, el benchmark PosteriorDB y tareas a gran escala del mundo real:

• Datos Sintéticos: En objetivos Gaussianos y mezclas de Gaussianas, el método propuesto demostró una convergencia más rápida y una mejor robustez ante la elección inicial de parámetros (específicamente en relación con pequeños valores propios de covarianza) en comparación con ADVI, GSM y BaM. Manejó con éxito el desajuste del modelo y espacios de estado de alta dimensión donde GSM/BaM fallaron debido a problemas numéricos.
• MNIST (BNN a pequeña escala): En regímenes de pocos datos, el método logró tasas de error de prueba comparables a ADVI, pero demostró una cuantificación de incertidumbre superior, evidenciada por un menor Error de Calibración Esperado (ECE) y una Negativa de Verosimilitud Logarítmica (NLL) más baja.
• Reconocimiento Visual a Gran Escala (ResNet-101 y ViT-L-32): En conjuntos de datos como Pets, Flowers, Aircraft y DTD, el método igualó la precisión de predicción de prueba de ADVI mientras lo superaba consistentemente en métricas de incertidumbre (ECE y NLL). Los métodos competidores (GSM, BaM, coincidencia de puntuación tradicional) fallaron debido a restricciones de memoria o explosiones numéricas.
• Predicción de Series Temporales: Utilizando la arquitectura de red Koopa, el método superó a la IV basada en ELBO y a las líneas base no bayesianas en MSE, MAE y NLL en múltiples conjuntos de datos (ECL, ETTh2, Ex. Rate, ILI, Traffic, Weather).
• Flujos Normalizantes: El método se extendió a densidades variacionales de Flujos Normalizantes, mostrando mejoras consistentes sobre ADVI en tareas de series temporales.
• Eficiencia: A pesar de mantener tanto las iteraciones anteriores como las actuales, el tiempo de reloj y la huella de memoria de la GPU fueron comparables a ADVI, demostrando su viabilidad práctica para el entrenamiento a gran escala.

Significado y Afirmaciones

El artículo afirma que este trabajo proporciona una alternativa escalable y flexible a la IV basada en ELBO dominante para el aprendizaje profundo bayesiano a gran escala. Al abordar el problema de colapso de modos inherente a los objetivos de KL inversa y superar las limitaciones computacionales de los métodos anteriores basados en puntuación, el algoritmo propuesto permite una inferencia bayesiana rigurosa en arquitecturas modernas de alta capacidad como los Transformadores de Visión.

Los autores enfatizan que su enfoque es empíricamente efectivo en diversos benchmarks, ofreciendo una mejor cuantificación de la incertidumbre sin sacrificar el rendimiento predictivo. Notan modestamente que, si bien la convergencia teórica depende de suposiciones fuertes sobre la expresividad de la familia variacional, los resultados empíricos apoyan la utilidad del método en la práctica. El trabajo posiciona la IV basada en puntuación como una vía viable y escalable para el aprendizaje profundo bayesiano, avanzando más allá de las restricciones de solo Gaussianas o requisitos de lote completo de la tecnología anterior.