Resumen Técnico: ¿Es suficiente la última capa para la cuantificación de la incertidumbre?

Planteamiento del problema

La adopción de Redes Neuronales Profundas (DNN) en entornos de misión crítica se ve obstaculizada por la falta de una cuantificación fiable de la incertidumbre epistémica (UQ). Aunque los métodos bayesianos ofrecen un marco natural para modelar la incertidumbre mediante la distribución posterior predictiva, los enfoques predominantes como los Ensembles Profundos (DE) o las Redes Neuronales Bayesianas (BNN) completas suelen incurrir en costes computacionales prohibitivos. Una aproximación común para reducir esta carga es linealizar la DNN alrededor de sus parámetros para formar un Modelo Lineal Generalizado (GLM) Bayesiano. Sin embargo, un supuesto prevalente en el campo es que linealizar únicamente la última capa conectada (la "aproximación de la última capa") degrada significativamente el rendimiento de la UQ en comparación con la linealización de toda la red (linealización completa). Este artículo investiga si este compromiso entre eficiencia computacional y fidelidad de la UQ es necesario.

Metodología

Los autores emplean un enfoque dual que combina el análisis teórico mediante la Teoría de Matrices Aleatorias (RMT) y la evaluación empírica a gran escala.

Análisis Teórico

El estudio compara dos GLM Bayesianos derivados de una DNN inicializada aleatoriamente:

1. LL-GLM (GLM de la última capa): Formado al linealizar únicamente los parámetros de la última capa conectada. Esto induce el Kernel Conjugado (CK).
2. DNN-GLM (GLM de linealización completa): Formado al linealizar todos los parámetros de la red. Esto induce el Kernel de Tangente Neuronal (NTK).

Los autores analizan la Energía Libre Bayesiana (BFE), definida como el logaritmo negativo de la verosimilitud marginal, que sirve como una medida clave de la calidad del modelo y la capacidad de UQ. Derivan la BFE límite bajo un régimen de doble asíntota donde tanto el ancho de la red como el número de puntos de datos de entrenamiento tienden al infinito proporcionalmente. Este análisis utiliza herramientas de RMT para caracterizar los espectros deterministas límite de las matrices Gram de CK y NTK.

Evaluación Empírica

Para validar los hallazgos teóricos en entornos prácticos, los autores introducen LinearSampling, un paquete de PyTorch ligero. Este marco proporciona un algoritmo para muestrear de la posterior predictiva tanto de LL-GLMs como de DNN-GLMs. El algoritmo emplea una estrategia de "muestrear-luego-optimizar", inicializando los parámetros cerca de los pesos entrenados y utilizando el Descenso de Gradiente (Estocástico) para resolver el problema de optimización linealizado. Este enfoque escala a modelos grandes, incluyendo aquellos con millones de parámetros.

La evaluación empírica abarca tareas de regresión, clasificación y modelado de lenguaje, comparando LL-GLMs contra DNN-GLMs y otros métodos de estado del arte (por ejemplo, Deep Ensembles, SWAG, MC-Dropout, LLA). Las métricas incluyen el Error de Calibración Esperado (ECE), la Log-Verosimilitud Negativa Gaussiana (NLL) y métricas de incertidumbre especializadas como VARROC (AUC-ROC usando la varianza de las predicciones softmax) tanto para la detección de datos en distribución (in-distribution) como fuera de distribución (OOD).

Contribuciones Clave

1. Caracterización Teórica de la BFE Límite: Los autores derivan la B la BFE límite para procesos gaussianos (GPs) empleando kernels CK y NTK bajo supuestos leves sobre los componentes de la red y las distribuciones de datos. Caracterizan este límite como la solución única de un sistema de ecuaciones de punto fijo implícitas.
2. Análisis de Robustez y Descenso: Teóricamente, demuestran que ambos kernels exhiben robustez ante perturbaciones de entrada y fenómenos de descenso (incluyendo comportamientos de doble y triple descenso). Identifican que el NTK solo muestra una ventaja teórica en el régimen altamente sobre-muestreado (donde los puntos de datos exceden significativamente el ancho de la red) en la inicialización, pero esta ventaja desaparece tras el entrenamiento.
3. Validación Empírica a Gran Escala: A través de experimentos extensos en diversas tareas, el artículo demuestra que los LL-GLMs igualan consistentemente el rendimiento de UQ de los DNN-GLMs completos. En muchos casos, el LL-GLM se comporta de manera comparable a los Deep Ensembles, que son el estándar de oro.
4. Marco de Muestreo Escalable: La publicación de LinearSampling, un algoritmo escalable que permite el muestreo eficiente de la posterior de GLMs Bayesianos para grandes modelos de lenguaje y arquitecturas de visión profundas, superando los cuellos de botella computacionales previos.

Resultados

• Hallazgos Teóricos: El análisis revela que no hay una mejora significativa en las capacidades de UQ (medidas por BFE) para la linealización completa sobre la linealización de la última capa en la mayoría de los regímenes. El único régimen donde el NTK completo supera teóricamente al CK es cuando el número de puntos de datos es significativamente mayor que el ancho de la red, un escenario que los autores sugieren puede ser un artefacto de la inicialización aleatoria más que una limitación fundamental, ya que la brecha desaparece después del entrenamiento.
• Hallazgos Empíricos:
  • Regresión: En conjuntos de datos de regresión UCI, los LL-GLMs lograron un rendimiento comparable o incluso superior a los DNN-GLMs en términos de RMSE, NLL y ECE.
  • Clasificación: En tareas de clasificación de imágenes (MNIST, FashionMNIST, CIFAR-10/100, ImageNet), los LL-GLMs igualaron el rendimiento de UQ de los DNN-GLMs. Ambos métodos capturaron con éxito la varianza "consciente de la distancia", aumentando la incertidumbre para puntos alejados de los datos de entrenamiento.
  • Modelado de Lenguaje: En un modelo GPT-2 ajustado para el análisis de sentimiento, el LL-GLM logró puntuaciones VARROC-ID estadísticamente indistinguibles de las del DNN-GLM.
  • Eficiencia: Crucialmente, el LL-GLM logró estos resultados con costos computacionales y de memoria sustancialmente menores. En el experimento de GPT-2, el LL-GLM requirió 10 veces menos tiempo de ejecución (wall-time) y 3.5 veces menos memoria que el DNN-GLM.

Significado y Conclusión

El artículo concluye que la aproximación de la última capa es suficiente para una cuantificación de la incertidumbre de alta calidad. Los autores argumentan que este hallazgo sugiere que la incertidumbre epistémica en las DNN surge principalmente de la incertidumbre en cómo se utilizan las características de la red (pesos de la última capa) en lugar de la incertidumbre en las características mismas (las representaciones aprendidas por las capas anteriores).

Este resultado desafía el supuesto de que la linealización de toda la red es necesaria para una UQ precisa. Al validar que el enfoque de la última capa, computacionalmente más económico, ofrece un rendimiento comparable, este trabajo aboga por la adopción más amplia de la linealización de la última capa en el aprendizaje profundo bayesiano, particularmente para modelos a gran escala donde la inferencia bayesiana de la red completa es inviable. Los autores señalan las limitaciones respecto a la conexión teórica con redes entrenadas bajo supuestos realistas y el enfoque actual en la incertidumbre epistémica en lugar de la aleatoria, identificándolos como vías para trabajos futuros.