Escalonando Modelos XGBoost Baseados em Difusão e Fluxo

Scaling Up Diffusion and Flow-based XGBoost Models

August 28, 2024

Autores: Jesse C. Cresswell, Taewoo Kim

cs.AI

Resumo

Novos métodos de aprendizado de máquina para geração de dados tabulares são frequentemente desenvolvidos em conjuntos de dados pequenos que não correspondem à escala necessária para aplicações científicas. Investigamos uma proposta recente de usar o XGBoost como o aproximador de função em modelos de difusão e *flow-matching* para dados tabulares, que se mostrou extremamente intensivo em memória, mesmo em conjuntos de dados minúsculos. Neste trabalho, conduzimos uma análise crítica da implementação existente sob uma perspectiva de engenharia e mostramos que essas limitações não são fundamentais para o método; com uma implementação melhor, ele pode ser dimensionado para conjuntos de dados 370 vezes maiores do que os utilizados anteriormente. Nossa implementação eficiente também permite dimensionar os modelos para tamanhos muito maiores, o que mostramos levar diretamente a um desempenho melhorado em tarefas de referência. Também propomos melhorias algorítmicas que podem beneficiar ainda mais o uso de recursos e o desempenho do modelo, incluindo árvores de múltiplas saídas, que são bem adequadas para modelagem generativa. Por fim, apresentamos resultados em conjuntos de dados científicos de grande escala derivados da física experimental de partículas, como parte do *Fast Calorimeter Simulation Challenge*. O código está disponível em https://github.com/layer6ai-labs/calo-forest.

English

Novel machine learning methods for tabular data generation are often developed on small datasets which do not match the scale required for scientific applications. We investigate a recent proposal to use XGBoost as the function approximator in diffusion and flow-matching models on tabular data, which proved to be extremely memory intensive, even on tiny datasets. In this work, we conduct a critical analysis of the existing implementation from an engineering perspective, and show that these limitations are not fundamental to the method; with better implementation it can be scaled to datasets 370x larger than previously used. Our efficient implementation also unlocks scaling models to much larger sizes which we show directly leads to improved performance on benchmark tasks. We also propose algorithmic improvements that can further benefit resource usage and model performance, including multi-output trees which are well-suited to generative modeling. Finally, we present results on large-scale scientific datasets derived from experimental particle physics as part of the Fast Calorimeter Simulation Challenge. Code is available at https://github.com/layer6ai-labs/calo-forest.