Flexible Modeling and Multitask Learning using Differentiable Tree Ensembles

TL;DR

This paper introduces a flexible, differentiable tree ensemble framework supporting arbitrary loss functions, missing data, and multi-task learning, with efficient GPU training and significantly more compact models.

Contribution

It presents a novel tensor-based formulation for scalable differentiable trees enabling broader modeling capabilities beyond existing toolkits.

Findings

Achieved up to 100x more compact tree ensembles.

Demonstrated 23% increase in model expressiveness.

Validated on 28 diverse datasets.

Abstract

Decision tree ensembles are widely used and competitive learning models. Despite their success, popular toolkits for learning tree ensembles have limited modeling capabilities. For instance, these toolkits support a limited number of loss functions and are restricted to single task learning. We propose a flexible framework for learning tree ensembles, which goes beyond existing toolkits to support arbitrary loss functions, missing responses, and multi-task learning. Our framework builds on differentiable (a.k.a. soft) tree ensembles, which can be trained using first-order methods. However, unlike classical trees, differentiable trees are difficult to scale. We therefore propose a novel tensor-based formulation of differentiable trees that allows for efficient vectorization on GPUs. We perform experiments on a collection of 28 real open-source and proprietary datasets, which demonstrate that our framework can lead to 100x more compact and 23% more expressive tree ensembles than those by popular toolkits.