MultiTab: A Scalable Foundation for Multitask Learning on Tabular Data

AI summary

Overview

• Research area: Multitask learning (MTL) on tabular data; transformer architectures for structured/tabular machine learning.
• Technical level: Intermediate. The paper assumes familiarity with transformers, attention mechanisms, and standard multitask benchmarks, but the core ideas (masking certain attention connections, sharing one model across tasks) are explained in accessible terms.
• Scope: The paper proposes MultiTab-Net, a transformer-based multitask architecture for tabular data, along with MultiTab-Bench, a synthetic multitask dataset generator, and evaluates them against MLP-based multitask models, single-task tabular transformers, and XGBoost across three public datasets.

What This Paper Is About

Tabular data is described as the most abundant data type in the world, and many real datasets now have several related prediction targets that could be learned together. Existing multitask learning methods for tabular data are dominated by MLP-based backbones developed mostly for recommendation systems, and the paper argues these struggle to capture complex feature interactions and to scale when data is abundant. The goal is to bring transformer-style attention to multitask tabular learning, control the interference between tasks, and provide a synthetic benchmark that lets researchers systematically vary task count, task correlation, and task difficulty.

Key Contributions

1. MultiTab-Net — described as the first multitask transformer architecture for tabular data, which the authors report outperforms existing MLP-based multitask models and single-task transformer models on widely used tabular benchmarks.
2. A multitask masked-attention mechanism — a method that selectively removes attention connections (masking F↛T, T↛T, or both) to mitigate task competition during inter-feature relationship learning, described as the first such mechanism applied to multitask transformers.
3. MultiTab-Bench — a synthetic multitask tabular dataset generator enabling fine-grained control over pairwise task correlations and relative task difficulty for an arbitrary number of tasks, extending a prior generator that was limited to exactly two tasks and had no notion of relative task difficulty.
4. A broad empirical evaluation spanning three public datasets (AliExpress, ACS Income, Higgs) with varying feature counts, feature types, task counts, task types, and class balance, plus a computational efficiency analysis and an ablation over masking schemes and task-token designs.

Main Findings

• MultiTab-Net achieves the highest multitask gain on every public dataset. Reported Δm values are 0.5512 on AliExpress, 0.1064 on ACS Income, and 1.2337 on Higgs, compared with STEM (0.1763, 0.0725, 0.0571), PLE (0.2778, 0.0892, −0.0314), MMoE (0.0873, 0.0893, −0.3525), SAINT (0.1146, 0.0948, −1.6514), XGBoost (0.2598, −0.1948, −18.5526), and TabT (0.1956, −0.0336, −24.7917).
• Masking task tokens from each other (T↛T) works best. In the ablation, the multi-token + T↛T configuration gives Δm of 0.5512 (AliExpress), 0.1064 (ACS Income), and 1.2337 (Higgs), beating no masking (0.2579, 0.0783, 1.1182) and F↛T (0.3698, 0.0951, 0.9626). The paper reports that F↛T offers inconsistent benefits, implying features should retain access to task context.
• Multiple task tokens matter more as task count grows. With two tasks (AliExpress, ACS Income) the difference between single and multiple task tokens is described as minimal, but on the eight-task Higgs dataset the multi-token setup yields clear gains over a single shared task token (which produced Δm of −6.3491 with no masking and −12.5587 with F↛T).
• MultiTab-Net scales well in parameters and FLOPs among MTL methods. Reported Params/FLOPs in millions: MultiTab-Net 1.80/4.85 (AliExpress), 0.28/0.77 (ACS Income), 0.70/1.90 (Higgs); MMoE 1.53/3.07, 0.64/1.25, 1.05/2.11; PLE 1.55/3.11, 0.65/1.28, 1.22/2.46; STEM 1.55/3.11, 0.69/1.29, 1.25/2.51; SAINT 3.62/9.70, 0.49/1.35, 5.50/15.02.
• Efficiency versus the single-task architecture scales with task count. Against SAINT, the paper reports an efficiency multiplier roughly equal to the number of tasks — roughly 2x lower cost on AliExpress and ACS Income, and 8x on Higgs. MultiTab-Net is noted to be slightly more expensive than other MTL methods on AliExpress due to its high number of features.
• Robustness on synthetic benchmarks. Using MultiTab-Bench, the paper reports that MultiTab-Net consistently achieves higher multitask gain than MMoE, PLE, and STEM across controlled variations of task correlation (p), task complexity (polynomial degree), and task count, and also outperforms them by a large margin when tasks have markedly different levels of complexity.
• Tree-based and categorical-only models struggle on Higgs. XGBoost's poor performance is said to align with findings in the original Higgs paper where tree-based models struggled to capture the complex targets, and TabT performs poorly because Higgs has only four categorical features, leaving most of that architecture as a shallow MLP.

Methodology in Plain English

The team builds on BERT-style tabular transformers (FT-Transformer, SAINT) that attach a single task token to a sequence of feature tokens. MultiTab-Net instead appends a separate task token for every task, so all tasks are processed in parallel while sharing the feature representations. Both categorical and numerical features, plus the task tokens, are embedded to size e and fed through N encoder blocks. Each block contains an inter-feature attention layer (relationships between columns within a sample, where tasks can attend to features and vice versa) and an inter-sample attention layer (relationships across rows in a batch, following SAINT), each followed by feed-forward networks with layer norm and residual connections. After the blocks, each task token goes through its own small MLP to produce that task's prediction.

The distinctive step is multitask masked attention. Because task tokens create extra quadrants in the attention matrix, dominant tasks could distort attention and cause the "seesaw phenomenon." The authors test three masking schemes that add −∞ to selected attention scores before softmax: masking features from attending to task tokens (F↛T), masking task tokens from attending to other task tokens (T↛T), and both combined. They do not mask task tokens from attending to features (T↛F), since tasks need feature information. Masking cannot be applied to inter-sample attention because it pools information across entire samples.

For MultiTab-Bench, they generalize a prior two-task generator by defining a symmetric correlation matrix P with pairwise correlation p off the diagonal, performing its eigendecomposition, and constructing task weight vectors from the eigenvectors and square-root eigenvalues so that pairwise cosine similarity exactly equals P. Inputs are drawn from a standard multivariate normal; labels come from a polynomial transformation of the projections plus task-specific Gaussian noise, where the polynomial degree controls task complexity and the noise variance controls relative task difficulty. Each plot in the paper uses 100 generated datasets of 10k samples each.

All baselines were implemented from scratch with comparable model capacity, inputs embedded for every baseline, a grid-search hyperparameter sweep with WANDB, the Adam optimizer with weight decay, results aggregated over five random seeds, and all experiments on a single RTX A5500 GPU.

Why This Matters

• Research impact: The paper positions transformer-based multitask learning for tabular data as a largely unexplored direction and provides both a model and a tunable synthetic benchmark. It also argues the multitask masked-attention framework has potential beyond tabular data, in domains such as computer vision and NLP where managing task competition while capturing intricate feature interactions is equally critical.
• Healthcare: The paper's own example is patient records containing age, BMI, and lab results used to predict both the presence of diabetes and the risk of high blood pressure.
• E-commerce and recommendation: Predicting not only whether a user will click a product but also whether they will add it to their cart and make a purchase — the AliExpress tasks of click and conversion.
• Socioeconomic and survey analysis: Datasets like ACS Income support simultaneous prediction of income level and marital status from demographic and employment attributes.
• Scientific data: The Higgs dataset from particle collision simulation, extended from a single classification target to seven additional physics-derived regression targets.
• Industry relevance: Multitask learning is described as already proven in large-scale production recommendation systems for metrics such as click-through rate and conversion rate, and the paper emphasizes computational efficiency — one model handling multiple tasks in parallel reduces training and inference time.

Future Directions

• Scaling to larger task counts and more heterogeneous task sets. The paper notes that the advantage of multiple task tokens grows with task count, but its public evaluation covers 2, 2, and 8 tasks; how the masking design behaves at much larger task counts is an open question.
• Extending multitask masked attention beyond tabular data. The authors explicitly state that the framework has broader potential in vision and NLP, but no experiments in those domains are reported here.
• Making the synthetic benchmark more realistic. MultiTab-Bench offers tunable correlations and difficulty for an arbitrary number of tasks, but the paper notes that existing tabular foundation models such as TabPFN and TabICL handle only up to 10k and 60k samples respectively, leaving large-scale in-context or pretrained multitask tabular modeling as an open direction.
• Controlling the cost on high-feature-count data. MultiTab-Net is reported as slightly more expensive than other MTL methods on AliExpress because of its high number of features, suggesting further efficiency work on wide tables.

Target Audience

Researchers and practitioners working on tabular machine learning, multitask learning, or recommendation systems, especially those already comfortable with transformer architectures and looking for a transformer-native alternative to MLP-based multitask backbones. It is also relevant for applied scientists in healthcare, finance, e-commerce, and the physical sciences who have datasets with several related prediction targets, and for benchmark designers interested in a synthetic generator with tunable task count, correlation, and difficulty. Readers without background in attention mechanisms will need to consult the referenced transformer literature, since the masking formulation is presented in equations.

Authors’ abstract

Tabular data is the most abundant data type in the world, powering systems in finance, healthcare, e-commerce, and beyond. As tabular datasets grow and span multiple related targets, there is an increasing need to exploit shared task information for improved multitask generalization. Multitask learning (MTL) has emerged as a powerful way to improve generalization and efficiency, yet most existing work focuses narrowly on large-scale recommendation systems, leaving its potential in broader tabular domains largely underexplored. Also, existing MTL approaches for tabular data predominantly rely on multi-layer perceptron-based backbones, which struggle to capture complex feature interactions and often fail to scale when data is abundant, a limitation that transformer architectures have overcome in other domains. Motivated by this, we introduce MultiTab-Net, the first multitask transformer architecture specifically designed for large tabular data. MultiTab-Net employs a novel multitask masked-attention mechanism that dynamically models feature-feature dependencies while mitigating task competition. Through extensive experiments, we show that MultiTab-Net consistently achieves higher multitask gain than existing MTL architectures and single-task transformers across diverse domains including large-scale recommendation data, census-like socioeconomic data, and physics datasets, spanning a wide range of task counts, task types, and feature modalities. In addition, we contribute MultiTab-Bench, a generalized multitask synthetic dataset generator that enables systematic evaluation of multitask dynamics by tuning task count, task correlations, and relative task complexity. Our code is publicly available at https://github.com/Armanfard-Lab/MultiTab.