EquiTabPFN: A Target-Permutation Equivariant Prior Fitted Networks

TL;DR

EquiTabPFN introduces a fully target-equivariant architecture for tabular data models, improving stability and performance on unseen class datasets by eliminating the equivariance gap.

Contribution

This paper presents a novel target-equivariant architecture for tabular models, addressing the equivariance gap and enhancing generalization and efficiency.

Findings

Matches or surpasses existing methods on classification benchmarks with many classes

Reduces computational overhead compared to prior models

Eliminates the equivariance gap to improve prediction stability

Abstract

Recent foundational models for tabular data, such as TabPFN, excel at adapting to new tasks via in-context learning, but remain constrained to a fixed, pre-defined number of target dimensions-often necessitating costly ensembling strategies. We trace this constraint to a deeper architectural shortcoming: these models lack target equivariance, so that permuting target dimension orderings alters their predictions. This deficiency gives rise to an irreducible "equivariance gap", an error term that introduces instability in predictions. We eliminate this gap by designing a fully target-equivariant architecture-ensuring permutation invariance via equivariant encoders, decoders, and a bi-attention mechanism. Empirical evaluation on standard classification benchmarks shows that, on datasets with more classes than those seen during pre-training, our model matches or surpasses existing methods while incurring lower computational overhead.