Escalando Modelos XGBoost Basados en Difusión y Flujo

Scaling Up Diffusion and Flow-based XGBoost Models

August 28, 2024

Autores: Jesse C. Cresswell, Taewoo Kim

cs.AI

Resumen

Los nuevos métodos de aprendizaje automático para la generación de datos tabulares a menudo se desarrollan en conjuntos de datos pequeños que no coinciden con la escala requerida para aplicaciones científicas. Investigamos una propuesta reciente de utilizar XGBoost como aproximador de funciones en modelos de difusión y emparejamiento de flujos (flow-matching) para datos tabulares, que resultó ser extremadamente intensivo en memoria, incluso en conjuntos de datos diminutos. En este trabajo, realizamos un análisis crítico de la implementación existente desde una perspectiva de ingeniería y demostramos que estas limitaciones no son fundamentales para el método; con una mejor implementación, puede escalarse a conjuntos de datos 370 veces más grandes que los utilizados previamente. Nuestra implementación eficiente también permite escalar modelos a tamaños mucho mayores, lo que demostramos que conduce directamente a un mejor rendimiento en tareas de referencia. También proponemos mejoras algorítmicas que pueden beneficiar aún más el uso de recursos y el rendimiento del modelo, incluyendo árboles de múltiples salidas (multi-output trees) que son adecuados para el modelado generativo. Finalmente, presentamos resultados en conjuntos de datos científicos a gran escala derivados de la física experimental de partículas como parte del Fast Calorimeter Simulation Challenge. El código está disponible en 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.