Scaling Up Diffusion and Flow-based XGBoost Models
Authors: name: Jesse C. Cresswell, name: Taewoo Kim
Organization: Layer 6 AI

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.