[Paper Review] MT-GBM: A Multi-Task Gradient Boosting Machine with Shared Decision Trees

ZhenZhe Ying, Zhuoer Xu+3 more|arXiv (Cornell University)|Jan 17, 2022

Advanced Neural Network Applications4 citations

TL;DR

MT-GBM is a novel multi-task gradient boosting machine that learns shared decision tree structures across multiple tasks by jointly optimizing multi-task losses, enabling efficient and effective multi-task learning in tabular data. It improves performance on main tasks by up to 23% in RMSE and 9% in MAPE compared to single-task GBDT models, with significant gains in fraud detection tasks using shared representations.

ABSTRACT

Despite the success of deep learning in computer vision and natural language processing, Gradient Boosted Decision Tree (GBDT) is yet one of the most powerful tools for applications with tabular data such as e-commerce and FinTech. However, applying GBDT to multi-task learning is still a challenge. Unlike deep models that can jointly learn a shared latent representation across multiple tasks, GBDT can hardly learn a shared tree structure. In this paper, we propose Multi-task Gradient Boosting Machine (MT-GBM), a GBDT-based method for multi-task learning. The MT-GBM can find the shared tree structures and split branches according to multi-task losses. First, it assigns multiple outputs to each leaf node. Next, it computes the gradient corresponding to each output (task). Then, we also propose an algorithm to combine the gradients of all tasks and update the tree. Finally, we apply MT-GBM to LightGBM. Experiments show that our MT-GBM improves the performance of the main task significantly, which means the proposed MT-GBM is efficient and effective.

Motivation & Objective

• To address the challenge of applying Gradient Boosted Decision Trees (GBDT) to multi-task learning, where shared representations are essential but difficult to learn.
• To enable GBDT models to jointly optimize multiple loss functions by sharing tree structures across tasks, improving generalization and efficiency.
• To stabilize training across tasks with varying residual error scales by adaptively adjusting learning rates per task.
• To develop a method that integrates seamlessly into existing GBDT frameworks like LightGBM with minimal computational overhead.
• To demonstrate that shared tree structures can significantly improve performance on both main and auxiliary tasks in real-world tabular data applications.

Proposed method

• Assign multiple outputs to each leaf node, enabling each tree to predict for multiple tasks simultaneously.
• Compute task-specific gradients for each output (task) during the boosting process.
• Combine gradients from all tasks using a weighted aggregation strategy that accounts for residual error scales to stabilize training.
• Update tree structure and leaf values using the combined gradient, preserving shared splits across tasks.
• Integrate the method into LightGBM by modifying its training loop to support multi-task learning with shared trees.
• Apply adaptive learning rate adjustment per task based on the magnitude of residual errors to balance convergence across tasks.

Experimental results

Research questions

• RQ1 Can shared decision tree structures be effectively learned in gradient boosting for multi-task learning in tabular data?
• RQ2 How can multi-task gradients be combined in GBDT to stabilize training when tasks have different residual error magnitudes?
• RQ3 Does learning shared representations across tasks improve performance on the main task compared to single-task GBDT models?
• RQ4 Can MT-GBM be efficiently implemented within existing GBDT frameworks like LightGBM without significant computational overhead?
• RQ5 How does the inclusion of low-correlation or high-correlation auxiliary tasks affect the performance of the main task in MT-GBM?

Key findings

• MT-GBM reduced the main task's RMSE from 400.32 to 308.21 and MAPE from 4.3% to 3.9% on the first financial dataset, showing a 23% improvement in RMSE and 9% in MAPE.
• On the fraud detection dataset, MT-GBM with four auxiliary tasks achieved a 3-fold average ROC-AUC of 0.9454, outperforming LightGBM (0.9369) and XGBoost (0.9443).
• The 4-task MT-GBM model achieved the highest performance, indicating that combining multiple related subtasks enhances robustness and predictive power.
• Even with low-correlation tasks (e.g., task 2 with 0.309 correlation), MT-GBM maintained strong performance, suggesting effective knowledge transfer.
• Neural networks performed poorly on the same datasets, highlighting the superiority of MT-GBM for tabular data with heterogeneous features and noisy inputs.
• The method showed stable convergence and improved generalization, especially when using adaptive learning rates based on residual error scales.