A Framework for Variational Inference of Lightweight Bayesian Neural Networks with Heteroscedastic Uncertainties

A continuación se presenta un resumen técnico detallado del artículo "Un marco para la inferencia variacional de redes neuronales bayesianas ligeras con incertidumbres heterocedásticas" de Schodt et al.

1. Planteamiento del Problema

Las Redes Neuronales Bayesianas (BNN) son esenciales para cuantificar la incertidumbre predictiva, distinguiendo entre incertidumbre epistémica (incertidumbre del modelo debido a datos limitados) e incertidumbre aleatoria (ruido de los datos).

• El Desafío: En muchas aplicaciones, particularmente aquellas que requieren incertidumbres heterocedásticas (donde los niveles de ruido varían con la entrada), los enfoques estándar requieren que la BNN outpute la varianza aleatoria como una cabeza de red adicional junto a la media predictiva.
• La Limitación: Añadir un nodo de salida para la varianza incrementa el número de parámetros aprendibles y la complejidad arquitectónica. Esto es perjudicial para las BNN ligeras destinadas a su despliegue en sistemas con restricciones de hardware o donde el tiempo de cómputo es limitado.
• La Brecha: Los métodos existentes a menudo fuerzan una compensación entre el tamaño del modelo y la capacidad de capturar estructuras de ruido complejas dependientes de la entrada.

2. Metodología

Los autores proponen un marco de inferencia novedoso que elimina la necesidad de un nodo de salida separado para la incertidumbre aleatoria, incrustando tanto las incertidumbres aleatorias como las epistémicas directamente en las varianzas de los parámetros aprendibles de la BNN.

A. Reinterpretación de las Distribuciones de Parámetros

En lugar de tratar la varianza de los parámetros de las neuronas bayesianas como representativa únicamente de la incertidumbre epistémica, los autores reinterpretan estas distribuciones para codificar la incertidumbre predictiva total (aleatoria + epistémica).

• Enfoque Estándar: $Var(\hat{y}) = \sigma^2_{epistémica}$. La varianza aleatoria ($\sigma^2_{aleatoria}$) se predice como una salida separada.
• Enfoque Propuesto: La varianza total $\sigma^2_k$ para una entrada dada $k$ se define como la suma de las varianzas aleatoria y epistémica:
  $$\sigma^2_k = \sigma^2_{a,k} + \sigma^2_{e,k}$$
  Durante el entrenamiento, la red minimiza la Cota Inferior de la Evidencia (ELBO), donde la pérdida de reconstrucción asume que la distribución de salida es Gaussiana con varianza $\sigma^2_k$. La varianza $\sigma^2_k$ se calcula como la varianza de la salida de la BNN derivada de las distribuciones de parámetros, "absorbiendo" efectivamente el ruido aleatorio en las varianzas de los parámetros.

B. Inferencia Variacional Sin Muestreo mediante Propagación de Momentos

Para mantener la eficiencia computacional y evitar el costo del muestreo de Monte Carlo, los autores emplean un esquema de propagación de momentos.

• Mecanismo: La media y la varianza de la entrada a cada capa se propagan analíticamente a través de la red para calcular la media y la varianza de la salida.
• Manejo de Capas:
  • Totalmente Conectadas y Convolucionales: Propagación analítica exacta de medias y varianzas asumiendo independencia entre neuronas.
  • Leaky-ReLU: Propagación aproximada asumiendo que la entrada a la función de activación sigue una distribución normal. Los autores derivan expresiones en forma cerrada para la media y la varianza de la salida utilizando la ley de la esperanza total y la varianza, involucrando distribuciones normales truncadas.
• Beneficio: Esto permite una inferencia determinista sin necesidad de muestreo repetido, haciendo que el marco sea adecuado para aplicaciones ligeras y en tiempo real.

3. Contribuciones Clave

1. Incrustación de Varianza de Parámetros: Un cambio conceptual donde las varianzas de los pesos y sesgos de la BNN se entrenan para representar la suma de las incertidumbres aleatoria y epistémica, eliminando la necesidad de una cabeza de salida dedicada a la varianza.
2. Arquitectura Ligera: El marco permite que las BNN capturen incertidumbres heterocedásticas con menos parámetros aprendibles en comparación con los enfoques estándar de "varianza aprendida".
3. Inferencia Sin Muestreo: Integración de una técnica simplificada de propagación de momentos que proporciona estimaciones analíticas de incertidumbre, evitando la sobrecarga computacional del dropout de Monte Carlo o el muestreo.
4. Marco Unificado: Un pipeline completo que combina la reinterpretación de las fuentes de incertidumbre con una propagación de momentos eficiente para el entrenamiento y la inferencia.

4. Resultados Experimentales

Los autores evaluaron su enfoque de "Varianza Incrustada" frente al enfoque estándar de "Varianza Aprendida" (donde la varianza es una salida separada) utilizando una tarea de regresión polinómica con ruido Gaussiano heterocedástico.

• Rendimiento en Regímenes Ligeros:
  • Con un número pequeño de neuronas ocultas (por ejemplo, 4), el modelo de Varianza Incrustada predijo con precisión tanto la tendencia como la incertidumbre para datos dentro de la distribución y expresó correctamente una alta incertidumbre para datos fuera de la distribución.
  • El modelo de Varianza Aprendida falló en recuperar la incertidumbre real para datos dentro de la distribución y subestimó significativamente la incertidumbre para datos fuera de la distribución con la misma arquitectura pequeña.
• Eficiencia de Parámetros:
  • Para un número fijo de neuronas, el modelo de "Varianza Aprendida" requirió aproximadamente un 34% más de parámetros aprendibles que el modelo de "Varianza Incrustada" para lograr un rendimiento comparable.
  • El enfoque de "Varianza Incrustada" logró una pérdida de reconstrucción menor (log-verosimilitud negativa) en el régimen de bajo número de parámetros.
• Generalización:
  • Los modelos de "Varianza Incrustada" demostraron una mejor generalización a datos fuera de la distribución, mostrando un aumento más gradual en la pérdida de reconstrucción a medida que los datos se alejaban de la distribución de entrenamiento en comparación con los modelos de "Varianza Aprendida".
• Escalabilidad: A medida que aumentaba el tamaño de la red (por ejemplo, 1024 neuronas), la brecha de rendimiento se estrechó y ambos métodos funcionaron de manera similar, lo que sugiere que el método propuesto es más crítico para modelos ligeros con restricciones de recursos.

5. Significado

Este trabajo aborda un cuello de botella crítico en el despliegue de Redes Neuronales Bayesianas para la computación en el borde y sistemas críticos para la seguridad.

• Eficiencia: Al eliminar la necesidad de nodos de salida adicionales para predecir el ruido, el método reduce significativamente el tamaño del modelo y el costo computacional, haciendo viables las BNN para hardware con restricciones estrictas.
• Robustez: La capacidad de capturar incertidumbres heterocedásticas sin aumentar la complejidad del modelo asegura que los modelos ligeros no sacrifiquen la fiabilidad o las evaluaciones de seguridad.
• Practicidad: La naturaleza sin muestreo de la inferencia permite tiempos de predicción más rápidos, lo cual es esencial para aplicaciones en tiempo real donde la latencia es una preocupación.

En resumen, el artículo demuestra que la incertidumbre total puede ser efectivamente "incrustada" en las distribuciones de parámetros de una BNN, ofreciendo una alternativa superior a las cabezas de salida múltiples tradicionales para la cuantificación de incertidumbre heterocedástica ligera.