TabH2O: A Unified Foundation Model for Tabular Prediction

2026-05-18

AI Summary

This work proposes TabH2O, a unified in-context learning framework for tabular data that simultaneously handles both classification and regression tasks within a single forward pass, addressing limitations of conventional approaches that require separate modeling, incur high pretraining costs, and exhibit sensitivity to noise. The method introduces a unified dual-head architecture, a single-stage pretraining strategy, and an explicit noise dimension to enhance robustness, complemented by bounded scalable softmax, inter-stage normalization, learnable residual scaling, and soft logit clipping. Evaluated on the TALENT benchmark comprising 300 datasets, TabH2O achieves an average rank of 2.55 and places among the top three methods on 81% of test sets, significantly outperforming established baselines such as CatBoost and LightGBM.

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.

Problem

• tabular prediction
• foundation model
• classification and regression
• in-context learning
• unified model

Innovation

• unified foundation model
• in-context learning
• single-stage pretraining
• noise-aware pretraining
• tabular prediction