TabH2O: A Unified Foundation Model for Tabular Prediction
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.1 (29.2M parameters) on the TALENT benchmark (300 datasets), where it achieves an average rank of 2.54 out of 6 evaluated methods, outperforming tuned CatBoost (4.00), H2O AutoML (4.28), LightGBM (4.98), and TabPFN v2.6 (2.80), and competitive compared to TabICL v2 (2.13), while placing in the top-3 on 80% of the testing datasets across classification and regression tasks. We also achieve state-of-the-art results on TabArena. Moreover, our single model is well-suited to tabular datasets up to roughly 500,000 rows and 100 features.