Scaling Up Diffusion and Flow-based XGBoost Models
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.
66d123e173bc1d219a6a559a
Jesse C. Cresswell
admin_assigned
2024-08-30T09:09:43.553Z
65adf6fbae773a8112f4b9ad
/avatars/38d2eb81bf1f9b8c4a1903d21c48f365.svg
Jesse Cresswell
user
JesseCresswell
66d123e173bc1d219a6a559b
Taewoo Kim
admin_assigned
2024-08-30T09:10:01.584Z
666d972cf98a92bc89418c73
/avatars/1d81926bce5a8760e9e79321fada62f2.svg
Taewoo Kim
user
TaewooKim
Scaling Up Diffusion and Flow-based XGBoost Models
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.
Improved implementation and algorithmic enhancements of XGBoost in diffusion and flow-matching models enable scaling to larger datasets and superior performance on scientific benchmarks.
XGBoost diffusion models flow-matching models multi-output trees scientific datasets Fast Calorimeter Simulation Challenge