bde: A Python Package for Bayesian Deep Ensembles via MILE

Resumo Técnico de "bde: Um Pacote Python para Ensembles Bayesianos Profundos via MILE"

Declaração do Problema
A quantificação de incerteza (UQ) confiável é um componente crítico dos sistemas modernos de aprendizado de máquina. Embora o Aprendizado Profundo Bayesiano (BDL) ofereça um framework principiado para alcançar isso, ferramentas acessíveis para inferência baseada em amostragem em redes neurais bayesianas permanecem escassas. Frameworks de programação probabilística existentes geralmente suportam Cadeias de Markov Monte Carlo (MCMC), mas exigem configuração manual substancial para alcançar desempenho competitivo em modelos de redes neurais. Além disso, há uma falta de implementações amigáveis ao usuário que conectem pesquisas de MCMC de alto desempenho com fluxos de trabalho práticos de ciência de dados, particularmente para tarefas de aprendizado supervisionado em dados tabulares.

Metodologia: A Abordagem MILE
O artigo apresenta o bde, um pacote Python que implementa Ensembles Langevin Microcanônicos (MILE), uma técnica de amostragem híbrida projetada para Ensembles Bayesianos Profundos (BDEs). A metodologia segue um processo de inferência em duas etapas:

1. Etapa de Otimização: O sistema otimiza $n_{members}$ instâncias independentes de uma rede neural feed-forward configurável de forma flexível. Isso é alcançado por meio de minimização de risco empírico regularizada (usando o log-verossimilhança negativa como função de perda) com o otimizador AdamW. O objetivo é encontrar modos diversos de alta verossimilhança.
2. Etapa de Amostragem: O processo transita para uma fase de amostragem usando Monte Carlo Langevin Microcanônico (MCMC). Esta fase é aprimorada com uma fase de ajuste especificamente adaptada para Redes Neurais Bayesianas.

A combinação dessas etapas permite que o ensemble localize primeiro modos diversos e, subsequentemente, explore estruturas locais do posterior, gerando um ensemble de amostras que constitui uma aproximação implícita do posterior.

Design e Implementação de Software
O pacote bde é projetado para aproveitar a velocidade e as capacidades de alto desempenho do JAX e do blackjax, combinados com a API amigável ao usuário do scikit-learn. As principais características de design incluem:

• Paralelização: Como a otimização e a amostragem entre os membros do ensemble são independentes, o pacote explora o paralelismo do JAX e a compilação just-in-time para escalar eficientemente através de CPUs, GPUs e TPUs.
• Interface: Fornece uma interface unificada para previsões pontuais, intervalos de credibilidade, estimativas de cobertura e outras métricas de incerteza.
• Fluxo de Trabalho: O pacote automatiza a orquestração de otimização, amostragem, paralelização e inferência preditiva, não exigindo conhecimento especializado de dinâmicas de MCMC, estratégias de inicialização ou internos do JAX por parte do usuário.

Principais Contribuições

• Primeira Implementação Amigável ao Usuário de MILE: O bde fornece a primeira implementação acessível de MILE, uma técnica demonstrada como capaz de entregar alta precisão preditiva e incerteza calibrada para redes neurais bayesianas.
• Compatibilidade com Scikit-learn: Ao aderir à API do scikit-learn, o pacote permite integração perfeita em fluxos de trabalho existentes de aprendizado de máquina.
• Foco em Dados Tabulares: Visa especificamente tarefas de aprendizado supervisionado em dados tabulares, incluindo regressão distribucional e classificação (multiclasse).
• Reprodutibilidade: O pacote inclui configurações automatizadas e setups experimentais para garantir UQ reprodutível em aprendizado de máquina tabular.

Resultados
Os autores avaliaram o bde nos conjuntos de dados airfoil e bikesharing, comparando-o com Modelos Lineares, Random Forests, XGBoost, CatBoost, Ensembles Profundos e TabPFN (V2).

• Desempenho: No conjunto de dados airfoil, o bde alcançou um RMSE de $0.1215 \pm 0.0042$, superando a maioria dos concorrentes, incluindo TabPFN ($0.1359$) e Ensembles Profundos ($0.3650$). No conjunto de dados bikesharing, alcançou um RMSE de $0.2261 \pm 0.0016$.
• Quantificação de Incerteza: Os resultados destacam desempenho competitivo pronto para uso em UQ, particularmente através de sua capacidade nativa de regressão distribucional. O bde demonstrou forte desempenho nas métricas de Log-Verossimilhança Negativa (NLL), alcançando $-0.9126$ (regressão distribucional) no airfoil e $-0.7315$ no bikesharing.
• Custo Computacional: Os autores reconhecem que, embora BDEs altamente otimizados incorram em custos computacionais mais altos do que concorrentes baseados em otimização devido à natureza iterativa da amostragem MCMC, esse investimento permite a construção de uma aproximação de posterior flexível e não paramétrica.

Significado e Afirmações
O artigo posiciona o bde como uma ponte entre pesquisas de MCMC de alto desempenho e ciência de dados prática. Sua inclusão na organização scikit-learn-contrib garante a adesão a padrões rigorosos de software e consistência de API, tornando-o uma ferramenta confiável e pronta para a comunidade. Os autores afirmam que o bde permite que pesquisadores e praticantes apliquem facilmente Ensembles Bayesianos Profundos para dados tabulares sem necessidade de conhecimento especializado das dinâmicas complexas subjacentes, facilitando assim a quantificação de incerteza reprodutível em aprendizado de máquina.