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

AI summary

Overview

• Research area: Machine learning, specifically exact (globally optimal) decision-tree learning for regression, formulated through mixed-integer programming and reduced-space branch-and-bound.
• Technical level: Advanced. The paper relies on two-stage stochastic-style optimization, branch-and-bound theory, and mixed-integer formulations, though the core ideas are described here in plain language.
• Scope: The paper introduces RS-ORT, a branch-and-bound algorithm that trains globally optimal regression trees directly on continuous features by branching only on tree-structure variables, and benchmarks it against MIP and heuristic baselines on ten regression datasets.

What This Paper Is About

Existing mixed-integer programming methods for learning optimal decision trees work well for classification but struggle with regression: they either require purely binary features or become computationally intractable on large, continuous datasets. Naively binarizing continuous features sacrifices global optimality and tends to produce unnecessarily deep trees. The paper recasts optimal regression-tree training as a two-stage optimization problem and solves it with a branch-and-bound algorithm whose search space does not grow with the number of training samples, allowing it to handle datasets with continuous features at the million-row scale within a four-hour time limit.

Key Contributions

1. A two-stage formulation that handles continuous features directly. The first stage encodes only tree structure (split features, thresholds, leaf assignments, split indicators), and the second stage captures sample-specific routing and loss. Because the algorithm branches exclusively on the first-stage variables, the search space — and memory footprint — is independent of the number of training samples, a property the authors state is not shared by previous MIP formulations for regression trees.

2. A convergence guarantee for branching on structural variables only. The authors prove (Theorem 2) that the reduced-space branch-and-bound algorithm converges to the global optimum when branching only on the tree-structure variables, even though the second stage contains both binary and continuous decisions, extending the argument of Hua et al. (2022) from classification trees to regression trees.

3. Three bound-tightening techniques tailored to regression. Closed-form leaf predictions (each leaf's optimal value is the empirical mean of its assigned targets, removing all leaf prediction variables from the search space); empirical threshold discretization (restricting each split threshold to observed feature values without loss of generality); and exact depth-1 subtree parsing (resolving terminal subtrees by fitting a depth-one CART model once the upper tree structure is fixed).

4. A parallelizable implementation that scales to multi-million-row data. The decomposable lower and upper bounds permit trivial parallel execution across branch-and-bound nodes, which the authors use to solve datasets with over 2 million rows and continuous features.

Main Findings

• Combinatorial explosion motivates pruning. Theorem 1 bounds the number of distinct tree structures at at most P^(2^D − 1) × ∏_{t=0}^{D−1} (⌊n/2^t⌋ − 1)^(2^t). For the Concrete dataset (n = 1,030, P = 8) at depth D = 2, this bound evaluates to 8³ · 1029 · 514² = 139,191,134,208, approximately 1.39 × 10^11 structures.

• Bound tightening produces negligible gaps on most datasets. Across Concrete, Insurance, Airfoil Noise, Abalone, Seoul Bike, CC Power Plant, Steel Industry, and Micro Gas Turbine, RS-ORT reported optimality gaps below 0.01%. On CPU ACT it reported a gap of 8.08%, and on Household 33.47%.

• RS-ORT matched or beat exact baselines on training and test RMSE. On most datasets (Concrete, Insurance, Airfoil Noise, Abalone, Seoul Bike, Micro Gas Turbine) RS-ORT's RMSE matched the Bertsimas MIQP formulation. On CPU ACT it achieved Train RMSE 5.99 versus 6.01 for Bertsimas, and on CC Power Plant 6.34 versus 6.35. On Steel Industry it reached Train RMSE 7.66 and Test RMSE 7.61, versus 8.42 and 8.33 for both Bertsimas and CART.

• The MIP baseline often failed to close its gap. Bertsimas's formulation reported optimality gaps of 50.90% (Airfoil Noise), 99.98% (Abalone), 100.00% (CPU ACT, Seoul Bike, Steel Industry, Micro Gas Turbine), and 98.90% (CC Power Plant) after hitting the 14,400-second limit. It produced out-of-memory errors on Household.

• The current state-of-the-art exact method hit resource limits. OSRT (Zhang et al., 2023) ran out of memory on CPU ACT, CC Power Plant, and Steel Industry, and timed out without producing a valid tree on Micro Gas Turbine and Household. It was faster than RS-ORT on several small datasets — for example 4.63 s on Airfoil Noise, 26.06 s on Insurance, and 115.67 s on Concrete.

• Scaling to 2 million rows. RS-ORT solved the Household dataset (n = 2,049,280, P = 5), with Train and Test RMSE of 0.29 and a runtime of 14,400 s; CART reached 0.30 but Bertsimas ran out of memory and OSRT timed out without a tree. The authors state this is the first exact method capable of handling continuous features at this scale, and the abstract reports that on datasets of up to 2,000,000 samples with continuous features RS-ORT obtains guaranteed training performance with a simpler tree structure within four hours, reporting optimality gaps below 0.1 (Table 1 lists 33.47% for Household).

• CART was competitive in accuracy but has no guarantees. CART matched other methods on many datasets but is a heuristic; on Concrete it produced Test RMSE 12.57 versus 11.80 for the exact methods. Because CART offers no optimality guarantee, no gap is reported for it.

• Continuous data is where the advantage concentrates. The paper attributes OSRT's out-of-memory and timeout failures to its binary encoding of dense, high-cardinality continuous features, which inflates the number of binary columns.

Methodology in Plain English

The authors start from the standard way of writing an optimal regression tree as a single large optimization problem: binary variables decide which feature each internal node splits on and whether the node is active, continuous variables hold the split thresholds and leaf predictions, and binary variables assign each training sample to a leaf. The problem is NP-complete even for binary classification, and the sample-to-leaf assignment variables grow with the dataset size.

Their reformulation splits the problem in two. The first stage contains only the tree-structure variables. The second stage handles everything sample-specific — which leaf each sample lands in and the resulting squared error. The branch-and-bound procedure then branches only on the first-stage variables, partitioning the space of possible tree structures. For each region, it computes a lower bound by relaxing the requirement that all samples share one tree (letting each sample use its own tree decouples the problem across samples and makes the bound easy to parallelize), and an upper bound from feasible solutions. A proof shows this branching scheme still converges to the global optimum.

Three tricks make this practical. First, since a leaf's best prediction under squared error is just the mean of the targets assigned to it, leaf predictions need never be treated as decision variables. Second, because any threshold between two consecutive observed feature values produces the same split, thresholds can be restricted to observed values, turning an uncountable variable domain into a finite one; the implementation then bisects the remaining sorted values. Third, once everything above a parent node is fixed, the decision to split that parent or keep it as a leaf depends only on the samples reaching it, so the whole bottom layer can be resolved exactly by fitting a depth-one CART model locally and comparing the gain to a regularization penalty. This prunes branches, tightens upper bounds, and preserves optimality.

Experiments compared RS-ORT against the MIQP formulation of Bertsimas and Dunn (2019) solved with CPLEX 22.1.2, CART from scikit-learn v1.7.1, and the official OSRT Python implementation, on ten public regression datasets with a 70/30 train/test split, a regularization parameter of 0.0005, tree depth fixed at 2 except where noted, and a four-hour wall-clock limit per dataset. RS-ORT was written in Julia and run across 40 to 200 physical cores for most datasets and up to 1,000 cores for datasets over 100,000 observations; the baselines were limited to the largest single-node machines available (20 threads, 128 GB RAM).

Why This Matters

Research impact. The paper pushes exact optimal regression-tree training from binary-feature, small-scale settings toward continuous, million-row datasets. It shows that the reduced-space branch-and-bound idea, previously applied to optimal classification trees, transfers to regression, and it contributes regression-specific bound tightening where classification-style sample reduction does not apply. This gives the optimal-tree community a new baseline and a formulation whose memory footprint is decoupled from sample count.

Real-world applications. The benchmark suite spans several concrete domains where interpretability and accuracy both matter:

• Civil engineering and materials, via the Concrete compressive-strength dataset.
• Power generation and energy monitoring, via the CC Power Plant and Micro Gas Turbine datasets.
• Urban mobility and demand forecasting, via the Seoul Bike dataset.
• Industrial process control, via the Steel Industry dataset, and household-level prediction, via the 2-million-row Household dataset.

Industry relevance. In high-stakes settings such as medicine, where transparent predictions can be validated and trusted by practitioners, an exactly optimal and compact tree is more defensible than a heuristically grown one. Because RS-ORT runs on distributed CPU cores and its nodes decompose independently, it fits existing cluster infrastructure. Its main practical cost is runtime: several benchmarks consumed the full 14,400-second budget, and the largest case still left a 33.47% gap at termination.

Future Directions

• Closing the gap on large datasets. The Household dataset terminated with a 33.47% optimality gap on its regularized training objective, and CPU ACT with 8.08%, so stronger bounds or better branching rules for large, continuous problems remain an open target.
• Reducing wall-clock time relative to specialized baselines. OSRT solved Concrete in 115.67 s, Insurance in 26.06 s, and Airfoil Noise in 4.63 s, while RS-ORT took 631.00 s, 87.20 s, and 196.00 s respectively on those datasets, so the reduced-space framework's overhead on small, binarizable data is worth investigating.
• Extending beyond depth 2 and squared error. The experiments fix depth at 2 except where noted, and the closed-form leaf prediction result depends on squared-error loss; other loss functions and deeper trees could be examined.
• Quantifying the benefit of parallelism and the depth advantage. The paper reports using between 40 and 1,000 cores but does not isolate how much of the speedup comes from parallel execution, and the abstract mentions trees often two to three levels shallower than competing models without reporting a per-dataset depth comparison in the truncated content.

Target Audience

This paper is for researchers and practitioners in optimization and machine learning who work on exact decision-tree learning, mixed-integer programming formulations, or interpretable regression models. It will be most useful to readers already comfortable with branch-and-bound methods and MIP formulations who want to understand how reduced-space techniques and regression-specific bound tightening extend to continuous, large-scale data. Data scientists in regulated domains who need regression trees with provable optimality guarantees rather than heuristic ones are the secondary audience.

Authors’ 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.