Skalierung von Diffusions- und Flow-basierten XGBoost-Modellen

Scaling Up Diffusion and Flow-based XGBoost Models

August 28, 2024

papers.authors: Jesse C. Cresswell, Taewoo Kim

cs.AI

papers.abstract

Neuartige maschinelle Lernmethoden zur Generierung tabellarischer Daten werden oft an kleinen Datensätzen entwickelt, die nicht die für wissenschaftliche Anwendungen erforderliche Größe erreichen. Wir untersuchen einen kürzlich vorgeschlagenen Ansatz, XGBoost als Funktionsapproximator in Diffusions- und Flow-Matching-Modellen für tabellarische Daten zu verwenden, der sich selbst bei winzigen Datensätzen als extrem speicherintensiv erwiesen hat. In dieser Arbeit führen wir eine kritische Analyse der bestehenden Implementierung aus einer ingenieurtechnischen Perspektive durch und zeigen, dass diese Einschränkungen nicht grundlegend für die Methode sind; mit einer besseren Implementierung kann sie auf Datensätze skaliert werden, die 370-mal größer sind als bisher verwendet. Unsere effiziente Implementierung ermöglicht es auch, Modelle auf viel größere Größen zu skalieren, was nachweislich direkt zu einer verbesserten Leistung bei Benchmark-Aufgaben führt. Wir schlagen außerdem algorithmische Verbesserungen vor, die die Ressourcennutzung und Modellleistung weiter steigern können, darunter Multi-Output-Bäume, die sich gut für die generative Modellierung eignen. Schließlich präsentieren wir Ergebnisse zu groß angelegten wissenschaftlichen Datensätzen aus der experimentellen Teilchenphysik im Rahmen der Fast Calorimeter Simulation Challenge. Der Code ist verfügbar unter 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.