Robust Tabular Foundation Models

AI summary

Overview

Research area: Machine learning for tabular (structured) data, specifically tabular foundation models (TFMs) that learn via in-context learning and are pretrained on synthetic datasets.

Technical level: Advanced. The paper builds a distributionally robust optimization (DRO) formulation over structural causal model (SCM) parameters and proves a softmax form for the optimal adversary distribution, then applies it as a fine-tuning pipeline.

Scope in one sentence: The paper proposes RTFM, a model-agnostic adversarial training framework that adapts the synthetic data generator toward regions where a tabular foundation model underperforms strong baselines, and demonstrates it by improving TabPFN V2 on two tabular benchmarks.

What This Paper Is About

TFMs such as TabPFN V2 are pretrained on large quantities of synthetic datasets drawn from structural causal models, using a fixed prior distribution over the generator parameters. Because that prior is fixed, some regions of the parameter space are underrepresented, and the resulting model can still lag behind boosted-tree methods on datasets with those characteristics. The paper reframes this as an adversarial robustness problem: keep a model-agnostic training loop that searches for generator settings where the model has the largest gap relative to strong baselines (XGBoost, CatBoost, Random Forests, and others), then train the model preferentially on data drawn from those regions.

Key Contributions

1. Formalizing adversarial training over the SCM parameter space. The authors define a parameterized joint distribution over generator hyperparameters (e.g., number of layers, hidden size, activation function, ratio of categorical features), introduce an optimality gap measure, and use it to target regions where the TFM underperforms relative to the best achievable performance.

2. A model-agnostic two-stage adversarial training algorithm. RTFM alternates a maximization stage (black-box parameter search over generator configurations, using a frozen model) with a minimization stage (model training on datasets sampled from a softmax distribution over the discovered parameters).

3. A practical estimator for the optimality gap plus a theoretical result. Because the conditional entropy of a generator is hard to compute, the gap is lower-bounded using the minimum cross-entropy loss over a set of fitted baseline estimators. The paper also proves that the optimal maximizer distribution in the entropy-constrained DRO problem is a softmax.

4. An empirical demonstration on TabPFN V2. Applying RTFM to TabPFN V2 improves benchmark performance, with up to a 6% increase in mean normalized AUC over the original TabPFN and other baseline algorithms, using fewer than 100k additional synthetic datasets.

Main Findings

• RTFM improves TabPFN V2 on both benchmarks. On TabPertNet the mean normalized AUC rises from 0.7483 for TabPFN (no ensembling) to 0.8167 for TabPFN (RTFM), with mean rank AUC improving from 3.2 to 2.7 and rank-1 wins increasing from 11 to 17. On TabArena the mean normalized AUC (OVO) rises from 0.9031 to 0.9298, with mean rank AUC (OVO) improving from 2.2 to 1.9 and rank-1 wins increasing from 5 to 12.

• RTFM outperforms all baseline estimators on the reported metrics. The comparison table includes Logistic Regression, MLP, Random Forest, CatBoost, and XGBoost. For example, on TabPertNet the best baseline mean normalized AUC shown is 0.6663 (CatBoost), and on TabArena the best baseline shown is 0.7749 (CatBoost), both below the RTFM numbers. RTFM also has the best mean rank AUC and rank-1 wins on both benchmarks.

• Improvements are statistically significant. A Friedman test on median normalized AUC scores gave p = 2.2 × 10⁻¹⁴ for the TabPertNet datasets and p = 1.6 × 10⁻¹² for TabArena. Wilcoxon signed-rank tests comparing TabPFN (RTFM) to the original TabPFN gave p = 0.0023 (TabPertNet) and p = 0.0103 (TabArena); the paper states all other Wilcoxon tests yielded smaller p-values.

• Gains concentrate on datasets where the original model was weakest. For TabPertNet, among datasets where the original TabPFN has rank > 2, TabPFN (RTFM) has a mean rank of 3.26 — the highest among all models and a 1.24 point improvement over the original. On 21% of those datasets, where the original TabPFN is strictly worse than the baseline models, TabPFN (RTFM) becomes the top ranked model.

• Similar "leapfrogging" occurs on TabArena. Among datasets where the original TabPFN has rank > 1, RTFM shows a 1.8 point improvement in mean rank and becomes the top ranked model on 20% of those datasets.

• The extra training data is a small fraction of the original pretraining budget. These results were achieved with 90k additional synthetic datasets, described as about 1% of the more than 9 million datasets generated to pretrain the original TabPFN model.

• The adversary does find larger gaps after the original model is added as a baseline. After epoch 5, the original TabPFN is incorporated as an additional baseline model, and the parameter search uncovers larger optimality gaps for a few epochs, suggesting the model had some performance degradation in certain regions during the initial training epochs.

Methodology in Plain English

Pretraining a TFM means generating millions of synthetic tables by sampling from randomly initialized SCMs — in this work, randomized multi-layer perceptrons. Each SCM is drawn from a distribution whose parameters (mean hidden size, mean number of layers, categorical ratio, and so on) are themselves settings that can be tuned.

RTFM treats those generator settings as the adversary's action. First, in the maximization stage, the current model is frozen and a black-box optimizer proposes generator configurations. For each proposed configuration, the authors generate a fixed number of datasets, evaluate the model's cross-entropy loss on each, and compare it against the best loss achieved by several fitted baseline estimators. The difference is the estimated optimality gap. Because each (dataset, baseline) pair can be fitted independently, this step parallelizes trivially.

Second, in the minimization stage, the searched configurations are turned into a sampling distribution: parameters with larger estimated gaps get exponentially higher sampling probability (a softmax), with the temperature chosen by one-dimensional root finding so that the distribution's entropy equals a fraction of its maximum possible entropy. Training batches are generated by drawing a configuration from this distribution and then drawing a dataset from it, so the model's loss is effectively reweighted toward hard regions. The model is trained for a fixed number of steps, then the loop returns to the maximization stage with updated weights, repeating until the objective converges. After a set number of max-min iterations (five epochs in this implementation), the original pretrained model is added as one more baseline, which the authors describe as a way to mitigate any unlearning.

The experiments use a learning rate of 1e-5, batch size 64, 3000 training steps per iteration, 100 search trials, and 20 sampled datasets per trial, run for 30 max-min epochs on a single node with one A100 GPU and 256 CPU cores. TFM results are reported without ensembling, and all algorithms use their default settings.

Why This Matters

Impact on research. The paper shows that the synthetic data generator itself can be treated as a tunable object for targeted robustness training, rather than a fixed prior to be hand-crafted. It connects TFM pretraining to the adversarial training and distributionally robust optimization literature, provides a proof that the entropy-constrained adversary takes a softmax form, and demonstrates that a small amount of additional synthetic data (90k datasets, about 1% of the original pretraining set) can shift benchmark performance.

Real-world applications.

• Enterprise tabular workloads such as credit scoring, fraud detection, and risk modeling, where boosted trees are the incumbent and a model that matches or beats them on hard dataset slices has direct value.
• Any pipeline where a foundation model must be fine-tuned without access to proprietary or sensitive customer data, since RTFM trains only on synthetic datasets.
• Settings where inference speed matters, because TFMs make zero-shot predictions in milliseconds when GPU-accelerated, as noted in the paper.
• Benchmark suites and model-selection workflows, where the optimality-gap diagnostic can identify dataset regimes (for example, high categorical ratios or high missingness) where a model is weak.

Industry relevance. Tabular data remains the dominant data type in applied machine learning, and the paper's framing — use synthetic data alone to steer a foundation model toward the regions where it loses to gradient-boosted trees — is an inexpensive, model-agnostic recipe that does not require collecting new real data.

Future Directions

• Broaden the generator space. The authors state they plan to expand the parameter space and include tree-based SCMs in the framework; the current experiments only use MLP-based SCMs.
• Apply RTFM to other TFMs. The paper notes the framework is model-agnostic and could be applied to any TFM, naming Mitra and TabICL explicitly, not just TabPFN.
• Extend to regression. The framework is stated to be easily extended to regression tasks, though this work focuses on classification.
• Understand and control degradation. The finding that larger optimality gaps appear once the original model is added as a baseline suggests open questions about how much a model unlearns in specific parameter regions during adversarial training and how best to prevent it.
• Scale and refine the search. Because the pipeline is highly parallelizable in both dataset generation and parameter search, larger-scale parameter searches and more baseline estimators with varying strengths and inductive biases remain to be explored.

Target Audience

Researchers and practitioners working on tabular machine learning, foundation models, in-context learning, and synthetic data generation will get the most from this paper. It is also relevant to engineers deploying tabular models who want a fine-tuning method that does not require real training data, and to readers interested in adversarial training and distributionally robust optimization applied to model pretraining. A background in optimization and probabilistic modeling helps, since the theoretical sections derive a softmax adversary distribution from an entropy-constrained DRO problem.

Authors’ abstract

The development of tabular foundation models (TFMs) has accelerated in recent years, showing strong potential to outperform traditional ML methods for structured data. A key finding is that TFMs can be pretrained entirely on synthetic datasets, opening opportunities to design data generators that encourage desirable model properties. Prior work has mainly focused on crafting high-quality priors over generators to improve overall pretraining performance. Our insight is that parameterizing the generator distribution enables an adversarial robustness perspective: during training, we can adapt the generator to emphasize datasets that are particularly challenging for the model. We formalize this by introducing an optimality gap measure, given by the difference between TFM performance and the best achievable performance as estimated by strong baselines such as XGBoost, CatBoost, and Random Forests. Building on this idea, we propose Robust Tabular Foundation Models (RTFM), a model-agnostic adversarial training framework. Applied to the TabPFN V2 classifier, RTFM improves benchmark performance, with up to a 6% increase in mean normalized AUC over the original TabPFN and other baseline algorithms, while requiring less than 100k additional synthetic datasets. These results highlight a promising new direction for targeted adversarial training and fine-tuning of TFMs using synthetic data alone.