How TabPFN Learns: Retrieval and Feature Selection
Abstract
Tabular foundation models such as TabPFN per- form in-context learning by conditioning on labeled training examples at inference time. In this work, we aim to conceptualize how TabPFN-family models organize information at the attention-head level during inference. We show that a subset of datapoint-attention heads implements implicit class-conditional retrieval; selectively attending to training examples that share the test example’s class label. To quantify this behavior, we introduce the adaptive-k Class Alignment Score (CAS). We further show through targeted multi-head ablations that these retrieval heads are causally load-bearing for model perfor- mance. Extending the analysis to feature atten- tion, we identify highly selective feature-attention heads using a Feature Selectivity Score (FSS), and demonstrate that ablating heads based on FSS ranking causes sharp performance degra- dation. Across NanoTabPFN and TabPFN v2, our results suggest that tabular in-context learn- ers internally organize computation into special- ized retrieval and feature-selection circuits. These findings provide a more mechanistic understand- ing of how tabular foundation models perform in-context learning.