Source-linked AI summary
Mitra: Mixed Synthetic Priors for Enhancing Tabular Foundation Models
Xiyuan Zhang, Danielle C. Maddix, Junming Yin, Nick Erickson, Abdul Fatir Ansari, Boran Han, Shuai Zhang, Leman Akoglu, Christos Faloutsos, Michael W. Mahoney, Cuixiong Hu, Huzefa Rangwala, George Karypis, Bernie Wang
TL;DR
Synthetic priors enable tabular foundation models to generalize through in-context learning, but principles for designing effective priors and mixtures remain poorly understood. This paper analyzes prior performance, diversity, and distinctiveness, then introduces Mitra, which uses a curated mixture of structural causal and tree-based priors. Mitra consistently outperforms existing TFMs across classification and regression benchmarks and is more sample-efficient, while remaining weaker than TabPFNv2 on some large-feature regression tasks.
Problem
Prior TFM work relies largely on fixed or heuristic synthetic priors, leaving open which prior properties support generalization and how priors should be combined.
Method
The paper systematically evaluates prior performance, diversity, and distinctiveness, then pretrains Mitra on a model-agnostic mixture of structural causal and tree-based priors.
Results
Mitra consistently outperforms existing TFMs across classification and regression tasks and demonstrates better sample efficiency with fewer in-context examples.
Takeaways & Limitations
Prior-mixture analysis provides a basis for a diverse, high-performing, model-agnostic TFM that generalizes across tabular tasks and attention architectures.
Takeaways & Limitations
MITRA does not consistently outperform TabPFNv2 on large-feature regression tasks, motivating larger-scale pretraining for high-dimensional settings.
Abstract
from arXiv · showhide
Since the seminal work of TabPFN, research on tabular foundation models (TFMs) based on in-context learning (ICL) has challenged long-standing paradigms in machine learning. Without seeing any real-world data, models pretrained on purely synthetic datasets generalize remarkably well across diverse datasets, often using only a moderate number of in-context examples. This shifts the focus in tabular machine learning from model architecture design to the design of synthetic datasets, or, more precisely, to the prior distributions that generate them. Yet the guiding principles for prior design remain poorly understood. This work marks the first attempt to address the gap. We systematically investigate and identify key properties of synthetic priors that allow pretrained TFMs to generalize well. Based on these insights, we introduce Mitra, a TFM trained on a curated mixture of synthetic priors selected for their diversity, distinctiveness, and performance on real-world tabular data. Mitra consistently outperforms state-of-the-art TFMs, such as TabPFNv2 and TabICL, across both classification and regression benchmarks, with better sample efficiency.
1 Introduction
Tabular foundation models use synthetic-data pretraining and in-context learning to transfer across datasets, shifting attention toward the design of synthetic priors. Mitra studies which prior properties support generalization and builds a mixture that improves performance across tasks and attention architectures.
- Background: TFMs pretrained on synthetic tabular tasks make downstream predictions by conditioning on labeled examples from the target dataset.Synthetic data provides broad task coverage and supports adaptation without large amounts of real-world downstream data.
- Motivation: Existing TFM work leaves open what makes a synthetic prior effective and how priors should be mixed for better generalization.Prior efforts commonly rely on fixed or heuristic data priors while emphasizing architectural innovation.
- Contribution: Mitra identifies standalone real-data performance, diversity, and within-mixture distinctiveness as key properties of effective synthetic priors.Diversity concerns resistance to overfitting the prior’s own distribution, while distinctiveness concerns transfer between different priors.
- Method: The proposed mixture combines structural causal models with gradient boosting, random forest, decision tree, and extra tree priors.SCMs are selected for diversity and standalone performance, while tree-based priors contribute distinctiveness.
- Results: Mitra achieves state-of-the-art results across classification and regression, improves strong TFM baselines, and shows better sample efficiency with fewer in-context examples.The mixture is model-agnostic and improves both row-based 1D and element-based 2D attention architectures.
2 Related Work
Related work on tabular prediction spans traditional tree-based models, semantically rich language-model approaches, and ICL-based tabular foundation models. ICL-based work primarily advances synthetic priors or model architectures.
- Traditional and deep learning-based models: Gradient-boosted decision trees and ensemble systems have historically dominated tabular prediction because of strong performance, robustness, interpretability, and practical usability.These methods are generally tailored to individual datasets and have limited transfer across datasets.
- TFMs: Semantically-Rich Models: Semantically rich approaches serialize tables into text or adapt language models for zero-shot and few-shot tabular inference.Examples include TabLLM, GTL, Tabula, and TP-BERTa.
- TFMs: ICL-Based Models: ICL-based tabular models pretrain on synthetic or real datasets and use downstream datasets as in-context examples.This line of work focuses mainly on designing better data priors and improving model architectures.
3 MITRA: TFM Pretrained on a Mixture of Priors
Mitra pretrains a model-agnostic TFM on a mixture of structural causal and tree-based synthetic priors selected for standalone performance, diversity, and distinctiveness. The mixture is evaluated through cross-prior generalizability and real-data performance, then used to generate in-context training tasks.
- Prior characterization: Effective priors are characterized by strong real-dataset performance, diversity, and distinctiveness within the mixture.These criteria jointly guide the development of MITRA.
- Prior construction: A data-generating prior samples hyperparameters such as feature size, sample count, class count, and categorical-feature count before producing a synthetic dataset.The dataset contains feature-label pairs with continuous or categorical attributes.
- Prior construction: MITRA mixes structural causal models with tree-based priors including decision tree, extra tree, gradient boosting, and random forest generators.SCMs generate features and targets through randomly constructed causal graphs, while tree-based generators model conditional targets using threshold-based splits.
- Prior analysis: The generalizability matrix records cross-prior evaluation, while the performance vector records each prior’s performance on real-world datasets.Higher performance-vector values indicate stronger real-data performance; lower diagonal and off-diagonal matrix values indicate greater diversity and distinctiveness, respectively.
- Prior analysis: The illustrative prior analysis links high prior quality to both strong real-data performance and favorable diversity or distinctiveness properties.The example highlights GSCM’s combined strengths and shows why a prior with high standalone performance need not be the best mixture addition.
- Pretraining: The mixture samples support and query entries from generated tables and trains on the likelihood of masked query labels using either 1D row-wise or 2D element-wise attention.The pipeline is model-agnostic across these attention architectures.
4 Empirical Results
MITRA is evaluated across classification, regression, model architectures, sample sizes, ensembling strategies, and fine-tuning settings. It consistently achieves strong or state-of-the-art performance, including under limited data and competitive ensemble budgets.
- Classification: MITRA wins across three merged classification benchmarks and achieves the best performance with fine-tuning and ensembling under ranking-based and aggregated metrics.The merged evaluation contains 137 unique datasets spanning TabRepo, TabZilla, and AMLB.
- Regression: MITRA achieves the best performance across the TabRepo regression benchmark, showing that the prior mixture is task-agnostic.
- Model agnosticism: MITRA 1D outperforms TabPFN and TabForestPFN when all models use the same prior mixture, supporting model-agnostic performance gains.The comparison is made under the capabilities available to each baseline, including differences in native ensemble and fine-tuning logic.
- Sample efficiency: MITRA consistently performs better than TabPFNv2 and TabICL when classification ICL examples are reduced to 10%, 25%, 50%, or 75% of the original size.The reported improvement with ensemble and fine-tuning is attributed to increased prior diversity.
- Ensembling: MITRA (bagging) consistently outperforms TabPFNv2 PHE and the AutoGluon ensemble across 300-, 600-, 900-, and 3600-second training budgets.The bagging ensemble fine-tunes one MITRA instance per fold of an 8-fold stratified cross-validation ensemble and averages predictions.
- Fine-tuning: MITRA shows better fine-tuning performance across ensemble sizes, whereas fine-tuning and ensembling barely improve TabPFNv2's ensemble-alone performance.The authors hypothesize that MITRA's diverse priors and adaptation from a maximum of 16 pretraining features contribute to this result.
5 Conclusion and Discussion
The paper frames tabular foundation-model performance around synthetic-prior design and describes MITRA's mixture-based data-generation approach. Its generation process supports varied feature types and target-generation mechanisms for synthetic pretraining.
- Conclusion and Discussion: MITRA combines structural causal-model and tree-based priors to cover diverse real-world tabular distributions.The mixture includes gradient boosting, random forest, decision tree, and extra-tree priors.
- Data generation and ICL: Synthetic pretraining samples datasets from a prior distribution with varying feature counts, sample counts, class counts, and categorical attributes.After pretraining, the model predicts query labels from labeled support rows and query features in a single forward pass without updating parameters.
- Feature and target generation: MITRA generates features using continuous Gaussian components and categorical components whose class counts follow a geometric distribution.The categorical proportion and feature indices are sampled during dataset generation.
- Feature and target generation: Targets are generated either directly from p(y | x) or indirectly by fitting a classifier or regressor on a synthetic training dataset.
A.2.2 Synthetic Data-Generating Priors
MITRA combines structural causal models and tree-based priors, using both indirectly and directly sampled procedures to generate synthetic tabular datasets for pretraining.
- Prior mixture: MITRA mixes SCMs with tree-based priors, including gradient boosting, random forests, decision trees, and extra trees.The mixture includes indirectly sampled ET, GB, DT, and RF priors, alongside directly sampled SCM and DSRF priors.
- Indirect sampling: Indirectly sampled priors fit a task-specific estimator on generated training data, then predict targets for newly sampled features.For classification, MITRA fits a classifier; for regression, it fits a regressor.
- Generated data: The section presents generated two-dimensional and high-dimensional datasets whose features come from the MITRA prior mixture and whose classification labels are shown by color.The accompanying figures visualize mixture-generated data and compare training accuracy across pretraining steps.
- Direct sampling: Directly sampled priors first sample a function from a function space and then generate targets by evaluating that function on sampled features.In DSRF, the sampled function is represented by randomly constructed decision trees and evaluated through tree traversal.
- Direct sampling: DSRF constructs random trees by sampling feature split indices, thresholds, terminal-node behavior, and task-dependent leaf targets.Classification leaves receive cyclically assigned class labels, while regression leaves are sampled from a Gaussian distribution.
- Pretraining procedure: MITRA pretraining repeatedly samples a generator, partitions its generated table into support and query sets, predicts query labels, and updates model parameters.The population likelihood decomposes across conditionally independent query rows; classification uses cross-entropy and regression uses MSE.
B.1 Benchmarking Datasets
The evaluation uses diverse classification and regression benchmarks and combines standard aggregated metrics with ranking-based metrics to compare models across datasets.
- Benchmarking datasets: Small-feature evaluation covers 66 TabRepo classification, 75 TabZilla classification, and 10 TabRepo regression datasets with up to 3,000 rows and 100 features.The benchmark design also distinguishes small-feature datasets from large-feature benchmarks for comparing 1D and 2D models.
- Ranking-based metrics: Ranking-based evaluation reports average rank, Elo, winrate, rescaled accuracy, and champion delta to capture relative performance across datasets.The paper uses these metrics because aggregated scores can be disproportionately influenced by a small number of extreme datasets.
- Ranking-based metrics: Rescaled accuracy adjusts for dataset difficulty by scaling performance relative to the best- and worst-performing models on each dataset.Its underlying error uses 1−AUC for classification and 1−R2 for regression.
- Aggregated metrics: Classification is evaluated with AUC, accuracy, and cross-entropy, while regression uses R2, RMSE, and MAE.These aggregated metrics are averaged across datasets.
- Implementation: MITRA uses a 12-layer, 512-dimensional Transformer with four attention heads, row- and column-wise attention, and 72M parameters.The MITRA 1D variant uses row-wise attention and has 37M parameters.
C.1 Generalizability Matrix and Performance Vector Metrics
The generalizability matrix evaluates models trained on each prior against data from every prior, while the performance vector evaluates them on real-world datasets.
- Metric construction: Each matrix entry Gij records a TFM’s performance when pretrained on prior Gi and evaluated on data generated from prior Gj.The supplied table descriptions specify ACC for Table 9 and CE for Table 10.
- Findings: The accuracy results further emphasize TBP distinctiveness: the off-diagonal Gij for a model pretrained on GET is 0.577.The passage identifies this as an off-diagonal generalization result rather than a standalone-prior score.
- Metric construction: Each performance-vector element Pi records the performance of a TFM pretrained on prior Gi when evaluated on real-world datasets.The table descriptions define this using ACC in Table 9 and CE in Table 10.
C.2 Sample Efficiency
MITRA’s more diverse prior mixture improves sample efficiency relative to ablations, with the largest Elo gains appearing under the strongest downsampling.
- Comparison: MITRA is compared with MITRA-Mix2, using SCM + ET + GB, and Attic, a variant using SCM + DT.The comparison evaluates performance across multiple downsampling ratios.
- Results: The Elo differences increase as the downsampling ratio decreases, with the largest gains occurring at ds = 0.1.The result is presented as evidence of stronger generalizability for MITRA’s broader prior mixture in data-scarce scenarios.
C.3 Ablations
The ablation studies evaluate priors individually and incrementally, showing that SCM provides the strongest standalone performance while tree-based priors add diversity and distinctiveness. On TabRepo, mixtures with balanced SCM and TBP weighting outperform either prior family alone.
- The forward ablation ranks priors as {SCM, ET, GB, DT, RF, DSRF}, aligning with performance, diversity, and distinctiveness analyses.
- On TabRepo, SCM and TBP mixtures at p = 0.5, 0.4, and 0.6 tie for the best average ranking of 3.2.
- At p = 0.5, the mixture achieves Elo 1040, compared with 1030 for the p = 0.4 and p = 0.6 variants.
- SCM alone outperforms TBPs alone on TabRepo, with average rankings of 3.7 versus 4.4.
- The ablation attributes the mixture’s benefit to SCM’s performance and TBPs’ distinctiveness within the prior mixture.
C.4 Classification
MITRA achieves strong classification performance across benchmark datasets and remains competitive on broader evaluations, while bagging further improves its results. Timing measurements also show a favorable performance–runtime trade-off against CatBoost.
- MITRA (+ef) has the best performance across TabRepo, TabZilla, and AMLB classification benchmarks.
- MITRA (bagging) further improves MITRA (+ef) and produces larger gains over other baselines.
- On TabArena, MITRA is the strongest single model and remains state-of-the-art, while TabPFNv2 with HPO and ensembling is significantly slower.
- Critical-difference plots compare MITRA and baselines across three classification and two regression benchmarks.
- MITRA (+ef) gains 138 Elo over CatBoost while running approximately 3.5× faster on TabRepo.
C.8 Decision Boundary Visualizations
Decision-boundary visualizations examine MITRA and baselines on representative simulated distributions using few-shot ICL. MITRA generally produces regular boundaries and shows robustness to noise, while some smooth or spiral structures favor Gaussian Processes.
- The visualization protocol uses 1,000 samples per dataset, with 10% support samples and 90% query samples.
- On axis-aligned data, MITRA produces more regular and less fragmented decision boundaries than TabPFNv2.
- On spiral data, MITRA is an exception to its mostly comparable or superior generalization, but it is more robust as noise increases.
- MITRA’s boundaries fall between those of tree-based classifiers and TabPFNv2 across several representative distributions.
- Gaussian Process classifiers perform strongly on spiral and Swiss roll examples, motivating future inclusion of GP-based priors.
D Limitations and Future Work
The paper identifies several directions for extending MITRA, including adaptive mixture weighting, additional continuous priors, and larger-scale pretraining. A current boundary is weaker consistency on large-feature regression tasks.
- Hyperparameter optimization could adapt mixture weights to specific downstream tasks or domains.
- Future mixtures may incorporate Gaussian Processes to model smooth boundaries and improve generalization beyond tabular tasks such as time-series forecasting.
- MITRA does not consistently outperform TabPFNv2 on large-feature regression tasks.
- The planned response to large-feature regression limitations is scaling pretraining to datasets with more rows and features.