TabH2O: A Unified Foundation Model for Tabular Prediction

TL;DR

TabH2O is a unified foundation model for tabular data that efficiently handles classification and regression tasks through in-context learning, with novel training techniques for robustness and stability.

Contribution

It introduces a single, dual-head model trained with a simplified, noise-aware pretraining process for improved tabular prediction performance.

Findings

Achieves top-3 performance on 81% of datasets in TALENT benchmark.

Outperforms traditional models like CatBoost, H2O AutoML, and LightGBM.

Competitive with state-of-the-art TabPFN v2, placing second in average rank.

Abstract

We present TabH2O, a foundation model for tabular data that performs classification and regression in a single forward pass via in-context learning. TabH2O builds on the TabICL architecture with several key modifications: (1) unified training, a single model handles both classification and regression via a dual-head architecture, eliminating the need for separate models and reducing total pretraining cost; (2) single-stage pretraining, training stability improvements (bounded scalable softmax, inter-stage normalization, learnable residual scaling, logit soft-capping) eliminate the need for multi-stage curriculum learning, enabling training with full-length sequences from the start; and (3) noise-aware pretraining, synthetic datasets include explicit noise dimensions to teach the model robustness to irrelevant features. We evaluate TabH2O v1 (29.2M parameters) on the TALENT benchmark (300 datasets), where it achieves an average rank of 2.55 out of 6 evaluated methods, outperforming tuned CatBoost (4.07), H2O AutoML (4.18), and LightGBM (5.08), competitive with TabPFN v2.6 (2.74), and behind TabICL v2 (2.12), while placing in the top-3 on 81% of the testing datasets across classification and regression tasks.