Causal Data Augmentation for Robust Fine-Tuning of Tabular Foundation Models

TL;DR

The paper tackles the challenge of fine-tuning tabular foundation models when labeled data are scarce, where traditional validation-based early stopping can be unreliable. It introduces CausalMixFT, which learns a Structural Causal Model from the target data using PC/FCI causal discovery and DoWhy additive-noise SCMs to generate synthetically augmented samples that preserve causal dependencies, then trains on a 1:1 mix of real and synthetic data with validation performed on real data. Across 33 TabArena datasets and 2,310 fine-tuning runs, CausalMixFT achieves a median ROC-AUC improvement of $+0.12$ over the pre-trained baseline (versus $+0.10$ for standard fine-tuning) and outperforms purely statistical generators like CTGAN, TabEBM, and TableAugment, while reducing the validation-test gap from $0.67$ to $0.30$. This demonstrates that causal-structure-informed data augmentation provides a principled and effective regularization for low-data fine-tuning of tabular foundation models, enabling more stable and reliable adaptation in real-world scarce-label scenarios.

Abstract

Fine-tuning tabular foundation models (TFMs) under data scarcity is challenging, as early stopping on even scarcer validation data often fails to capture true generalization performance. We propose CausalMixFT, a method that enhances fine-tuning robustness and downstream performance by generating structurally consistent synthetic samples using Structural Causal Models (SCMs) fitted on the target dataset. This approach augments limited real data with causally informed synthetic examples, preserving feature dependencies while expanding training diversity. Evaluated across 33 classification datasets from TabArena and over 2300 fine-tuning runs, our CausalMixFT method consistently improves median normalized ROC-AUC from 0.10 (standard fine-tuning) to 0.12, outperforming purely statistical generators such as CTGAN (-0.01), TabEBM (-0.04), and TableAugment (-0.09). Moreover, it narrows the median validation-test performance correlation gap from 0.67 to 0.30, enabling more reliable validation-based early stopping, a key step toward improving fine-tuning stability under data scarcity. These results demonstrate that incorporating causal structure into data augmentation provides an effective and principled route to fine-tuning tabular foundation models in low-data regimes.

Figures (18)

• Figure 1: Overview of the SCM-augmented fine-tuning process. Real training data ($D^{\text{real}}_{\text{train}}$) are used to fit a Structural Causal Model (SCM) via PC/FCI and DoWhy dowhy. The SCM samples synthetic data ($D^{\text{syn}}_{\text{train}}$) that preserve the discovered causal dependencies among features. Real and synthetic samples are mixed in equal proportion to fine-tune the tabular foundation model (TFM), which is optimized by cross-entropy loss. Validation is performed only on real data, and fine-tuning continues as long as validation performance improves.
• Figure 2: Performance comparison across data generation strategies. (left) Normalized ROC-AUC improvements relative to the pre-trained baseline (dashed line). Whiskers represent $1.5\times$ IQR; medians and standard deviations are annotated. Higher score is better. (right) Critical difference diagram (significance level = 0.05) Lower rank is betterdemvsar2006ranking_method.
• Figure 3: Validation--test performance correlation across datasets and data generators. Each cell shows the Pearson correlation between validation and test log-loss for a given dataset and generator configuration. Low or negative correlations indicate that validation performance is not a reliable proxy for generalization under small-data conditions. The columns and rows are sorted by average correlation coefficients from left (higher) to right (lower) and top (higher) to bottom (lower). Incomplete runs due to the time limit of 1h or too many features are marked with "!" (TabEBM uses TabPFNv1 internally, which only allows for 100 features)
• Figure 4: Validation performance comparison across generators. The results reveal that validation-based ranking can be misleading under small-data conditions, with methods such as TableAugment and Default showing strong validation performance but large decrease in test performance and thus generalization. The CausalMixFT configuration demonstrates a smaller validation-test discrepancy, suggesting more stable and generalizable fine-tuning behavior.
• Figure 5: Normalized test ROC-AUC performance across datasets and generator configurations. Each cell reports the normalized ROC-AUC (mean ± standard deviation) for a given generator on a specific dataset. The observed heterogeneity indicates that fine-tuning performance varies considerably across generators and datasets, highlighting the need for dataset-adaptive augmentation strategies.