[Paper Review] Transfer Learning with Deep Tabular Models

R. V. Levin, Valeriia Cherepanova+6 more|arXiv (Cornell University)|Jun 30, 2022

Artificial Intelligence in Healthcare and Education18 citations

TL;DR

This paper demonstrates that transfer learning with deep tabular models significantly outperforms gradient boosted decision trees (GBDTs) in low-data medical diagnosis scenarios. By leveraging supervised pre-training and a novel pseudo-feature method for feature set mismatches, the approach achieves state-of-the-art performance, especially when downstream data is scarce.

ABSTRACT

Recent work on deep learning for tabular data demonstrates the strong performance of deep tabular models, often bridging the gap between gradient boosted decision trees and neural networks. Accuracy aside, a major advantage of neural models is that they learn reusable features and are easily fine-tuned in new domains. This property is often exploited in computer vision and natural language applications, where transfer learning is indispensable when task-specific training data is scarce. In this work, we demonstrate that upstream data gives tabular neural networks a decisive advantage over widely used GBDT models. We propose a realistic medical diagnosis benchmark for tabular transfer learning, and we present a how-to guide for using upstream data to boost performance with a variety of tabular neural network architectures. Finally, we propose a pseudo-feature method for cases where the upstream and downstream feature sets differ, a tabular-specific problem widespread in real-world applications. Our code is available at https://github.com/LevinRoman/tabular-transfer-learning .

Motivation & Objective

• Address the challenge of limited labeled data in tabular machine learning, particularly in medical applications where data collection is costly.
• Investigate whether deep tabular models can effectively transfer knowledge across tasks, unlike traditional GBDT approaches.
• Develop practical solutions for real-world issues such as feature set mismatches between upstream and downstream datasets.
• Compare supervised and self-supervised pre-training strategies for tabular transfer learning to determine which yields more transferable representations.

Proposed method

• Pre-train deep tabular models (e.g., FT-Transformer, TabMLP) on large upstream tabular datasets to learn generalizable representations.
• Fine-tune the pre-trained models on downstream tasks with limited labeled data, enabling knowledge transfer.
• Propose a pseudo-feature method: pre-train on upstream data missing downstream features, then use the model to predict missing features on upstream data before re-pre-training and transferring.
• Implement both supervised and self-supervised pre-training strategies to compare their transferability in tabular settings.
• Use the MetaMIMIC repository to construct realistic medical diagnosis benchmark tasks with varying data availability.
• Apply end-to-end training with learnable heads (linear or MLP) to optimize downstream performance after transfer.

Experimental results

Research questions

• RQ1 Can deep tabular models with transfer learning outperform strong GBDT baselines when downstream data is limited?
• RQ2 Does supervised pre-training lead to more transferable representations than self-supervised pre-training in the tabular domain?
• RQ3 How effective is the proposed pseudo-feature method in handling feature set mismatches between upstream and downstream datasets?
• RQ4 What is the impact of pre-training strategy (supervised vs. self-supervised) on downstream performance in low-data settings?
• RQ5 Can transfer learning with deep tabular models generalize effectively across diverse medical diagnosis tasks with minimal labeled examples?

Key findings

• Deep tabular models with transfer learning consistently outperform GBDT baselines, even those using stacking to leverage upstream data, especially with limited downstream data.
• Supervised pre-training yields more transferable features than self-supervised pre-training in tabular data, contrary to trends in computer vision.
• The proposed pseudo-feature method achieves performance comparable to models pre-trained with ground-truth values of missing features, significantly outperforming methods that discard missing features.
• On the MetaMIMIC medical diagnosis benchmark, transfer learning with FT-Transformer achieved ROC-AUC scores of 0.878 for 'Alcohol' and 0.836 for 'Purpura' with limited downstream data.
• Even with only 1% of downstream data, transfer learning with deep tabular models achieved 0.692 AUC on the 'Alcohol' task, outperforming XGBoost+stacking (0.692) and Catboost+stacking (0.692), demonstrating robustness in low-data regimes.
• The pseudo-feature method improved performance by up to 0.05 AUC in some cases compared to discarding missing features, validating its practical utility in real-world settings.