Opschalen van op diffusie en stroming gebaseerde XGBoost-modellen

Scaling Up Diffusion and Flow-based XGBoost Models

August 28, 2024

Auteurs: Jesse C. Cresswell, Taewoo Kim

cs.AI

Samenvatting

Nieuwe machine learning-methoden voor het genereren van tabulaire gegevens worden vaak ontwikkeld op kleine datasets die niet overeenkomen met de schaal die nodig is voor wetenschappelijke toepassingen. Wij onderzoeken een recent voorstel om XGBoost te gebruiken als de functiebenadering in diffusie- en flow-matchingmodellen voor tabulaire gegevens, wat extreem geheugenintensief bleek te zijn, zelfs op minuscule datasets. In dit werk voeren we een kritische analyse uit van de bestaande implementatie vanuit een technisch perspectief en tonen we aan dat deze beperkingen niet fundamenteel zijn voor de methode; met een betere implementatie kan deze worden opgeschaald naar datasets die 370 keer groter zijn dan voorheen gebruikt. Onze efficiënte implementatie maakt het ook mogelijk modellen op te schalen naar veel grotere omvang, wat direct leidt tot verbeterde prestaties op benchmarktaken, zoals we aantonen. We stellen ook algoritmische verbeteringen voor die het resourcegebruik en de modelprestaties verder kunnen bevorderen, waaronder multi-output bomen die bijzonder geschikt zijn voor generatief modelleren. Tot slot presenteren we resultaten op grootschalige wetenschappelijke datasets afkomstig uit experimentele deeltjesfysica als onderdeel van de Fast Calorimeter Simulation Challenge. Code is beschikbaar op 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.