RS-ORT: A Reduced-Space Branch-and-Bound Algorithm for Optimal Regression Trees

TL;DR

RS-ORT introduces a novel branch-and-bound algorithm for optimal regression trees that efficiently handles large-scale continuous data, guaranteeing global optimality and superior performance compared to existing methods.

Contribution

The paper proposes RS-ORT, a reduced-space branch-and-bound algorithm that exclusively branches on tree-structural variables, ensuring convergence and scalability for large datasets.

Findings

RS-ORT outperforms state-of-the-art methods on regression benchmarks.

It guarantees training performance with simpler trees on large datasets.

The algorithm scales to datasets with up to 2 million samples within four hours.

Abstract

Mixed-integer programming (MIP) has emerged as a powerful framework for learning optimal decision trees. Yet, existing MIP approaches for regression tasks are either limited to purely binary features or become computationally intractable when continuous, large-scale data are involved. Naively binarizing continuous features sacrifices global optimality and often yields needlessly deep trees. We recast the optimal regression-tree training as a two-stage optimization problem and propose Reduced-Space Optimal Regression Trees (RS-ORT) - a specialized branch-and-bound (BB) algorithm that branches exclusively on tree-structural variables. This design guarantees the algorithm's convergence and its independence from the number of training samples. Leveraging the model's structure, we introduce several bound tightening techniques - closed-form leaf prediction, empirical threshold discretization, and exact depth-1 subtree parsing - that combine with decomposable upper and lower bounding strategies to accelerate the training. The BB node-wise decomposition enables trivial parallel execution, further alleviating the computational intractability even for million-size datasets. Based on the empirical studies on several regression benchmarks containing both binary and continuous features, RS-ORT also delivers superior training and testing performance than state-of-the-art methods. Notably, on datasets with up to 2,000,000 samples with continuous features, RS-ORT can obtain guaranteed training performance with a simpler tree structure and a better generalization ability in four hours.