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

2026-10-05

AI Summary

This study addresses the challenge that tabular foundation models face in accurately inferring task-specific predictive relationships from context during inference. To this end, we propose TAFFY, a framework that introduces a context diversity prior and leverages causal intervention to generate multi-environment synthetic data for pre-training, thereby constructing richer contextual representations. Furthermore, we design a task-conditioned gated recurrent Transformer architecture that enables dynamic, iterative refinement of context, enhancing the model's task adaptability. Experimental results demonstrate that TAFFY achieves the best average ranking across multiple classification and regression benchmarks, significantly outperforming existing methods.

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.

Problem

• tabular foundation model
• in-context learning
• synthetic pretraining
• task-adaptive

Innovation

• Tabular Foundation Model
• In-Context Diversity Prior
• Task-Conditioned Looped Transformer
• In-Context Learning
• Synthetic Pretraining