MultiTab-Net Transformer for Tabular Data

• The paper presents a novel multitask transformer architecture featuring a multitask masked-attention mechanism that dynamically models complex dependencies in tabular data.
• MultiTab-Net employs a dual-attention paradigm by integrating inter-feature and inter-sample attention, delivering superior multitask performance over traditional MLP approaches.
• Empirical results demonstrate consistent multitask gains across diverse datasets, highlighting its efficacy in managing task interference and scalability.

MultiTab-Net is a multitask transformer architecture tailored specifically for learning on large-scale tabular data. Designed to address deficiencies in previous multitask learning (MTL) approaches—particularly those using multi-layer perceptron (MLP) backbones—MultiTab-Net integrates a novel multitask masked-attention mechanism to dynamically model complex feature-feature dependencies and systematically prevent adverse task interactions. Its modular construction, scalable regularization, and empirical superiority across diverse domains mark it as a foundation model for multitask tabular prediction (Sinodinos et al., 13 Nov 2025).

1. Architectural Design and Input Processing

MultiTab-Net employs a transformer-based backbone in contrast to conventional MLP architectures. Each of the $d$ input features, whether categorical or numerical, is individually embedded into a vector of dimension $e$. Categorical features use embedding lookups, while numerical features are passed through a linear projection followed by LayerNorm. For multitask settings with $t$ tasks, each task is allocated a distinct, learnable "task token" $T_i \in \mathbb{R}^e$. The concatenation of $d$ feature tokens and $t$ task tokens yields an input matrix $x \in \mathbb{R}^{(d+t)\times e}$.

The architecture comprises $N$ identical transformer encoder blocks. Each block features two distinct self-attention modules:

• Inter-Feature Attention: Attends over all $(d + t)$ tokens within a given sample, enabling explicit modeling of within-sample feature and cross-task interactions.
• Inter-Sample Attention: Operates across samples in the batch, treating each flattened $(d + t) \times e$ vector as a token, facilitating learning of sample-sample relationships.

Each encoder block adheres to the standard transformer order: layer-normalized multi-head self-attention, residual connection, layer normalization, feed-forward network, and another residual/normalization step.

2. Multitask Masked-Attention Mechanism

The introduction of multiple task tokens prompts division of the attention matrix $A \in \mathbb{R}^{(d+t)\times(d+t)}$ into blocks representing Feature$\to$Feature, Feature$\to$Task, Task$\to$Feature, and Task$\to$Task interactions. Unconstrained Task$\to$Task attention introduces "task competition," leading to instability or the "seesaw phenomenon."

For the $i$-th attention head, with $d_k = e/h$,

• $Q_i = x W_{Q_i}$
• $K_i = x W_{K_i}$
• $V_i = x W_{V_i}$
• $$A_i = \mathrm{softmax}\left( \frac{Q_i K_i^\top}{\sqrt{d_k}} + M_A \right)$$
• $x_{att, i} = A_i V_i$
• $$x_{out} = \mathrm{Concat}(x_{att, 1}, ..., x_{att, h}) W_O$$

Here, the mask matrix $M_A \in \{0, -\infty\}^{(d+t)\times(d+t)}$ selectively blocks attention flows: $M_A[j, k] = -\infty$ disables, $0$ enables. Several masking schemes are considered:

• $F \not\to T$: Mask Feature$\to$Task block.
• $T \not\to T$: Mask Task$\to$Task block.
• Both: Mask Feature$\to$Task and Task$\to$Task.

Empirical investigations find that employing multiple task tokens (one per task) combined with $T \not\to T$ masking yields the best multitask gain, effectively reducing destructive interference and enabling stable multitask sharing.

3. Modeling of Feature–Feature and Sample–Sample Dependencies

MultiTab-Net’s attention mechanisms enable explicit and dynamic modeling of dependencies across both feature and sample axes:

• Inter-Feature Attention: Every feature token attends to all others (and all task tokens); explicit identity encoding replaces positional embeddings.
• Inter-Sample Attention: The batch of $n$ samples, each with flattened feature-task representations, undergoes standard multi-head self-attention. A dedicated set of projection matrices $(W_Q, W_K, W_V, W_O)$ is used separate from in-sample modules. This allows for modeling correlations and dependencies across entire rows.

This dual-attention paradigm is unattainable with classic MLPs, which encode feature and row dependencies only implicitly.

4. Mitigating Task Competition and Loss Formulation

To systematically counteract negative transfer and the seesaw effect, MultiTab-Net incorporates several architectural and optimization strategies:

• Multi-token Design: Each task receives a unique learnable token.
• Task$\to$Task Masking: Prevents direct attention-based interference among task tokens.
• Task-Specific Output Heads: After the final encoder layer, each outputted task token $T'_i \in \mathbb{R}^e$ is processed by a task-specific MLP to yield $\hat y_i$ for the $i$-th task.
• Loss Aggregation: The total training loss is $$L_{\mathrm{total}} = \sum_{i=1}^t L_i(\hat y_i, y_i)$$, with $L_i$ being cross-entropy for classification or mean square error for regression. All tasks are equally weighted ($w_i=1$) during training.

5. Scalability, Regularization, and Hyperparameterization

The architecture is designed for scalability and robust training through:

• LayerNorm post each sub-layer to stabilize deep updates.
• Separate Dropout in attention and feed-forward modules, with rate selection in $\{0.0, 0.1, 0.2, 0.3\}$.
• Optional RoPE: Rotary positional embeddings in inter-sample attention to counter representation collapse with large batches.
• Capacity Control: Embedding sizes and hidden dimensions are tuned to match strong baselines (MLP, MMoE).

Hyperparameters and training settings are dataset-specific. For example:

• Embedding dimension $e=8$ (AliExpress, 75 features), $e=16$ (ACS Income, Higgs).
• Transformer hidden size of $256$, $h=4$ attention heads, $d_k=64$ or $32$.
• Encoder depth $N=3$--$6$, optimized per dataset.
• Adam optimizer (weight decay $1\mathrm{e}{-5}$), batch size $2048$, learning rate grid $$\{1\mathrm{e}{-4}, 1\mathrm{e}{-3}, 1\mathrm{e}{-2}, 1\mathrm{e}{-1}\}$$, and early stopping (patience 3--5 epochs).

6. Empirical Evaluations and Benchmarks

MultiTab-Net’s performance is evaluated on public benchmarks and synthetic multitask tasks (MultiTab-Bench). Multitask gain $\Delta_m$ is the average percent improvement over the best single-task MLP.

Results on public datasets:

Dataset (tasks)	MultiTab-Net $\Delta_m$	Best MLP-MTL Baseline	Single-Task Transformer
AliExpress (2, binary)	0.5512	PLE (0.2778)	SAINT ($\approx$0.11)
ACS Income (binary, multi-cls)	0.1064	PLE (0.0892)	—
Higgs (1 bin. + 7 reg.)	1.2337	SAINT (0.0948)	—

On synthetic MultiTab-Bench tasks, MultiTab-Net achieves the highest $\Delta_m$ under various task counts ($t \in \{3,5,7\}$), task correlations ($p \in \{0.2, 0.6, 1.0\}$), and complexity (polynomial degree), demonstrating robust generalization and resistance to negative transfer.

7. Factors Underpinning MultiTab-Net Performance

Several inductive and architectural strategies explain the observed improvements:

• Explicit attention across feature-feature and sample-sample relations exposes patterns missed by MLP or Mixture-of-Experts.
• The multi-token approach enables each task to maintain its unique contextual representation while leveraging shared signals through the feature-attention mechanism.
• Task$\to$Task attention masking promotes minimal crosstalk and suppresses the seesaw effect.
• Inter-sample attention enables detection and exploitation of batch-level tabular patterns.
• Control over model capacity and careful regularization ensure observed task improvements derive from data and architectural innovation, not merely increased parameter count.

In sum, MultiTab-Net integrates transformer-based dynamic interaction modeling, multi-token task representation, and targeted task-interference mitigation, yielding substantial and consistent multitask gains on real and synthetic, large-scale tabular workloads (Sinodinos et al., 13 Nov 2025).