Mise à l'échelle des modèles XGBoost basés sur la diffusion et les flux

Scaling Up Diffusion and Flow-based XGBoost Models

August 28, 2024

papers.authors: Jesse C. Cresswell, Taewoo Kim

cs.AI

papers.abstract

Les nouvelles méthodes d'apprentissage automatique pour la génération de données tabulaires sont souvent développées sur de petits jeux de données qui ne correspondent pas à l'échelle requise pour les applications scientifiques. Nous étudions une proposition récente visant à utiliser XGBoost comme approximateur de fonction dans les modèles de diffusion et d'appariement de flux sur des données tabulaires, qui s'est avérée extrêmement gourmande en mémoire, même sur des jeux de données minuscules. Dans ce travail, nous menons une analyse critique de l'implémentation existante d'un point de vue ingénierie, et montrons que ces limitations ne sont pas fondamentales à la méthode ; avec une meilleure implémentation, elle peut être mise à l'échelle sur des jeux de données 370 fois plus grands que ceux utilisés précédemment. Notre implémentation efficace permet également de passer à l'échelle des modèles de tailles bien supérieures, ce qui, comme nous le démontrons, améliore directement les performances sur des tâches de référence. Nous proposons également des améliorations algorithmiques pouvant further bénéficier à l'utilisation des ressources et aux performances du modèle, notamment des arbres à sorties multiples bien adaptés à la modélisation générative. Enfin, nous présentons des résultats sur des jeux de données scientifiques à grande échelle issus de la physique des particules expérimentale dans le cadre du Fast Calorimeter Simulation Challenge. Le code est disponible à l'adresse 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.