bde: A Python Package for Bayesian Deep Ensembles via MILE

Resumen Técnico de "bde: Un Paquete Python para Ensembles Profundos Bayesianos vía MILE"

Enunciado del Problema
La cuantificación fiable de la incertidumbre (UQ) es un componente crítico de los sistemas modernos de aprendizaje automático. Si bien el Aprendizaje Profundo Bayesiano (BDL) ofrece un marco principiado para lograrlo, las herramientas accesibles para la inferencia basada en muestreo en redes neuronales bayesianas siguen siendo escasas. Los marcos de programación probabilística existentes suelen soportar la Cadena de Markov Monte Carlo (MCMC), pero requieren una configuración manual sustancial para lograr un rendimiento competitivo en modelos de redes neuronales. Además, existe una falta de implementaciones amigables para el usuario que conecten la investigación de alto rendimiento en MCMC con flujos de trabajo prácticos de ciencia de datos, particularmente para tareas de aprendizaje supervisado con datos tabulares.

Metodología: El Enfoque MILE
El artículo introduce bde, un paquete Python que implementa Ensembles Langevin Microcanónicos (MILE), una técnica de muestreo híbrida diseñada para Ensembles Profundos Bayesianos (BDEs). La metodología sigue un proceso de inferencia en dos etapas:

1. Etapa de Optimización: El sistema optimiza $n_{members}$ instancias independientes de una red neuronal de alimentación hacia adelante configurada de manera flexible. Esto se logra mediante la minimización del riesgo empírico regularizado (utilizando la verosimilitud negativa como función de pérdida) con el optimizador AdamW. El objetivo es encontrar modos de alta verosimilitud diversos.
2. Etapa de Muestreo: El proceso transiciona a una fase de muestreo utilizando Monte Carlo Langevin Microcanónico (MCMC). Esta fase se mejora con una fase de ajuste adaptada específicamente para Redes Neuronales Bayesianas.

La combinación de estas etapas permite que el ensemble localice primero modos diversos y posteriormente explore estructuras locales del posterior, generando un ensemble de muestras que constituyen una aproximación implícita del posterior.

Diseño e Implementación del Software
El paquete bde está diseñado para aprovechar la velocidad y las capacidades de alto rendimiento de JAX y blackjax, combinados con la API amigable para el usuario de scikit-learn. Las características clave de diseño incluyen:

• Paralelización: Dado que la optimización y el muestreo a través de los miembros del ensemble son independientes, el paquete explota la paralelización de JAX y la compilación just-in-time para escalar eficientemente a través de CPUs, GPUs y TPUs.
• Interfaz: Proporciona una interfaz unificada para predicciones puntuales, intervalos de credibilidad, estimaciones de cobertura y otras métricas de incertidumbre.
• Flujo de trabajo: El paquete automatiza la orquestación de la optimización, el muestreo, la paralelización y la inferencia predictiva, no requiriendo que el usuario posea conocimientos especializados sobre dinámicas de MCMC, estrategias de inicialización o internos de JAX.

Contribuciones Clave

• Primera Implementación Amigable de MILE: bde proporciona la primera implementación accesible de MILE, una técnica que ha demostrado ofrecer una alta precisión predictiva y una incertidumbre calibrada para redes neuronales bayesianas.
• Compatibilidad con scikit-learn: Al adherirse a la API de scikit-learn, el paquete permite una integración fluida en flujos de trabajo existentes de aprendizaje automático.
• Enfoque en Datos Tabulares: Se dirige específicamente a tareas de aprendizaje supervisado con datos tabulares, incluyendo regresión distribucional y clasificación (multiclase).
• Reproducibilidad: El paquete incluye configuraciones automatizadas y diseños experimentales para garantizar una UQ reproducible en el aprendizaje automático con datos tabulares.

Resultados
Los autores evaluaron bde en los conjuntos de datos airfoil y bikesharing, comparándolo con Modelos Lineales, Bosques Aleatorios, XGBoost, CatBoost, Ensembles Profundos y TabPFN (V2).

• Rendimiento: En el conjunto de datos airfoil, bde logró un RMSE de $0.1215 \pm 0.0042$, superando a la mayoría de los competidores, incluido TabPFN ($0.1359$) y Ensembles Profundos ($0.3650$). En el conjunto de datos bikesharing, logró un RMSE de $0.2261 \pm 0.0016$.
• Cuantificación de la Incertidumbre: Los resultados destacan un rendimiento competitivo fuera de la caja en UQ, particularmente a través de su capacidad nativa de regresión distribucional. bde demostró un fuerte rendimiento en las métricas de Verosimilitud Negativa (NLL), alcanzando $-0.9126$ (regresión distribucional) en airfoil y $-0.7315$ en bikesharing.
• Costo Computacional: Los autores reconocen que, si bien los BDEs altamente optimizados incurren en costos computacionales más altos que los competidores basados en optimización debido a la naturaleza iterativa del muestreo MCMC, esta inversión permite la construcción de una aproximación posterior flexible y no paramétrica.

Significado y Afirmaciones
El artículo posiciona a bde como un puente entre la investigación de alto rendimiento en MCMC y la ciencia de datos práctica. Su inclusión en la organización scikit-learn-contrib garantiza el cumplimiento de estándares de software rigurosos y consistencia en la API, convirtiéndolo en una herramienta confiable y lista para la comunidad. Los autores afirman que bde permite a investigadores y profesionales aplicar fácilmente Ensembles Profundos Bayesianos para datos tabulares sin necesidad de conocimientos especializados sobre las dinámicas complejas subyacentes, facilitando así la cuantificación de la incertidumbre reproducible en el aprendizaje automático.