Scaling Up Diffusion and Flow-based XGBoost Models
Jesse C. Cresswell, Taewoo Kim
2024-08-30
Summary
This paper discusses a new approach to using XGBoost, a popular machine learning model, for generating data from large datasets more efficiently.
What's the problem?
Many new machine learning methods are tested on small datasets, which isn't enough for real-world scientific applications. Current implementations of XGBoost can be very memory-intensive, making it hard to use on larger datasets, which limits their effectiveness.
What's the solution?
The authors analyze the existing implementation of XGBoost and show that the memory limitations can be overcome with better engineering. They present a new implementation that can handle datasets that are 370 times larger than those previously used. They also introduce algorithmic improvements, like using multi-output trees, which help the model perform better while using fewer resources. Their results demonstrate improved performance on large-scale scientific datasets, particularly in particle physics.
Why it matters?
This research is important because it makes it possible to apply machine learning techniques to much larger datasets, which is crucial for fields like scientific research where large amounts of data are common. By improving the efficiency and performance of XGBoost, it opens up new possibilities for data analysis and modeling in various scientific disciplines.
Abstract
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.