TabPFN-Wide: Continued Pre-Training for Extreme Feature Counts

TL;DR

The paper addresses the challenge of high-dimensional, low-sample (HDLSS) tabular data in biomedicine by extending an existing tabular foundation model (TabPFNv2) through continued pre-training on synthetic HDLSS data. It introduces a causal prior-based HDLSS data generation process with feature widening to tens of thousands of features and demonstrates that TabPFN-Wide maintains or improves predictive performance while remaining robust to noise and avoiding feature reduction. A key finding is that feature-wise attention scores correlate with predictive importance, enabling inherent interpretability, including biologically relevant gene signals in cancer datasets. Overall, TabPFN-Wide represents a scalable, interpretable approach for HDLSS tabular data and opens avenues for adapting other foundation models to extreme feature counts, with practical impact for biomedical discovery and precision medicine.

Abstract

Revealing novel insights from the relationship between molecular measurements and pathology remains a very impactful application of machine learning in biomedicine. Data in this domain typically contain only a few observations but thousands of potentially noisy features, posing challenges for conventional machine learning approaches. While prior-data fitted networks emerge as foundation models for tabular data, they are currently not suited to handle large feature counts (>500). Although feature reduction enables their application, it hinders feature importance analysis. We propose a strategy that extends existing models through continued pre-training on synthetic data sampled from a customized prior. The resulting model, TabPFN-Wide, matches or exceeds its base model's performance while exhibiting improved robustness to noise. It seamlessly scales beyond 50,000 features, regardless of noise levels, while maintaining inherent interpretability, which is critical for biomedical applications. Our results show that prior-informed adaptation is suitable to enhance the capability of foundation models for high-dimensional data. On real-world biomedical datasets many of the most relevant features identified by the model overlap with previous biological findings, while others propose potential starting points for future studies.

Figures (18)

• Figure 1: The performance of existing tabular foundation models decreases for a selected high-dimensional biomedical dataset. Further datasets are presented in \ref{['subsec:real_world_results']} to confirm generality.
• Figure 2: Feature Widening
• Figure 3: Feature correlation maps for (a) mRNA gene expression data and (b–f) synthetically generated datasets with different sparsity values $p$. We compute Pearson correlation for $100$ randomly sampled features and sort them by average absolute correlation.
• Figure 4: Average relative performance to random forest (pink) for up to 4 multiomics datasets. Higher is better. See Appendix \ref{['app:all_omic_dataset_results']} for detailed results.
• Figure 5: (a) AUROC for TabPFN-Wide-5k vs TabPFNv2 on $21$ TabArena classification tasks with $\le10{,}000$ samples and $\le500$ features. (b-c) Average AUROC (relative to TabPFNv2 evaluated on the original dataset) on a set of $13$ widened datasets: (b) needle-in-a-haystack and (c) features-smearing (see text for further details). (d) TabPFN-Wide-5k's performance for different sparsities. $p=0$ corresponds to TabPFN-Wide-5k's curve in (b), and $p=0.02$ in (c).