bartz

BART vetorizado

Uma implementação de Árvores de Regressão Aditiva Bayesiana (BART) em JAX.

Se você não sabe o que é BART, mas conhece o XGBoost, considere o BART como uma espécie de XGBoost bayesiano. bartz faz o BART rodar tão rápido quanto o XGBoost.

BART é uma técnica de regressão Bayesiana não paramétrica. Dados os preditores de treinamento $X$ e respostas $y$ , o BART encontra uma função para prever $y$ dado $X$ . O resultado da inferência é uma amostra de funções possíveis, representando a incerteza sobre a determinação da função.

Este módulo Python fornece uma implementação de BART que roda em GPU, para processar grandes conjuntos de dados com mais rapidez. Também é bom na CPU. A maioria das outras implementações do BART são para R e são executadas apenas na CPU.

Na CPU, bartz roda na velocidade de dbarts (a implementação mais rápida que conheço) se n > 20.000, mas usando 1/20 da memória. Na GPU, o prêmio de velocidade depende do tamanho da amostra; é conveniente na CPU apenas para n > 10.000. A aceleração máxima é atualmente de 200x, em uma Nvidia A100 e com pelo menos 2.000.000 de observações.

Este notebook Colab executa bartz com n = 100.000 observações, p = 1.000 preditores, 10.000 árvores, para 1.000 iterações MCMC, em 5 minutos.

Ligações

• Documentação (última versão)
• Documentação (versão de desenvolvimento)
• Repositório
• Cobertura de código
• Lista de pacotes BART

Citando Bartz

Artigo: Petrillo (2024), "Árvores de regressão aditiva bayesiana muito rápidas em GPU", arXiv:2410.23244.

Para citar o software diretamente, incluindo a versão específica, use zenodo.