# Scaling Up Diffusion and Flow-based XGBoost Models

arXiv: https://arxiv.org/abs/2408.16046 · HTML: https://mlsanity.com/paper/2408.16046/

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Builds on 33 of its 101 references that are in the map.
- First posted on arXiv 28 Aug 2024.

## Builds on (most cited first, 30 of 33)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Semi-Supervised Classification with Graph Convolutional Networks](https://mlsanity.com/paper/1609.02907/) (2016)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [A Unified Approach to Interpreting Model Predictions](https://mlsanity.com/paper/1705.07874/) (2017)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Generative Modeling by Estimating Gradients of the Data Distribution](https://mlsanity.com/paper/1907.05600/) (2019)
- [Improved Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2102.09672/) (2021)
- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [Building Normalizing Flows with Stochastic Interpolants](https://mlsanity.com/paper/2209.15571/) (2022)
- [Improving and generalizing flow-based generative models with minibatch optimal transport](https://mlsanity.com/paper/2302.00482/) (2023)
- [Revisiting Deep Learning Models for Tabular Data](https://mlsanity.com/paper/2106.11959/) (2021)
- [Modeling Tabular data using Conditional GAN](https://mlsanity.com/paper/1907.00503/) (2019)
- [TabDDPM: Modelling Tabular Data with Diffusion Models](https://mlsanity.com/paper/2209.15421/) (2022)
- [Consistent Individualized Feature Attribution for Tree Ensembles](https://mlsanity.com/paper/1802.03888/) (2018)
- [When Do Neural Nets Outperform Boosted Trees on Tabular Data?](https://mlsanity.com/paper/2305.02997/) (2023)
- [Exposing flaws of generative model evaluation metrics and their unfair treatment of diffusion models](https://mlsanity.com/paper/2306.04675/) (2023)
- [On the Importance of Noise Scheduling for Diffusion Models](https://mlsanity.com/paper/2301.10972/) (2023)
- [Missing Data Imputation using Optimal Transport](https://mlsanity.com/paper/2002.03860/) (2020)
- [STaSy: Score-based Tabular data Synthesis](https://mlsanity.com/paper/2210.04018/) (2022)
- [TabR: Tabular Deep Learning Meets Nearest Neighbors in 2023](https://mlsanity.com/paper/2307.14338/) (2023)
- [Generating and Imputing Tabular Data via Diffusion and Flow-based Gradient-Boosted Trees](https://mlsanity.com/paper/2309.09968/) (2023)
- [Deep Generative Models through the Lens of the Manifold Hypothesis: A Survey and New Connections](https://mlsanity.com/paper/2404.02954/) (2024)
- [CaloFlow: Fast and Accurate Generation of Calorimeter Showers with Normalizing Flows](https://mlsanity.com/paper/2106.05285/) (2021)
- [CaloFlow II: Even Faster and Still Accurate Generation of Calorimeter Showers with Normalizing Flows](https://mlsanity.com/paper/2110.11377/) (2021)
- [CaloMan: Fast generation of calorimeter showers with density estimation on learned manifolds](https://mlsanity.com/paper/2211.15380/) (2022)
- [Inductive Simulation of Calorimeter Showers with Normalizing Flows](https://mlsanity.com/paper/2305.11934/) (2023)
- [Attention to Mean-Fields for Particle Cloud Generation](https://mlsanity.com/paper/2305.15254/) (2023)
- [Generative Trees: Adversarial and Copycat](https://mlsanity.com/paper/2201.11205/) (2022)
- [SketchBoost: Fast Gradient Boosted Decision Tree for Multioutput Problems](https://mlsanity.com/paper/2211.12858/) (2022)

## Built on by (most cited first, 5 of 5)

- [A Survey on Deep Learning Approaches for Tabular Data Generation: Utility, Alignment, Fidelity, Privacy, Diversity, and Beyond](https://mlsanity.com/paper/2503.05954/) (2025)
- [Diffusion and Flow Matching Models for Tabular Data: A Survey](https://mlsanity.com/paper/2502.17119/) (2025)
- [Even Faster Simulations with Flow Matching: A Study of Zero Degree Calorimeter Responses](https://mlsanity.com/paper/2507.18811/) (2025)
- [A universal vision transformer for fast calorimeter simulations](https://mlsanity.com/paper/2601.05289/) (2026)
- [BitHEP -- The Limits of Low-Precision ML in HEP](https://mlsanity.com/paper/2504.03387/) (2025)

## Abstract (arXiv)

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.