PaperlayerView exact revision on arXiv

Source-linked AI summary

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, Guangyi Chen, Kun Zhang

arXiv:2610.07559v1cs.LG

TL;DR

Tabular foundation models must infer task-specific predictive relationships from context as feature meanings and dependencies shift across datasets. TAFFY combines related synthetic environments with a task-conditioned looped Transformer; it attains the lowest average rank across six classification and five regression suites.

  • Problem

    Because feature meanings and dependencies shift across datasets, tabular models must infer task-specific predictive relationships from the support context.

    Source: §1 · p. 1

  • Method

    TAFFY constructs synthetic tasks from related environments and repeatedly refines contextual representations with a shared Transformer whose final update is gated using task information.

    Source: §1 · p. 1, §4.4 · p. 6

  • Results

    TAFFY attains the lowest average rank across six classification and five regression benchmark suites.

    Source: §1 · p. 2, §6 · p. 9

  • Takeaways & Limitations

    Across the stated benchmark suites, the reported classification and regression configurations rank first on all eleven benchmarks under the stated protocols.

    Source: §6 · p. 9

    The model receives no environment identifiers and is trained on pooled prediction, not environment prediction or causal identification.

    Source: §4.2 · p. 5

Abstract

from arXiv · show

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.

1 INTRODUCTION

TAFFY combines multi-environment synthetic task construction with task-conditioned iterative context processing to learn task-specific predictive relationships. Across six classification and five regression benchmark suites, it achieves the lowest average rank on each suite.

1 INTRODUCTION

  • The In-Context Diversity Prior combines support and query observations from related environments created through distribution shifts and interventions on a shared causal prototype.The model is not given environment identifiers; training on the resulting contrasts encourages it to infer task-specific predictive relationships.Source: §1 · pp. 1–2
  • The Task-Conditioned Looped Transformer repeatedly applies a shared Transformer stack and uses a support-derived gate to modulate the final hidden-state update.The gate makes context refinement task-dependent.Source: §1 · p. 2
  • TAFFY obtains the lowest average rank on all six classification and five regression benchmark suites.The evaluation also compares selected models on OpenML-CC18, BCCO, and TALENT.Source: §1 · p. 2

2 RELATED WORK

Related work places TAFFY among synthetic-task tabular foundation models, recurrent Transformer approaches, and research on causal learning. TAFFY uses environmental variation to train pooled prediction rather than to recover the generating process.

2.1 TABULAR FOUNDATION MODELS

  • Tabular foundation models use synthetic tasks for in-context prediction, while related approaches vary their generators, representations, or training data.TAFFY differs from Drift-Resilient TabPFN by mixing related environments without supplying environment identifiers.Source: §2.1 · pp. 2–3

2.2 LOOPED TRANSFORMERS

  • TAFFY shares multi-block Transformer parameters across recurrent passes and gates only the final update, without token routing or changing the number of passes.Prior work explores shared processing depth, iterative algorithms, latent reasoning, and learned computation allocation.Source: §2.2 · p. 3

2.3 CAUSAL LEARNING

  • Causal-learning research motivates structured variation across environments, but the cited identifiability results do not establish identifiability for TAFFY’s mixtures without environment identifiers.TAFFY uses causal prototypes and interventions to construct prediction tasks, not to recover the generating process.Source: §2.3 · p. 3

3 PROBLEM FORMULATION

The formulation treats tabular prediction as conditioning query predictions on labeled support data. Pretraining draws synthetic support–query tasks from a prior and minimizes query prediction loss.

3 PROBLEM FORMULATION

  • A tabular task provides labeled support examples and query features, and the model predicts query labels conditioned on the support dataset.Outputs are class probabilities for classification or conditional quantile predictions for regression.Source: §3 · p. 3
  • The pretraining prior specifies the joint construction of support and query data for synthetic prediction tasks.The query prediction loss is minimized over tasks drawn from that prior.Source: §3 · pp. 3–4
  • The pipeline encodes support and query representations, repeatedly refines them, and uses the final representations to predict query labels.The query-label objective supervises pretraining, while query labels are not model inputs.Source: Figure 2 · p. 4, §3 · p. 4

4 TAFFY

TAFFY coordinates synthetic task construction with a prediction network that repeatedly integrates support-context information. Its prior assembles related environments into tasks, while a support-statistics gate controls the final recurrent refinement.

4.1 OVERALL FRAMEWORK

  • The overall framework trains prediction on sampled tasks whose context can combine observations from multiple causally related environments.At inference, the prediction network and task-conditioning branch are used without the synthetic generator or synthetic environment identifiers.Source: §4.1 · p. 4

4.2 IN-CONTEXT DIVERSITY PRIOR

  • The In-Context Diversity Prior creates related environments from a shared causal prototype and pools their support and query data into one task.Environments use distribution shifts or interventions; up to four may be assembled, while 50% of pretraining tasks remain single-environment.Source: §4.2 · pp. 4–5

4.3 COLUMN AND ROW ENCODERS

  • Column and row encoders form support-label-informed row representations for the looped learner, while query labels are never used as inputs.The encoders run once per forward pass; only the ICL learner is repeated.Source: §4.3 · pp. 5–6

4.4 TASK-CONDITIONED LOOPED TRANSFORMER

  • A gate derived from support-dataset statistics scales the final-pass refinement before it is added to the preceding hidden state.All passes execute; task conditioning changes the final representation rather than the number of computation steps.Source: §4.4 · p. 6
  • The looped Transformer repeatedly applies the same twelve-block ICL stack to support and query representations, with query examples not attending to one another.Each pass updates both support and query representations using support-only keys and values.Source: §4.4 · p. 6

5 EXPERIMENTS

Across six classification and five regression benchmarks, TAFFY attains the lowest average rank; the experiments compare diversity placement, task-conditioned gating, loop depth, and pretraining-data efficiency.

5.1 EXPERIMENTAL SETUP

  • Evaluation averages dataset-level ranks based on accuracy for classification and RMSE for regression; classification uses Taffy-4L and regression a separately trained model.Source: §5.1 · p. 7

5.2 OVERALL PERFORMANCE

  • TAFFY achieves the lowest average rank across all six classification and five regression benchmarks.Source: §5.2 · p. 7
  • Average rank falls from 2.40 to 2.17 with diversity, then to 2.07 with full TAFFY on TALENT classification.Source: §5.2 · p. 8, Figure 4 · p. 8

5.3 ANALYSIS OF CONTEXT CONSTRUCTION AND REFINEMENT

  • On TALENT classification, in-context diversity reaches lower mean rank than shuffling environments across separate single-environment tasks, with its advantage growing during training.Source: §5.3 · p. 8, Figure 3 · p. 8
  • Late-training task-specific gating beats global gating; α tracks categorical fraction negatively and median excess kurtosis positively, but not support-label entropy or weighted missingness stably.Source: §5.3 · pp. 8–9, Figure 3 · p. 8
  • Matched depth improves rank from 4.70 to 3.95 on TALENT and 6.18 to 5.62 on OpenML from two to four loops, while PFN and TabArena favor three.Source: Figure 3 · p. 8, §5.3 · p. 9
  • With two-loop training, a third inference loop helps OpenML and PFN but not TALENT or TabArena; across 428 shared entries, four loops worsen aggregate rank to 4.23.Source: Figure 3 · p. 8, §5.3 · p. 9

5.4 PRETRAINING DATA EFFICIENCY

  • At comparable budgets, Taffy-4L attains higher accuracy than TabICLv2, while Taffy-3L continues improving after TabICLv2 reaches its best; same-step accuracy orders Taffy-4L > Taffy-3L > Taffy-2L.Source: §5.4 · p. 9

6 CONCLUSION

TAFFY combines diverse synthetic contexts with task-adaptive iterative refinement; under the stated protocols, its reported configurations achieve the lowest average rank across six classification and five regression suites.

6 CONCLUSION

  • Across six classification and five regression suites, the reported configurations attain the lowest average rank on all eleven benchmarks under the stated protocols.The conclusion identifies context composition and task-adaptive iterative refinement as complementary design axes.Source: §6 · p. 9
  • The In-Context Diversity Prior places related environments in one pretraining task, exposing informative variation for learning predictive relationships from context.The Task-Conditioned Looped Transformer repeatedly integrates that evidence and uses support statistics to regulate the final refinement.Source: §6 · p. 9
  • Controlled comparisons favor in-context diversity, task-specific gating, and greater trained loop depth.Source: §6 · p. 9

A SYNTHETIC TASK GENERATION

The synthetic-task generator builds related environments from a shared causal process, using three distribution shifts and two interventions before combining environment-specific support and query data into tasks.

A SYNTHETIC TASK GENERATION

A.1 BASE CAUSAL PROTOTYPE

  • Each task begins with a causal model whose environments share its graph, node ordering, and readout, differing only in sampling distributions or selected mechanisms.The base joint distribution is induced by evaluating structural equations in topological order and applying the readout.Source: §A.1 · p. 15

A.2 DISTRIBUTION SHIFTS

  • Distribution shifts draw candidate rows from the base model and use reweighting calibrated by bisection to a common 0.1-nat KL target; unattainable targets trigger resampling.A shift strength of κ = 0.1 is described as mild; for example, it moves a balanced binary label distribution from 50/50 to approximately 72/28.Source: §A.2 · pp. 15–16

A.2.1–A.2.3 DISTRIBUTION SHIFTS

  • Label shift changes class proportions, covariate shift concentrates sampling around a random feature-space region, and conditional shift changes features within classes while preserving class proportions.For regression, labels are discretized into ten bins for label and conditional shifts; the conditional-shift invariances apply to environment distributions, not necessarily finite samples.Source: §A.2.1 · p. 16, §A.2.2 · p. 16, §A.2.3 · pp. 16–17

A.3 INTERVENTIONS

  • Interventions target non-label nodes and regenerate data under modified mechanisms, allowing effects to propagate to descendants, including the label when it is a descendant.Hard interventions replace a target equation with a constant, whereas soft interventions resample its mechanism while retaining parental dependence.Source: §A.3 · p. 17, §A.3.1 · p. 17, §A.3.2 · p. 17

A.4 CONTEXT ASSEMBLY AND MODEL INPUTS

  • A multi-environment task samples two to four environments, allocates equal support and query rows per environment, and mixes them in the same proportions in both sets.Half of pretraining tasks use the unmodified single-environment process; the model receives support features and labels plus query features, but not environment identifiers or causal information.Source: §A.4 · pp. 17–18

B.1 SUPPORT STATISTICS

  • The task-specific support gate uses a 51-entry descriptor whose feature summaries are aggregated across valid features by mean, standard deviation, minimum, and maximum.Padded features are excluded, and nonfinite descriptor entries are replaced by zero.Source: §B.1 · p. 18, Table 4 · p. 19

C.1 BASELINE ROSTER / C.2 DATA SPLITS AND PREPROCESSING

  • The evaluation roster groups baselines into tree-based, deep-tabular, foundation-model, and ensemble methods; preprocessing is fitted on support data, and benchmark suites may overlap.Source: §C.1 · p. 19, §C.2 · p. 19

C.3 TRAINING AND INFERENCE SETTINGS

The appendix describes separate classification and regression pretraining configurations, inference protocols, and the added cost associated with deeper looped models.

C.3 TRAINING AND INFERENCE SETTINGS

  • Classification pretraining uses 64 AMD MI210 GPUs, a global batch size of 1,024 tasks, and a configured 25,000-step schedule with a maximum context length of 4,096 rows.The schedule uses Muon, a peak learning rate of 6 × 10^-4, weight decay 0.01, gradient clipping at 10, label smoothing 0.02, and warmup followed by cosine decay.Source: §C.3 · pp. 19–20
  • Regression uses a separate model with quantile outputs and pinball-loss training, while the primary comparison uses Taffy-4L for classification and the separately trained regression model without per-dataset selection.The recorded inference protocols also vary by model; frozen-parameter inference does not imply identical inference-compute budgets.Source: §C.3 · p. 20

C.4 PRETRAINING COST

  • 1.80×: Taffy-4L requires 1.80× the single-pass backbone time per training step, with regression at comparable cost under the same setup.The cost is measured on 64 AMD MI210 GPUs; total time and GPU-hours are projected for the 25,000-step schedule.Source: §C.4 · p. 20, Table 5 · p. 20

C.5 METRICS AND AGGREGATION

The appendix specifies rank-based aggregation and details the reconstructed evaluation pool and two additional comparisons using matched or filtered datasets.

C.5 METRICS AND AGGREGATION

  • Methods are ranked per dataset using accuracy for classification and RMSE for regression, then averaged within each suite; lower average rank is better.Ties receive their average rank, and missing results are neither imputed nor treated as losses.Source: §C.5 · p. 20

C.6 PRIMARY EVALUATION POOL

  • The primary evaluation reconstructs TabICLv2's development-task exclusions from public materials because no machine-readable exclusion list is provided; missing results are not interpolated.Classification also removes TALENT tasks with more than ten classes from Table 1 and reports them separately.Source: §C.6 · p. 21

C.7 MANY-CLASS TALENT TASKS

  • Taffy has the best average rank and Elo among nine methods with observed results on all TALENT classification datasets with more than ten classes.The evaluation uses hierarchical classification with mixed-radix ensembling.Source: §C.7 · p. 21, Table 6 · p. 21

C.8 LIMIX-COMPATIBLE EVALUATION POOL

  • Taffy has the lowest overall average rank in the pool filtered using LimiX-2M dataset eligibility criteria.The analysis recomputes accuracy- and RMSE-based ranks using the complete applicable-method subset in each suite.Source: §C.8 · p. 21

C.9 PRETRAINING DATA BUDGET

C.9 defines an input-data budget proxy and describes budget-eligible evaluation comparisons; C.10 reports Elo outcomes under two evaluation pools.

C.9 PRETRAINING DATA BUDGET

  • The input-data budget is cumulative maximum feature elements, with each step based on task count, maximum support-and-query rows, and maximum input-feature columns.Multistage budgets include the initializing checkpoint; loop count is excluded, and the proxy does not measure compute or wall-clock throughput.Source: §C.9 · p. 21
  • Table 7 reports average ranks on the LimiX-2M-eligible, sample-compatible evaluation pool, where lower rank is better.Rank aggregation uses the complete applicable-method subset within each suite; dashes indicate incomplete task-type coverage.Source: Table 7 · p. 22
  • Figure 1(d) compares measured checkpoints at matched or lower input-data budgets, but inference settings are not fully controlled across models.The budget proxy is not a measure of compute or wall-clock throughput.Source: §C.9 · pp. 21–22

C.10 ELO COMPARISON

  • TAFFY has the highest mean Elo under both the primary evaluation pool and the LimiX-2M eligibility rule.The average is the mean of all eleven unrounded ratings, not a pooled Elo fit.Source: §C.10 · p. 23

C.11 GATE ANALYSIS DETAILS

The gate analysis relates reconstructed task-specific coefficients to dataset properties and tests the gate by changing only its coefficient at inference. The analyses report associations with task properties and a lowest average rank for the task-specific gate on matched TALENT datasets.

C.11 GATE ANALYSIS DETAILS

  • The gate analysis reconstructs coefficients from late-training checkpoints using saved FP32 parameters and statistics of processed training-plus-validation support data.It averages coefficients across two preprocessing views and uses TALENT datasets with at most 20 classes and one million support cells.Source: §C.11 · p. 23
  • The association analysis adjusts descriptor and gate midranks for support rows, effective columns, class count, and support-label entropy.Kurtosis analyses additionally control for categorical fraction; datasets without numerical columns are excluded from that analysis.Source: §C.11 · p. 24
  • Across three model depths, α is negatively associated with categorical fraction and positively associated with numerical excess kurtosis.The analysis reports pointwise 95% bootstrap intervals and applies separate multiple-testing corrections for the two property analyses.Source: §C.11 · p. 24

C.11.1 SUPPORT ROWS AND EFFECTIVE COLUMNS

  • Across TALENT datasets, α is positively associated with support-row count but shows no consistent association with effective column count.Table 11 reports unadjusted associations, with q-values corrected over the original 216 tests.Source: §C.11.1 · p. 25, Table 11 · p. 25

C.12 FROZEN-CHECKPOINT GATE INTERVENTION

  • On matched TALENT classification datasets, the task-specific gate has the lowest average rank among the three inference conditions.The frozen-checkpoint comparison holds model parameters and inference settings fixed while changing only the update coefficient.Source: §C.12 · p. 25, Table 12 · p. 25

D LIMITATIONS

TAFFY uses a single task-level scalar gate for the final refinement, which limits how finely refinement strength can vary.

D LIMITATIONS

  • The single task-level scalar gate cannot vary refinement strength across examples, features, layers, or earlier passes.Comparisons with finer-grained gating under matched parameter and compute budgets are identified as an important extension.Source: D LIMITATIONS · p. 25
Loading 2610.07559v1…