Scaling di Modelli XGBoost Basati su Diffusione e Flusso

Scaling Up Diffusion and Flow-based XGBoost Models

August 28, 2024

Autori: Jesse C. Cresswell, Taewoo Kim

cs.AI

Abstract

I nuovi metodi di machine learning per la generazione di dati tabellari sono spesso sviluppati su dataset di piccole dimensioni che non corrispondono alla scala richiesta per applicazioni scientifiche. Investigiamo una proposta recente di utilizzare XGBoost come approssimatore di funzione in modelli di diffusione e di flow-matching su dati tabellari, che si è rivelata estremamente dispendiosa in termini di memoria, anche su dataset minuscoli. In questo lavoro, conduciamo un'analisi critica dell'implementazione esistente da una prospettiva ingegneristica e dimostriamo che queste limitazioni non sono fondamentali per il metodo; con un'implementazione migliore, è possibile scalare a dataset 370 volte più grandi di quelli utilizzati in precedenza. La nostra implementazione efficiente permette inoltre di scalare i modelli a dimensioni molto maggiori, il che dimostriamo portare direttamente a prestazioni migliorate su task di benchmark. Proponiamo anche miglioramenti algoritmici che possono ulteriormente avvantaggiare l'uso delle risorse e le prestazioni del modello, inclusi alberi multi-output che sono particolarmente adatti per la modellazione generativa. Infine, presentiamo risultati su dataset scientifici su larga scala derivati dalla fisica delle particelle sperimentale, come parte della Fast Calorimeter Simulation Challenge. Il codice è disponibile all'indirizzo 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.