Towards Localization via Data Embedding for TabPFN

Mykhailo Koshil, Thomas Nagler, Matthias Feurer, Katharina Eggensperger

Keywords: PFN, TabPFN, In-Context Learning, Representation Learning

TL;DR: We scale TabPFN to arbitrary dataset sizes via localization.

Abstract: Prior-data fitted networks (PFNs), especially TabPFN, have shown significant promise in tabular data prediction. However, their scalability is limited by the quadratic complexity of the transformer architecture's attention across training points. In this work, we propose a method to localize TabPFN, which embeds data points into a learned representation and performs nearest neighbor selection in this space. We evaluate it across six datasets, demonstrating its superior performance over standard TabPFN when scaling to larger datasets. We also explore its design choices and analyze the bias-variance trade-off of this localization method, showing that it reduces bias while maintaining manageable variance. This work opens up a pathway for scaling TabPFN to arbitrarily large tabular datasets.

Submission Number: 72