TAFFY: A Task-Adaptive Tabular Foundation Model with In-Context Diversity

Zijian Li; Xiangchen Song; Gongxu Luo; Jie Qiao; Ruichu Cai; Zhenhao Chen; Xinshuai Dong; Fan Feng

cs.LGOct 6, 2026

TAFFY: A Task-Adaptive Tabular Foundation Model with In-Context Diversity

Authors: Zijian Li, Xiangchen Song, Gongxu Luo, Jie Qiao, Ruichu Cai, Zhenhao Chen, Xinshuai Dong, Fan Feng, +2 more

Organizations: Carnegie Mellon University · Mohamed bin Zayed University of Artificial Intelligence · Guangdong University of Technology · University of California, San Diego

Abstract

Recent progress in tabular foundation models suggests that training on synthetic tasks can substantially improve in-context learning capabilities, with overall performance largely depending on how well models can infer task-specific predictive relationships from the available context during inference. In this paper, we introduce TAFFY, a tabular foundation model with an In-Context Diversity Prior and a Task-Conditioned Looped Transformer that strengthen this ability. Specifically, to construct each synthetic pretraining context, the In-Context Diversity Prior samples from multiple related environments derived via controlled interventions and distribution shifts on a shared causal process. This in-context diversity encourages the model to learn a more comprehensive and task-specific representation. Moreover, the Task-Conditioned Looped Transformer iteratively and selectively applies a shared group of Transformer blocks to refine contextual representations, with a task-conditioned gate modulating the final hidden-state update. This enables task-adaptive iterative refinement. Together, these components encourage the model to identify predictive relationships from contextual contrasts during pretraining and dynamically modulate context integration for each task. Across six classification and five regression benchmark datasets, TAFFY attains the lowest average rank.

Figures & tables

Figure 1: Predictive performance and pretraining-data efficiency. (a–c) Average accuracy ranks on OpenML-CC18, BCCO, and TALENT using the ranking pools of Table 1 (lower is better; rank axes start at 2). (d) Accuracies on TALENT classification datasets versus cumulative maximum feature elements. Markers denote measured checkpoints connected by lines; the grey line marks TabICLv2’s best measured accuracy.

Figure 2: The pretraining pipeline of Taffy combines the in-context diversity prior with the task-conditioned looped transformer. (a) In-Context Diversity Prior proceeds through three steps: base-table generation, multiple-environment construction, and context assembly. (b) Column and row encoders produce representations for a looped Transformer, whose final update is modulated by a support-derived task gate before query prediction.

Model	Classification						Regression
Model	BCCO	OpenML	PFN	TALENT	TabArena	TabZilla	BCCO	CTR23	PFN	TALENT	TabArena
Ensemble methods
AutoGluon	10.62	8.29	9.88	9.96	9.98	8.52	5.96	5.85	5.96	5.89	6.15
Tree-based methods
CatBoost	10.29	9.88	9.30	10.70	10.78	11.02	7.20	7.39	7.25	7.02	7.77
XGBoost	11.66	10.52	9.95	11.81	11.03	11.07	11.30	10.55	10.93	10.39	10.92

Table 1: Average ranks across all evaluated classification and regression datasets (lower is better). The classification row uses Taffy-4L on every suite; regression uses the separately trained model.

Figure 3: Component analyses: (a) in-context versus cross-context diversity; (b) task-specific versus global gating; (c) matched training/inference depth; (d) extra inference loops after two-loop training. In (a,b), each variant is ranked by accuracy against four fixed baselines (TabICLv1, TabICLv2, LimiX-2M, and LimiX-16M) on each dataset, with average ranks for ties before averaging across datasets. In (c), ranks are computed on all datasets of each benchmark. Lower ranks are better; thin lines show checkpoints, thick lines show smoothed trends.

Figure 4: Average-rank ablation on TALENT classification datasets. Each variant is ranked separately against TabICLv1, TabICLv2, LimiX-2M, and LimiX-16M.

Property	2L	3L	4L
Categorical fraction	$-.533^{*}$	$-.453^{*}$	$-.452^{*}$
Excess kurtosis	$+.230^{*}$	$+.326^{*}$	$+.323^{*}$
Label entropy	$+.084$	$+.010$	$-.011$
Missing ratio	$+.021$	$+.042$	$+.031$

Table 2: Spearman correlations between

\alpha

and task attributes. ∗ indicates corrected

q<0.05

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

Distribution shifts (row sampling from $\mathcal{M}_{0}$ )
Operator	Implementation	Preserved	Changed	Strength
Label shift	Reweight rows by $a_{e}(y)=\exp\!\big(\tau_{e}u_{e}(y)\big)$	$p(x\mid y)$	$p(y)$ , hence $p(y\mid x)$	KL $=0.1$
Covariate shift	Reweight rows by $b_{e}(x)=\exp\!\big(-\tau_{e}\\|x-\mu_{e}\\|_{2}\big)$	$p(y\mid x)$	$p(x)$	KL $=0.1$
Conditional shift	Reweight rows within each class by $c_{e}(x,y)=\exp\!\big(-\tau_{e}\\|x-\mu_{e,y}\\|_{2}\big)$	$p(y)$	$p(x\mid y)$ , hence $p(y\mid x)$	KL $=0.1$
Interventions (mechanism changes in $\mathcal{M}_{0}$ )
Hard intervention	Set $z_{j}:=v_{e,j}$ with $v_{e,j}\sim\mathcal{U}[\ell_{j},u_{j}]$ for $j\in I_{e}$	Other mechanisms	$z_{j}$ and its descendants	$5\%$ of non-label nodes

Appendix

Table 3: Environment operators of the In-Context Diversity Prior. The three distribution shifts reweight rows sampled from the unchanged base model, with strength fixed by

\mathrm{KL}(p_{e}\,\|\,p_{0})=0.1

. The two interventions modify the mechanisms of

5\%

of the non-label nodes and regenerate the data.

Category	Statistics	Dim.
Size and validity	$\log(1+d)$ , $\log(1+n_{s})$ , $\log(1+N)$ , $n_{s}/N$ , $d/H$ , $1-d/H$ , and the fraction of finite support values over valid features.	7
Support labels	Observed class count divided by $K_{\max}$ ; label entropy divided by $\log K_{\max}$ ; largest and smallest proportions among observed classes, with $K_{\max}=20$ .	4
Support features	Per-feature mean, standard deviation, mean absolute value, minimum, maximum, and range.	$6\times 4$
	Per-feature fractions of zero, near-zero, near-integer, and near-binary values.	$4\times 4$
Total		51

Appendix

Table 4: Statistics used by the task-specific support gate. The descriptor contains

51

entries. Here,

H

is the input feature width before grouping, including padding, and

d\leq H

is the number of valid features. Feature summaries are aggregated across valid features by mean, standard deviation, minimum, and maximum.

Model	Time/step (s)	Relative	25k steps (h)	GPU-hours
Single-pass backbone	7.01	$1.00\times$	48.66 (2.03 days)	3,114
Taffy-2L	8.79	$1.25\times$	61.03 (2.54 days)	3,906
Taffy-3L	10.62	$1.51\times$	73.74 (3.07 days)	4,719
Taffy-4L	12.62	$1.80\times$	87.67 (3.65 days)	5,611

Appendix

Table 5: Pretraining cost on 64 AMD MI210 GPUs. Time per step is measured; total time and GPU-hours are projected for the 25,000-step schedule. Relative is the time per step divided by that of the single-pass backbone.

Method	Avg. rank $\downarrow$	Elo $\uparrow$
RF	8.42	1000.00
TabICLv1	5.17	1595.23
LimiX-2M	5.50	1552.44
TabPFN2	5.17	1595.23
Mitra	8.25	1044.05
LimiX-16M	5.12	1600.49

Appendix

Table 6: Results on TALENT classification datasets with more than ten classes. Only methods with observed results on all such datasets are included.

Ensemble methods
Model	Classification						Regression					Average
Model	BCCO	OpenML	PFN	TALENT	TabArena	TabZilla	BCCO	CTR23	PFN	TALENT	TabArena	Average
AutoGluon	10.53	8.29	9.88	10.00	9.98	8.52	5.86	5.85	5.96	6.00	6.15	7.91
Tree-based methods
CatBoost	10.47	9.88	9.30	10.72	10.78	11.02	7.12	7.39	7.25	6.99	7.77	8.97
XGBoost	11.68	10.52	9.96	11.71	11.03	11.09	11.28	10.55	10.93	10.29	10.92	10.91

Appendix

Table 7: Average ranks on the LimiX-2M-eligible, sample-compatible evaluation pool (lower is better). Rank aggregation uses the complete applicable-method subset within each suite; dashes denote incomplete task-type coverage.

	Classification						Regression
Method	TALENT	BCCO	OpenML	PFN	TabArena	TabZilla	TALENT	BCCO	CTR23	PFN	TabArena	Mean
AutoGluon	1139.13	1001.45	1217.39	1064.14	1022.49	1187.80	1241.40	1295.80	1310.03	1257.72	1332.51	1188.17
CatBoost	1105.99	1015.31	1150.20	1088.56	986.67	1086.68	1171.15	1210.85	1219.76	1178.34	1204.78	1128.94
XGBoost	1056.39	957.92	1123.70	1061.09	975.45	1084.46	981.73	956.97	1050.70	964.64	970.83	1016.72
ET	986.81	934.06	1080.91	1017.65	800.16	1053.96	1016.33	1027.80	1078.09	1002.07	1074.19	1006.55
RF	1000.00	1000.00	1000.00	1000.00	1000.00	1000.00	1000.00	1000.00	1000.00	1000.00	1000.00	1000.00

Appendix

Table 8: Corresponding observed-only Elo ratings under the operational development-set reconstruction (higher is better).

Ensemble methods
Model	Classification						Regression					Average
Model	BCCO	OpenML	PFN	TALENT	TabArena	TabZilla	BCCO	CTR23	PFN	TALENT	TabArena	Average
AutoGluon	1011.38	1217.39	1064.14	1138.25	1021.54	1187.80	1303.99	1310.03	1257.72	1236.25	1332.51	1189.18
Tree-based methods
CatBoost	1013.58	1150.20	1088.56	1106.17	973.79	1086.68	1216.33	1219.76	1178.34	1174.86	1204.78	1128.46
XGBoost	962.97	1123.70	1061.09	1061.68	976.48	1084.46	957.85	1050.70	964.64	986.98	970.83	1018.31

Appendix

Table 9: Elo ratings on the LimiX-2M-eligible, sample-compatible evaluation pool (higher is better; RF = 1000 per benchmark). Only observed pairwise outcomes are used; missing scores are not imputed, and dashes denote incomplete task-type coverage.

Descriptor	Model	Partial $\rho$ [95% CI]	$q$
Categorical fraction	Taffy-2L	$-0.533\;[-0.652,-0.394]$	$0.0003$
	Taffy-3L	$-0.453\;[-0.580,-0.299]$	$0.0003$
	Taffy-4L	$-0.452\;[-0.579,-0.296]$	$0.0003$
Numerical excess kurtosis	Taffy-2L	$0.230\;[0.057,0.383]$	$0.0338$
	Taffy-3L	$0.326\;[0.167,0.466]$	$0.000703$
	Taffy-4L	$0.323\;[0.167,0.463]$	$0.000703$

Appendix

Table 10: Adjusted associations between

\alpha

and TALENT dataset properties. Categorical fraction and numerical excess kurtosis use separate multiple-testing corrections.

	Support rows		Effective columns
Model	$\rho$	$q$	$\rho$	$q$
Taffy-2L	$0.192$	$0.0374$	$0.138$	$0.199$
Taffy-3L	$0.297$	$2.31\!\times\!10^{-4}$	$0.053$	$0.660$
Taffy-4L	$0.397$	$3.44\!\times\!10^{-7}$	$0.022$	$0.881$

Appendix

Table 11: Unadjusted associations between

\alpha

and support size on TALENT datasets. All

q

-values use Benjamini–Hochberg correction over the original 216 tests.

Inference condition	Update coefficient	Average rank $\downarrow$
Task-specific gate	$\alpha_{S}$	1.57
Support-independent gate	$\tanh(a)$	1.88
Gate bypass (full update)	$1$	2.55

Appendix

Table 12: Frozen-checkpoint gate intervention on matched TALENT classification datasets. Mean rank among the three inference conditions; lower is better.

Explore similar work

Is One Layer Enough? Understanding Inference Dynamics in Tabular Foundation Models
May 7, 2026Amir Rezaei Balef, Mykhailo Koshil, Katharina EggenspergerTabular Foundation Models Transformer Interpretability
TU Dortmund University, Dortmund, Germany · Lamarr Institute for Machine Learning and Artificial Intelligence, Dortmund, Germany · University of Tübingen, Tübingen, Germany
TabICLv2: A better, faster, scalable, and open tabular foundation model
Feb 11, 2026Jingang Qu, David Holzmüller, Gaël Varoquaux +1Synthetic Data Pretraining Tabular Foundation Models
SODA Team, INRIA Saclay, Palaiseau, France · Probabl, France
TabPFN-MT: A Natively Multitask In-Context Learner for Tabular Data
May 16, 2026Cormac Cureton, Narges ArmanfardTabular Foundation Models Efficient Neural Network Inference
McGill University · Mila - Quebec AI Institute · Montreal, QC, Canada

May 7, 2026, cs.LGJ/K move · Enter open · S save

Is One Layer Enough? Understanding Inference Dynamics in Tabular Foundation Models

Amir Rezaei Balef, Mykhailo Koshil, Katharina Eggensperger

TU Dortmund University, Dortmund, Germany · Lamarr Institute for Machine Learning and Artificial Intelligence, Dortmund, Germany · University of Tübingen, Tübingen, Germany

Tabular Foundation Models Transformer Interpretability Tabular ICL

Open paper AskPDF arXiv

Transformer-based tabular foundation models (TFMs) dominate small to medium tabular predictive benchmark tasks, yet their inference mechanisms remain largely unexplored. We present the first large-scale mechanistic study of layerwise dynamics in 6 state-of-the-art tabular in-context learning models. We explore how predictions emerge across depth, identify distinct stages of inference and reveal latent-space dynamics that differ from those of language models. Our findings indicate substantial depthwise redundancy across multiple models, suggesting iterative refinement with overlapping computations during inference stages. Guided by these insights, we design a proof-of-concept, looped single-layer model that uses only 20% of the original model's parameters while achieving comparable performance. The code is available at https://github.com/amirbalef/is_one_layer_enough.

Abstract

Explore similar work

Is One Layer Enough? Understanding Inference Dynamics in Tabular Foundation Models

TabICLv2: A better, faster, scalable, and open tabular foundation model

TabPFN-MT: A Natively Multitask In-Context Learner for Tabular Data