Масштабирование диффузионных и потоковых моделей на основе XGBoost

Scaling Up Diffusion and Flow-based XGBoost Models

August 28, 2024

Авторы: Jesse C. Cresswell, Taewoo Kim

cs.AI

Аннотация

Новые методы машинного обучения для генерации табличных данных часто разрабатываются на небольших наборах данных, которые не соответствуют масштабу, необходимому для научных приложений. Мы исследуем недавнее предложение использовать XGBoost в качестве аппроксиматора функции в моделях диффузии и согласования потоков для табличных данных, что оказалось чрезвычайно ресурсоемким даже на небольших наборах данных. В данной работе мы проводим критический анализ существующей реализации с инженерной точки зрения и показываем, что эти ограничения не являются фундаментальными для метода; с более эффективной реализацией его можно масштабировать на наборы данных, в 370 раз превышающие ранее использовавшиеся. Наша оптимизированная реализация также позволяет масштабировать модели до значительно больших размеров, что, как мы показываем, напрямую приводит к улучшению производительности на тестовых задачах. Мы также предлагаем алгоритмические улучшения, которые могут дополнительно оптимизировать использование ресурсов и производительность модели, включая деревья с множественными выходами, которые хорошо подходят для задач генеративного моделирования. Наконец, мы представляем результаты на крупномасштабных научных наборах данных, полученных из экспериментальной физики элементарных частиц в рамках задачи Fast Calorimeter Simulation Challenge. Код доступен по адресу 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.