Mixture of In-Context Prompters for Tabular PFNs

## Abstract

Abstract Recent benchmarks found In-Context Learning (ICL) outperforms both deep learning and tree-based algorithms on small tabular datasets. However, on larger datasets, ICL for tabular learning cannot run without severely compromising performance, due to its quadratic space and time complexity w.r.t. dataset size. We propose MixturePFN , which both extends nearest-neighbor sampling to the state-of-the-art ICL for tabular learning model and uses bootstrapping to finetune said model on the inference-time dataset. MixturePFN is the Condorcet winner across 36 diverse tabular datasets against 19 strong deep learning and tree-based baselines, achieving the highest mean rank among Top-10 aforementioned algorithms with statistical significance.

## 1 Introduction

Tabular data is a popular data format across various domains, consisting of column-wise features and row-wise data samples. Each feature can be either continuous, categorical, or ordinal. Thanks to the prevalence of relational databases, which ensure data integrity, consistency, and low redundancy, tabular data is widely used across various domains such as medicine, finance, and advertising. Hence, improving learning algorithms on tabular data is of interest to many researchers.

General tabular datasets remain unconquered by most deep learning algorithms (Popov et al., ; Gorishniy et al., ; Somepalli et al., ; Arik and Pfister, ; Yamada et al., ; Yoon et al., ; Chen et al., ). Instead, gradient-boosted decision trees (GBDTs) (Chen and Guestrin, ; Prokhorenkova et al., ), achieve better overall performance on tabular benchmarks (McElfresh et al., ; Shwartz-Ziv and Armon, ) considering a wide range of number of samples, numbers of features, feature types, and feature distributions. Recently, transformer-based prior-fitted networks, PFNs (Hollmann et al., ), have garnered interest, for their surprisingly strong and state-of-the-art performance on tabular datasets with $\leq 3,000$ samples (Hollmann et al., ; McElfresh et al., ).

Unlike SGD-based deep learning, PFNs learn the training algorithm itself (Müller et al., ; Von Oswald et al., ). Specifically, a PFN first pretrains a model by sampling labelled datasets from a predefined dataset prior (Müller et al., ), then performs inferece using In-Context Learning (ICL) (Brown et al., ), where each downstream dataset is tokenized into a “prompt”, sent to the PFN model, which returns test split predictions, as described in Section 2.1. By learning the learning algorithm itself, PFNs discover better inductive bias than conventional deep learning algorithms and gradient-boosted decision trees, which require extensive hyperparameter tuning to effectively fit the downstream dataset (Hollmann et al., ).

PFNs cannot scale to datasets with $>3000$ samples without compromising performance due to two key limitations. First, PFN inference is computationally expensive. Because the entire dataset is fed to the transformer as a “prompt”, the inference time and space complexity for $N_{train}$ training samples is $\mathcal{O}(N_{train}^{2})$. This is in stark contrast to GBDTs and traditional deep learning approaches where the inference time and space complexity is $\mathcal{O}(1)$. To fit larger datasets in memory, existing works (Hollmann et al., ; McElfresh et al., ) resort to randomly sampling the training dataset during inference when $N_{train}>3000$ to ensure the “prompt” has at most 3000 samples. Second, existing PFNs assume the dataset prior used during pretraining is always representative of the dataset prior used during inference, which we empirically find is untrue. Specifically, downstream performance improves when the PFN is finetuned on how datasets are sampled during inference rather than during pretraining.

In this work, we analyze recent claims (McElfresh et al., ) on PFN’s effectiveness, finding it does not scale w.r.t. dataset size. We improve PFN’s scalability by proposing Sparse “Mixture of In-Context Prompters” (MICP), which creates scalable “prompts” by routing new test samples to a group of prompters. We solve PFN’s alignment limitations with “Context-Aware PFN” (CaPFN), which finetunes PFNs for downstream datasets via bootstrapping. We call our combined model MixturePFN. To summarize:

- •
To improve scalability, we are the first to propose Sparse Mixture of In-Context Prompters (MICP) which routes new test samples to a pool of scalable prompters for In-Context Learning. MICP reduces PFN inference complexity from $\mathcal{O}(N_{train}^{2})$ memory and time to $\mathcal{O}(1)$ memory and $\mathcal{O}(log(N_{train}))$ time, w.r.t. the number of training samples, $N_{train}$.
- •
To improve performance, we finetune Context-Aware Prior-Fitted Network (CaPFN), which aligns pretrained PFNs with inference-time datasets using a novel bootstrapping policy.
- •
MixturePFN scales transformer PFNs from tabular datasets of 3,000 samples to those with much larger number of samples, with no performance deterioration w.r.t. dataset size.
- •
MixturePFN is the Condorcet winner across 36 diverse tabular datasets against 19 strong deep learning and tree-based baselines, achieving the top mean rank among Top-10 aforementioned algorithms with statistical significance.

## 2 Preliminaries

We consider tabular classification problems, where the inputs are numerical, ordinal, or categorical columns encoded as a $d$-dimensional feature vector, $x\in\mathbb{R}^{d}$, the output is the corresponding label, $y\in[1,...,C]$, and the dataset consists of labelled input output pairs, $D=\{(x^{(i)},y^{(i)})\}_{i=0}^{N}$.(^1^11We provide a table with all math notations in the Supplementary Material.) Given the training dataset, $D_{train}=\{(x^{(i)}_{train},y^{(i)}_{train})\}_{i=0}^{N_{train}}$, and test samples, $X_{test}=[x_{test}^{(i)}]_{i=0}^{N_{test}}$,
our goal is to correctly predict the corresponding test labels, $Y_{test}=[y_{test}^{(i)}]_{i=0}^{N_{test}}$. MixturePFN is inspired by Prior Fitted Networks, which we first introduce.

### 2.1 Prior Fitted Networks

Prior Fitted Network (PFN) (Müller et al., ) is a parameterized model, $q_{\theta}$, that learns to approximate Bayesian inference given the dataset prior, $p(D)$, via In-Context Learning (ICL) (Brown et al., ). Specifically, PFN inference approximates the posterior predictive distribution (PPD), $p_{\theta}(y|x,D)=\int_{\phi}p(y|x,\phi)p(D|\phi)p(\phi)d\phi$, where $\phi$ is the hypothesis mechanism behind how the tabular data is generated. For example, $\phi$ can be a structural causal model. Refer to Muller et. al (Müller et al., ) for further details.

Figure: Figure 1: We highlight the differences between In-Context Learning (ICL) on Prior Fitted Networks (ex. TabPFN), left, and Large Language Models (LLMs), right. TabPFN treats training data as tokens (where each token is a concatenation of feature and label), whereas LLMs use templates to convert training data into natural language prompts. TabPFN uses an attention pattern (blue and red arrows) supporting batch inference, whereas LLMs use generic encoder-decoder or decoder-only setups. TabPFN are pretrained on Equation 1, whereas LLMs are pretrained on a separate objective.
Refer to caption: extracted/5619985/figures/_PFN_LLM.png

#### 2.1.1 Pretraining

To approximate the PPD, PFNs are pretrained to minimize KL-Divergence between the parameterized model, $q_{\theta}(y|x,D)$, and the PPD, $p(y|x,D)$, over the dataset prior, $p(D)$, which was proven equivalent to optimizing the prior data negative log likelihood, $\mathcal{L}_{\textsc{PFN}}$. As shown in Equation 1 (Müller et al., ), this loss iteratively samples new datasets from a handcrafted dataset prior, $p(D)$, via Monte-Carlo.

$$ $\mathcal{L}_{\textsc{PFN}}=\mathop{\mathbb{E}}_{\begin{subarray}{c}x,y,D\sim p (D)\end{subarray}}[-log(q_{\theta}(y|x,D))]$ (1) $$

TabPFN (Hollmann et al., ) is the state-of-the-art pretrained PFN transformer for tabular data. It treats the hypotheses, $\phi$, as randomly sampled structural causal models (SCM) (Pearl, ; Peters et al., ) mixed with the original Bayesian Neural Network prior (Müller et al., ). Training dataset samples are generated by first sampling a SCM graph, $\phi\sim p(\phi)$, followed by sampling the SCM, $x,y,D\sim p(D|\phi)$.

Transformer-based (Vaswani et al., ) PFNs tokenize the sampled dataset, $(x,D)$ as input to the parameterized model, $q_{\theta}$, as shown in Figure 1 and discussed in Section 2.1.2. Note, PFN inputs are analogous to “prompts” from In-Context Learning (ICL) (Brown et al., ; Dong et al., ; Xu et al., ), hence they are called “prompts” in this work.

#### 2.1.2 Inference

During inference, transformer-based PFNs tokenize the downstream dataset, $(X_{test},D_{train})$, into batched “prompts”, consisting of $N_{train}$ encoder tokens and $N_{batch}$ decoder tokens, where each data sample corresponds with one token.(^2^22Our dataset is split into train/dev/test sets. During hyperparameter tuning, decoder tokens are taken from the dev set instead.) Because tabular columns are permutation invariant, TabPFN shuffles feature orderings and scalings, running $q_{\theta}$ on each permutation of the “prompt”, then returning an ensembled prediction, as depicted in Figure 1 and further detailed in Section 2.1.4. TabPFN does not perform finetuning, only inference, on downstream datasets.

#### 2.1.3 Fundamental Scalability Limitations

TabPFN cannot scale to datasets with large numbers of training samples, $N_{train}\geq 3,000$ (McElfresh et al., ). Because the “prompt” contains $N_{train}$ encoder tokens, the transformer’s, $q_{\theta}$, inference space and time complexities are quadratic w.r.t. dataset size, $\mathcal{O}(N_{train}^{2})$, which is too computationally expensive to run on large datasets.(^3^33While linear transformer approximations are $\mathcal{O}(N_{train})$, our goal is to scale PFNs to be comparable to other predictive algorithms which have $\mathcal{O}(1)$ space and time complexity w.r.t. the number of training samples.) Conventional deep learning (Popov et al., ; Arik and Pfister, ; Gorishniy et al., ) and GBDT (Chen and Guestrin, ; Prokhorenkova et al., ) algorithms have $\mathcal{O}(1)$ inference space and time complexities w.r.t. number of training samples. Our proposal, MixturePFN, reduces the space complexity of TabPFN inference to $\mathcal{O}(1)$ and time complexity to $\mathcal{O}(log(N_{train}))$, effectively scaling PFNs to larger datasets.

#### 2.1.4 PFN -Style Batching

Batching in TabPFN is unlike in Large Language Models (LLMs). LLMs for In-Context Learning (ICL) can fit $N_{batch}$ test samples across $N_{batch}$ “prompts” (Liu et al., ), where each test sample has its own “prompt”. TabPFN (Hollmann et al., ) must fit $N_{batch}$ test samples in $1$ “prompt”, because the ensembling process will augment each “prompt” into $N_{ensemble}$ inputs through shuffling feature orderings and scalings. We focus on TabPFN-style batching which fits $N_{batch}$ test samples in $1$ “prompt”.

Figure: Figure 2: Illustration of MixturePFN. MICP (Left): New test samples are passed to a router that picks 1 out of $K$ prompters to form a scalable “prompt” with $B$ training samples for the downstream PFN model. CaPFN (Right): TabPFN is frozen, fitted with adapters, then finetuned using data prior negative loss likelihood, Equation 1, on our bootstrapped data prior, $p(D|D_{train})$. This prior simulates the MICP inference mechanism. The finetuned model is called CaPFN.
Refer to caption: extracted/5619985/figures/capfn_micp.png

## 3 Method

### 3.1 Support Set Approximation for Scalable PFNs

To improve space and time complexity, MixturePFN hypothesizes each test sample, $x_{test}^{(i)}$, requires only a small support set, $D_{supp}(x_{test}^{(i)})\subset D_{train}$, of constant size, $|D_{supp}(x_{test}^{(i)})|=B$ $\forall i=[0,...,N_{test}-1]$, to effectively perform inference with In-Context Learning. This assumption is reasonable as training samples with drastically different features from the test sample should have little impact on its label (Khandelwal et al., ; Liu et al., ; Xu et al., ; Feuer et al., ). For example, the purchase trends of today share more in common with purchase trends from last week than the those from 10 years ago.

We define the support set, $D_{supp}$, for each individual test sample, $x_{test}^{(i)}$, to be the $B$ spatially closest training samples, which could be found via K-Nearest Neighbors: $D_{supp}(x)=\textsc{KNN}(x|D_{train},B)$. A scalable PFN “prompt” can thus be constructed with the test sample and support set, $(\{x_{test}^{(i)}\},D_{supp}(x_{test}^{(i)}))$. We call this approach KNN-Prompting (^4^44Further illustrations of KNN-Prompting and its relation to ICL for LLMs can be found in the Appendix.).

PFN inference on KNN prompts, $(\{x_{test}^{(i)}\},D_{supp}(x_{test}^{(i)}))$, can be computed in $\mathcal{O}(1)$ space and time complexity w.r.t. $N_{train}$ for fixed $B$. However, each test sample, $x_{test}^{(i)}$, may require a different support set, $D_{supp}(x_{test}^{(i)})$, and hence its own “prompt”. Thus, KNN-Prompting does not support TabPFN-style batching, where multiple test cases fit in 1 “prompt”. Furthermore, KNN-Prompting runs an expensive KNN search across the whole training dataset for each test sample. For these 2 reasons, KNN-Prompting is too expensive in practice when evaluating on larger tabular datasets.

### 3.2 Mixture of In-Context Prompters ( MICP )

In order to efficiently construct effective “prompts”, MixturePFN leverages cluster structures in the data. Inspired by Sparse Mixture of Experts (Shazeer et al., ; Lewis et al., ), where each test sample is routed to an specialized expert trained on a subset of the training dataset, Sparse Mixture of In-Context Prompters, MICP, routes each test sample to one of $K$ “In-Context Prompters” (ICP), $\{\mathcal{T}_{k}\}_{k=0}^{K}$, specializing on a cluster of the training dataset, using a routing module, $\mathcal{R}:\mathbb{R}^{d}\rightarrow\{0,...,K-1\}$. Each ICP then constructs a relevant “prompt” for incoming test samples, which are sent to the downstream PFN model in batch.

To reduce the “prompt”-construction overhead, each ICP, $\{\mathcal{T}_{k}\}_{k=0}^{K}$, precomputes its own support set, $D_{prompt}(k)\subset D_{train}$, of constant size, $|D_{prompt}(k)|=B$. During inference, ICP concatenates incoming test samples with its support set to form the scalable “prompt”: $\mathcal{T}_{k}(\{x_{test}^{(i)}:\mathcal{R}(x_{test}^{(i)}=k)\})=(\{x_{test}^
{(i)}:\mathcal{R}(x_{test}^{(i)}=k)\},D_{prompt}(k))$. The goal of MICP is efficiently approximate KNN-prompts, which can be accomplished by maximizing the overlap between the prompter’s and the test sample’s support sets: $|D_{prompt}(\mathcal{R}(x_{test}^{(i)}))\cap D_{supp}(x_{test}^{(i)})|$ $\forall i\in[0,...,N_{test}-1]$.

#### 3.2.1 Router and Prompter Initialization

Intuitively, each ICP, $\mathcal{T}_{k}$, specializes on a local cluster of the training dataset. The router, $\mathcal{R}$, routes each test sample to its nearest cluster. To capture local clusters, we initialize the router and ICP with K-Means on the training dataset: $\{D_{cluster}^{(k)}\}_{k=0}^{K},\{x_{center}^{(k)}\}_{k=0}^{K}=\textsc{KMeans}
(D_{train})$.
Our resulting router is the Nearest Neighbor Search algorithm: $\mathcal{R}(x)=\textsc{NNS}(x|\{x_{center}^{(k)}\}_{k=0}^{K})$.

Because efficient “prompts” have bounded number of entries, $B$, we subsample clusters with more than $B$ entries. Because PFN performance drastically increases with more training samples (Hollmann et al., ; Müller et al., ; McElfresh et al., ), we expand clusters with less than $B$ entries via K-Nearest Neighbors. We define this process in Equation 2, where $G(k)=|D_{cluster}^{(k)}|<B$ denotes whether the clusters are too big or small.

$$ $D_{prompt}(k)=\\ \begin{cases}\textsc{KNN}(x_{center}^{(k)}|D_{train},B),&\text{if }G(k)\\ \textsc{Sample}(D_{cluster}^{(k)},B),&\text{else}\end{cases}$ (2) $$

During inference, MICP routes each test sample, $x_{test}^{(i)}$, to its corresponding ICP. Similar to sparse mixture of experts, once each ICP receives enough test entries, PFN inference is performed on batched test “prompts”: $(X_{batch},D_{prompt}^{(k)})$, which contains multiple entries from the test set $X_{batch}\subseteq X_{test}$, because all said entries were routed to the same ICP cluster: $\mathcal{R}(x_{batch}^{(i)})=\mathcal{R}(x_{batch}^{(j)})\forall i,j$. Hence, MICP efficiently constructs bounded batched “prompts”.

**Table 1: MixturePFN is the Condorcet winner across 36 datasets against 19 baseline algorithms. MixturePFN achieves the top mean rank across 20 datasets where all algorithms successfully run and across 22 datasets where all Top-10 algorithms successfully run. To break ties, we rank algorithms based on their mean log-likelihoods following TabZilla (McElfresh et al., ). We report the Condorcet matrix, dataset breakdowns, and accuracy-metric results in the Appendix.**
| Method | Condorcet Statistics | All Algo. | Top-10 Algo. |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| #Votes$\uparrow$ | #Wins$\uparrow$ | #Ties | #Losses$\downarrow$ | Mean $\pm$ Std Rank$\downarrow$ |  |  |
| MixturePFN | 524 | 19 | 0 | 0 | 2.350 $\pm$ 1.824 | 2.273 $\pm$ 1.7106 |
| XGBoost | 500 | 18 | 0 | 1 | 5.500 $\pm$ 4.621 | 4.000 $\pm$ 2.663 |
| CatBoost | 474 | 17 | 0 | 2 | 4.900 $\pm$ 4.158 | 3.955 $\pm$ 2.688 |
| SAINT | 408 | 16 | 0 | 3 | 8.300 $\pm$ 5.367 | 4.045 $\pm$ 1.965 |
| TabPFN* | 381 | 13 | 1 | 5 | 4.550 $\pm$ 2.747 | 4.040 $\pm$ 1.311 |
| LightGBM | 373 | 14 | 1 | 4 | 9.150 $\pm$ 4.351 | 6.409 $\pm$ 2.839 |
| DANet | 312 | 14 | 0 | 5 | 9.050 $\pm$ 3.369 | 7.045 $\pm$ 1.988 |
| FTTransformer | 294 | 12 | 0 | 7 | 8.600 $\pm$ 3.541 | 6.773 $\pm$ 2.235 |
| ResNet | 286 | 11 | 0 | 8 | 8.400 $\pm$ 3.262 | 6.864 $\pm$ 1.961 |
| SVM | 285 | 9 | 0 | 10 | 11.300 $\pm$ 4.766 | 7.500 $\pm$ 2.482 |
| STG | 284 | 10 | 0 | 9 | 11.900 $\pm$ 4.549 | - |
| RandomForest | 247 | 7 | 0 | 12 | 11.600 $\pm$ 4.443 | - |
| NODE | 243 | 7 | 0 | 12 | 13.350 $\pm$ 3.410 | - |
| MLP-rtdl | 227 | 5 | 0 | 14 | 10.800 $\pm$ 5.046 | - |
| TabNet | 210 | 5 | 0 | 14 | 13.550 $\pm$ 5.296 | - |
| LinearModel | 202 | 3 | 1 | 15 | 12.400 $\pm$ 4.652 | - |
| MLP | 191 | 5 | 1 | 13 | 13.700 $\pm$ 3.621 | - |
| VIME | 134 | 2 | 0 | 17 | 15.350 $\pm$ 3.851 | - |
| DecisionTree | 114 | 1 | 0 | 18 | 16.800 $\pm$ 3.881 | - |
| KNN | 74 | 0 | 0 | 19 | 18.450 $\pm$ 1.936 | - |

In total, router and prompter initialization takes $\mathcal{O}(tN_{train}K+(N_{train}+KB)logN_{train})$ time and $\mathcal{O}(N_{train}+KB)$ space complexity and is done once before inference. Routing takes $\mathcal{O}(log(K))$ time and $\mathcal{O}(1)$ space complexity, using efficient nearest neighbor search with ball-tree for each test sample. PFN transformer inference takes $\mathcal{O}(B^{2}+BN_{batch})$ time and space complexity, as MICP prompts contain at most $B$ training samples and $N_{batch}=|X_{batch}|$ testing samples. We provide time and space complexity details in the Appendix. We illustrate MICP in Figure 2.

#### 3.2.2 Efficiency and Effectiveness Trade-Off

The effectiveness of MICP prompts depend on the number of ICPs used, $K$. As the complexity and size of data increase, more ICPs are needed to capture the entropy of the labels. This is natural as each router’s support set, $D_{prompt}(\mathcal{R}(x_{test}^{(i)}))$, should be representative of test samples routed to that cluster, $D_{support}(x_{test}^{(i)})$. If the true support set $D_{support}(x_{test}^{(i)})$ becomes more granular as the dataset size increases, more ICPs are required to maximize overlap: $|D_{prompt}(\mathcal{R}(x_{test}^{(i)}))\cap D_{supp}(x_{test}^{(i)})|$.

We theoretically characterize this relationship between $K$, $B$, and overlap by analyzing conditions required for nonzero overlap on the training data: $|D_{prompt}(\mathcal{R}(x_{train}^{(i)}))\cap D_{supp}(x_{train}^{(i)})|\geq 1$ $\forall i\in[0,...,N_{train}-1]$. Specifically, we encourage nonzero overlap by scaling the number of “prompts”, $K$, linearly with the size of each “prompt”, $B$, and training dataset size, $|N_{train}|$, as stated in Theorem 1: $K\geq\lceil N_{train}/B\rceil$.

This insight allows MixturePFN to trade-off efficiency and effectiveness with a single hyperparameter, $\gamma$, which controls the number of ICPs as a ratio of training and support set sizes: $K=\lceil\gamma N_{train}/B\rceil$. Intuitively, larger $\gamma$ improves effectiveness at the cost of efficiency. Assuming fixed $\gamma$, $N_{batch}$, and $B$, MixturePFN routing takes $\mathcal{O}(log(N_{train}))$ time and $\mathcal{O}(1)$ space complexity, and PFN inference takes $\mathcal{O}(1)$ time and space complexity.

###### Theorem 1 (Nonzero Overlap) .

If every K-Means cluster contains at most $B$ samples, $|D_{cluster}^{(k)}|\leq B$ $\forall k\in[0,...,K-1]$ and training points route to their assigned K-Means cluster $\mathcal{R}^{*}(x_{train}^{(i)})=k:x_{train}^{(i)}\in D_{cluster}^{(k)}$ (^5^55Thse conditions can be satisfied via constrained K-Means (Bradley et al., ), which ensures each cluster has at most $B$ entries, and a router that sends train points to their assigned clusters. In practice, we find the relationship with the tunable parameter $\gamma$ also holds for MixturePFN’s regular K-Means and Nearest-Neighbor Search router.), then nonzero overlap on the training data is guaranteed, $|D_{prompt}(\mathcal{R}^{*}(x_{train}^{(i)}))\cap D_{supp}(x_{train}^{(i)})|\geq
1$ $\forall i\in[0,...,N_{train}-1]$ $\forall D_{train}$.

Figure: (a) $N_{train}$ w.r.t. TabPFN*
Refer to caption: extracted/5619985/figures_new/improv_nll.png

###### Theorem 1 (Nonzero Overlap) .

### 3.3 Context-Aware Finetuning ( CaPFN )

PFNs are pretrained on the ICL task over a synthetic dataset prior, $p(D)=p(D|\phi)p(\phi)$ (Müller et al., ; Hollmann et al., ). Inspired by recent works which aligns Large Language Models on ICL “prompts” via finetuning (Thoppilan et al., ; Wei et al., ; Gu et al., ), we argue the pretraining data prior, $p(D)=p(D|\phi)p(\phi)$, is different than the true data generating mechanism during inference, $p(D_{prompt}|D_{train})$, which was described in Section 3.2. To better align the parameterized model, $q_{\theta}$, with the inference-time dataset, $D_{prompt}$, CaPFN uses bootstrapping on the downstream dataset, $D_{train}$, to simulate ICL “prompts”: $(X_{subtest},Y_{subtest},D_{subtrain})\sim p(D|D_{train})$, where $X_{subtest}\subset X_{train}$, $Y_{subtest}\subset Y_{train}$, and $D_{subtrain}\subset D_{train}$. Bootstrapped samples are used to tune adapters (Houlsby et al., ) via prior data negative log likelihood loss, as shown in Equation 1, except the dataset prior is now the bootstrap mechanism: $p(D)=p(D|D_{train})$.

#### 3.3.1 Bootstrapping Large MICP Datasets

The bootstrap procedure mimics MICP on large $N_{train}>3000$ datasets: $p(D|D_{train})=p(D_{support}|x)p(x|D_{train})$. Specifically, we sample a random training point from the training dataset, $x\sim p(x|D_{train})$, then run K-Nearest Neighbors from the sampled point, $p(D_{support}|x)=\textsc{KNN}(x|D_{train},B)$, as defined in Section 3.1, to obtain a bootstrap dataset, $D_{bootstrap}$. We randomly split the bootstrapped dataset $D_{bootstrap}\sim p(D|D_{train})$ into train/test splits to obtain the “labelled prompt”, $(X_{subtest},Y_{subtest},D_{subtrain})$.

#### 3.3.2 Bootstrapping Small Datasets

MICP does not run on smaller $N_{train}\leq 3000$ datasets. However, bootstrapping can still be used to finetune the model on said datasets to match the downstream dataset distribution. In this case, we sample from $p(D|D_{train})$ by randomly sampling $90\%$ of training samples without replacement to obtain $D_{subtrain}$ and treating the remaining $10\%$ of sample as $X_{subtest},Y_{subtest}$.

#### 3.3.3 Finetuning with Adapters

To prevent overfitting, we only train a small set of new adapter (Houlsby et al., ; Bapna et al., ; Hu et al., ; Liu et al., ) parameters, $\psi$, on $p(D|D_{train})$, without modifying in the pretrained transformer’s parameters, $\theta$.(^6^66Adapters are also efficient because only a small number of parameters are updated, $p(\phi)$ is not needed during finetuning, and different downstream datasets share a common pretrained model.) Specifically, we freeze a pretrained TabPFN transformer, $q_{\theta}(y|x,D)$. Next, for each downstream dataset, $D_{train}$, we add linear adapter layers (Houlsby et al., ), $\mathcal{A}_{\psi}^{(D_{train})}$, with parameters $\psi$, to form $q_{\theta,\psi}^{(D_{train})}(y|x,D,q_{\theta},\mathcal{A}_{\psi}^{(D_{train})})$. During finetuning, only $\psi$ is optimized. Intuitively, $q_{\theta}$ encodes the handcrafted prior, $p(D|\phi)p(\phi)$, and $\mathcal{A}_{\psi}^{(D_{train})}$ encodes the bootstrapped prior, $p(D|D_{train})$. We illustrate CaPFN in Figure 2.

Figure: Figure 4: Wilcoxon-Signed Rank Test shows MixturePFN significantly outperforms the Top-10 baselines on the 22 shared datasets. To break ties, we rank algorithms based on their mean log-likelihoods following TabZilla (McElfresh et al., ). We compute the rank across all 10 cross-validation splits. We report additional critical difference diagrams in the Appendix.
Refer to caption: extracted/5619985/figures/_critical_diff2.png

**Table 2: Ablation table results. MixturePFN (KNNv1) and MixturePFN (KNNv2) replace MICP with a scalable variant of KNN-Prompting. MixturePFN (CaPFN w. Full FT) uses full finetuning instead of adapters. MixturePFN (-CaPFN) and MixturePFN (-CaPFN-MICP) remove each component iteratively, where MICP is replaced by random sampling.**
| Method | Mean $\pm$ Std | Median | Min | Max |
| --- | --- | --- | --- | --- |
| Rank$\downarrow$ | Rank$\downarrow$ | Rank$\downarrow$ | Rank$\downarrow$ |  |
| MixturePFN | 2.75 $\pm$ 1.94 | 2 | 1 | 9 |
| MixturePFN (KNNv2) | 4.00 $\pm$ 1.56 | 4 | 2 | 7 |
| CatBoost | 4.70 $\pm$ 3.16 | 6 | 1 | 9 |
| MixturePFN (CaPFN w. Full FT) | 4.85 $\pm$ 2.03 | 4.5 | 2 | 9 |
| XGBoost | 4.95 $\pm$ 3.17 | 6 | 1 | 10 |
| MixturePFN (-CaPFN) | 5.10 $\pm$ 1.12 | 5 | 3 | 8 |
| MixturePFN (-CaPFN-MICP) = TabPFN* | 5.30 $\pm$ 1.59 | 5 | 2 | 9 |
| MixturePFN (KNNv1) | 7.25 $\pm$ 3.70 | 9 | 1 | 10 |
| MLP-rtdl | 7.25 $\pm$ 2.79 | 8 | 1 | 10 |
| MLP | 8.85 $\pm$ 0.99 | 9 | 7 | 10 |

## 4 Experiment Setup

We evaluate MixturePFN on the recently proposed TabZilla benchmark (McElfresh et al., ). TabZilla is the largest tabular benchmark, with 36 hardest datasets out of 176 tabular classification datasets and 19 baseline algorithms, covering both deep learning and GBDTs. The benchmark covers a diverse range of dataset properties, in number of samples, number of features, and feature distributions. MixturePFN’s goal is to (1) improve TabPFN* (McElfresh et al., ), which randomly samples $B$ training pairs so that TabPFN (Hollmann et al., ) runs on larger datasets, and (2) outperform both GBDTs (Chen and Guestrin, ; Prokhorenkova et al., ; Ke et al., ) which were found state-of-the-art by TabZilla, and recent deep learning models (Popov et al., ; Gorishniy et al., ; Arik and Pfister, ; Hollmann et al., ; Yamada et al., ; Yoon et al., ; Somepalli et al., ; Chen et al., ).

#### 4.0.1 Evaluation Protocol

Since TabZilla restricts the total runtime to 10 hours, not all algorithms run on the same datasets. To ensure a fair comparison(^7^77We cannot follow the same experimental settings as the October revision of TabZilla because they are unfair, as mentioned in a recent Github issue (abvesa, ). Further details are in the Appendix.), we first evaluate MixturePFN and baselines using Condorcet voting, where each dataset ranks the algorithms that successfully ran on it. An algorithm receives a vote whenever it achieves a higher rank than a baseline on each dataset. For each pair of algorithms, the winner received more votes than the loser across all datasets ranking both algorithms. The Condorcet winner is the algorithm that wins all pairwise comparisons.

After determining the Condorcet winner, We further evaluate MixturePFN by ranking it against all baseline algorithms and Top-10 Condorcet algorithms across their shared datasets. We evaluate the performance of each algorithm by its mean rank. Finally, we run a Wilcoxon-Signed Rank Test to check the statistical significance of said mean rank across the Top-10 Condorcet algorithms.

## 5 Results

### 5.1 MixturePFN : State-of-the-Art Performance

As shown in Table 1, MixturePFN achieves state-of-the-art performance on TabZilla across 36 datasets and 19 baseline algorithms. Specifically, MixturePFN is the Condorcet winner, receiving both the most votes and beating all other baselines in pairwise comparisons. Furthermore, MixturePFN achieves the top mean-rank across all subsets of fairly chosen datasets, followed by GBDTs (Chen and Guestrin, ; Prokhorenkova et al., ), then TabPFN* (McElfresh et al., ), then deep learning algorithms (Chen et al., ; Gorishniy et al., ; Somepalli et al., ). These results corroborate recent findings (McElfresh et al., ; Grinsztajn et al., ) that most deep learning algorithms fail on general tabular datasets. MixturePFN achieves its state-of-the-art results by scaling TabPFN*’s impressive performance to larger datasets. We provide additional metrics in the Appendix. To understand what dataset regimes each algorithm performs best at, we evaluate MixturePFN’s rankings w.r.t. dataset properties.

MixturePFN substantially improves the scalability of TabPFN and TabPFN*. As show in Figure 3a, unlike TabPFN which encounters memory bottlenecks on datasets with $>3000$ samples, MixturePFN successfully runs on all said datasets. As show in Figure 3a, MixturePFN substantially improves the performance of TabPFN* by improving how samples are chosen for the “prompt” and training on the downstream dataset.

MixturePFN encounters no performance deterioration w.r.t. dataset size. As shown in Figure 3b, unlike TabPFN*, whose performance deteriorates w.r.t dataset size, MixturePFN’s performance compared to the next best baseline is not correlated with dataset size. Hence, MixturePFN is necessary to scale TabPFN*’s impressive performance on to larger datasets.

MixturePFN is robust to irregular datasets. We measure the irregularity of datasets using the standard deviation of the kurtosis of all features. Deep learning algorithms are especially susceptible to irregular datasets with uninformative or heavy-tail features (Grinsztajn et al., ). As shown in Figure 3c, MixturePFN is the Top-1 algorithms on datasets with both high and low kurtosis standard deviation. Because it is finetuned on downstream datasets, MixturePFN is robust to dataset irregularity.

MixturePFN works better with fewer features. As shown in Figure 3c, MixturePFN loses against baselines on datasets with a large number of features. PFN transformers are known to face scalability challenges with number of features (Hollmann et al., ), due to their handling of column order invariance. We believe better tokenization practices and feature selection can improve feature size scalability, and leave such exploration to future work.

MixturePFN’s state-of-the-art performance is statistically significant. Specifically, we run the Wilcoxon Signed-Rank test with $p<0.05$ comparing the Top-10 Condorcet algorithms from Table 1 across their 22 shared datasets and 10 cross-validation splits. As shown in Figure 4, MixturePFN’s state-of-the-art performance is statistically significant.

### 5.2 Ablation Study

Both MICP and CaPFN contribute to MixturePFN state-of-the-art results. We perform ablation studies for MICP and CaPFN against common GBDTs and deep learning models across 10 shared datasets. As shown in Table 2, each component of the model, MICP and CaPFN, contributes to achieving state-of-the-art results. MICP helps by efficiently choosing an effective context for the “prompt”. CaPFN helps by aligning the dataset prior through finetuning the PFN on MICP’s prompting policy. Because overfitting is a well-known issue for deep learning models tackling tabular data (Kadra et al., ; Grinsztajn et al., ), adapters are a key component to ensure CaPFN aligns the pretrained TabPFN transformer with the downstream data, without destroying its pretraining prior, $p(D|\phi)p(\phi)$.

Under the same GPU resources, KNN-Prompting is much less effective than MICP on tabular datasets. As described in Section 2.1.4, KNN-Prompting does not support TabPFN-style batching. To empirically verify that MICP improves KNN-Prompting, we replace MICP in MixturePFN with 2 KNN-Prompting variants: MixturePFN (KNNv1): Because each prompt contains at most $B+N_{batch}$ tokens, we batch KNN-Prompts by considering $B/N_{batch}$ nearest neighbors instead of $B$-nearest neighbors; MixturePFN (KNNv2): Because LLM-batching fails due to TabPFN’s ensembling overhead, we remove the ensembling procedure and run KNN-Prompting following LLM-batching (Liu et al., ), as described in Section 2.1.4. As shown in Table 2, both KNN-Prompting variants perform substantially worse than MixturePFN because they compromise an essential component of PFN, either prompt size or ensembling, for performing efficient inference. The relative rankings suggest prompt size matters more than ensembling.

## 6 Limitations

As shown in Section 5.1, MixturePFN successfully improves TabPFN’s scalability w.r.t. dataset size to achieve state-of-the-art results on the TabZilla benchmark. However, we notice that TabPFN* does not scale well with feature and label count, relying on ensembling to capture feature and label order invariance. Scaling TabPFN* to datasets with large number of features and labels can further push ICL performance for tabular learning. While this work covers a large number of diverse datasets, we do not cover huge datasets with billions of samples (Zhu et al., ; Yang et al., ). We leave such investigation to future work.

## 7 Conclusion

In this work, we provide a scalable framework for In-Context Learning (ICL) on tabular datasets. To efficiently construct effective ICL “prompts”, we propose routing test samples through a Sparse Mixture of In-Context Prompters, MICP. To align the PFN with the inference-time datasets, we propose a novel finetuning policy using bootstrapping, CaPFN. Our framework scales PFNs from datasets with 3000 samples to those with much larger number of samples. Our framework, MixturePFN, achieves state-of-the-art performance against 19 deep learning and tree-based baselines across 36 general tabular datasets, establishing a new standard for general tabular learning.

## 8 Related Work

### 8.1 Tabular Learning Algorithms

Early tabular learning algorithms are based off decision trees, utilizing boosting, feature encoding, and ensembling [Shwartz-Ziv and Armon, , Borisov et al., , Chen and Guestrin, ]. Early deep learning algorithms are inspired by decision trees, making them end-to-end learnable [Popov et al., , Katzir et al., , Hazimeh et al., , Somepalli et al., , Arik and Pfister, ]; however, more thorough benchmarks find decison trees produce more reliable results. Hyperparameter tuning [Kadra et al., ] and inductive bias [Grinsztajn et al., ] were identified as key weaknesses in deep learning algorithms. Recent works focus on optimizing transformer models for specialized datasets [Huang et al., , Gorishniy et al., , ], and improving decision tree optimization [Joseph and Raj, ], or Bayesian learning [Hollmann et al., , Schäfl et al., , Feuer et al., ] for small datasets. With the rise of LLMs, pretrained tabular learning models also achieve impressive performance at the cost of billions of training points [Yang et al., , Zhu et al., ]. Of these methods, Prior-Fitted Networks [Feuer et al., ] were identified as a promising direction in recent benchmarks among deep learning approaches for general tabular learning problems [McElfresh et al., ].

### 8.2 Gradient-Boosted Decision Trees

Gradient-boosted decision trees (GBDTs) remain the preferred algorithm of tabular learning practitioners [Chen and Guestrin, , Prokhorenkova et al., ]. Deep learning algorithms [Popov et al., , Gorishniy et al., , Somepalli et al., , Arik and Pfister, , Yamada et al., , Yoon et al., ] fail on larger benchmarks considering different numbers of samples, numbers of features, feature types, feature distributions, and numbers of labels [McElfresh et al., , Shwartz-Ziv and Armon, ].
Because GBDTs utilize boosted gradients and are not rotationally invariant before training, GBDT learning algorithms have a better inductive bias than Stochastic Gradient Descent-based (SGD-based) deep learning algorithms Grinsztajn et al. []. Thus, GBDTs achieve state-of-the-art performance on medium to large datasets with 3,000 to 1,000,000 samples and competitive performance on smaller datasets McElfresh et al. [], Grinsztajn et al. [].

### 8.3 Prior Fitted Networks

Prior Fitted Networks (PFN) [Müller et al., ] approximates Bayesian inference using a data prior, where a parameterized model is trained to minimize the KL-Divergence between it and the posterior predictive distribution. The proof for PFNs is derived from meta-learning [Gordon et al., ]. PFNs fall under In-Context Learning [Brown et al., ], as the entire training dataset is fed to the model during inference. Hence, PFNs effectively learn the learning algorithm. This approach is particularly effective on tabular data [McElfresh et al., ], where the data prior effectively regularizes model predictions. TabPFN [Hollmann et al., ] is a PFN model specific to tabular data that achieves state-of-the-art performance on small datasets. Unlike preliminary works on scaling TabPFN [Feuer et al., ], that benchmark KMeans and Coreset “promting”, we propose a sparse mixture of KNN prompters, capable of forming scalable batched “prompts” and a novel finetuning protocol for the mixture of prompters. Furthermore, our approach not only improves TabPFN results on large datasets, but achieves state-of-the-art performance across tabular benchmarks [McElfresh et al., ].

### 8.4 Mixture of Experts

Sparsely gated MoE [Shazeer et al., , Lewis et al., ] shows a significant improvement in model capacity, training time, and accuracy with a gating mechanism. An expert is a sub-network, which better learns to predict similar data points. A gating mechanism, learnable or non-learnable method, decides to route each data point to the most suited experts [Shazeer et al., ]. Switch transformers [Fedus et al., ] is a learning to route approach, where it assigns one data point to only one expert, instead of top-k, which reduces computation, while preserving accuracy. However, learnable routing methods require auxiliary load balancing loss function, and further tuning. In [Roller et al., ], a non-learnable routing method is proposed, which uses a hashing method to assign similar data points to similar experts. They show that this procedure can be better or be competitive with learnable routing MoE methods. However, a hashing method is not necessarily flexible in assigning data points to suitable experts, as it can cause data skewness and choice of hashing function is sensitive to the downstream task. Inspired by these works, we proposed a non-learnable routing mechanism which assigns one data point to one expert, and our routing method finds the most suitable expert based on similarity of data points with a K-Means clustering method. Also, we used experts with shared weights as opposed to general MoE in most of the previous works, as the prompts, not the model, contains training examples for In-Context Learning. We highlight tackling MoE in the context of prompting is a newly emerging research interest [Anonymous, ], of which we are among the first.

### 8.5 Efficient Transformers

Long sequence inputs have long been studied by the efficient transformer community. Several linear time and space complexity transformers have been proposed [Katharopoulos et al., , Wang et al., , Choromanski et al., , Qin et al., ], primarily be SVD decomposing the attention computation. While efficient transformers can help scale the base PFN model, linear complexity is too expensive for the scale of tabular data in industry. Furthermore, many approaches require finetuning on downstream data [Katharopoulos et al., , Wang et al., ] which is nontrivial for PFN models. Constant time transformers [Zaheer et al., , Bulatov et al., , Chowdhery et al., ] exploit the sequential nature of text data. These methods also do not apply to PFNs for tabular data, as tabular training data is not inherently sequential. Hence, technologies outside of efficient transformers are needed to effectively scale PFNs for tabular learning.

### 8.6 In-Context Learning

In-Context Learning (ICL) prompts transformers with training examples prepended to the desired query [Brown et al., ]. Several works attempt to prompt engineer scalable in-context examples for better downstream performance [Hao et al., ], among which K-Nearest Neighbors emerge as a reliable choice [Liu et al., , Xu et al., ]. However, ICL for LLMs consider queries one-at-a-time instead of in batch fashion, prompts are encoded as natural language, and in-context examples can come from a large corpus of natural lanaguage data [Brown et al., ]. These properties are not afforded to PFN-style ICL, where inference is directly run on training set tokens. In addition to scaling “prompts”, LLMs can also be finetuned on ICL examples to reach better performance [Thoppilan et al., , Wei et al., ]. However, this is due to LLM pretraining objectives misaligning with the ICL task. Such misalignment is not as obvious in the PFN case, where the transformer is directly trained on the ICL task [Müller et al., ]. Hence, we propose Sparse Mixture of In-Context Prompters to support batching and our novel bootstrapping algorithm to finetune PFN models. In the spirit of multi-modal models [Radford et al., , Xu et al., ], TabPFN [Hollmann et al., ] extends ICL techniques from natural language LLMs to tabular data.

## 9 Broader Impact Statement

This paper presents work whose goal is to advance the field of tabular and in-context learning. Our work reaches a new state-of-the-art tabular classification accuracy, which has broad positive impact for many industries using relational databases and tabular datasets. We hope our impressive results inspire further research into PFNs and ICL for tabular learning. Our work is built on large transformer models, which are known to hallucinate in the natural language domain. While we observe no such behavior on our tabular datasets, we will open source our code, such that practitioners can plug in their own safe transformer models. We feel there are not any other noteworthy negative societal impacts.

## 10 Math Notations

We summarize our math notations below:

- •
$B$: Number of training samples in prompt
- •
$K$: Number of experts
- •
$D$: A generic dataset
- •
$D_{train}$: The training dataset
- •
$X_{test}$: The test samples
- •
$X_{batch}$: A batch of testing data
- •
$C$: Number of classes
- •
$q_{\theta}$: PFN Model
- •
$A_{\psi}^{D_{train}}$: PFN Model Adapters trained on $D_{train}$
- •
$\theta$: PFN Model Parameters
- •
$\psi$: PFN Model Adapter Parameters
- •
$\phi$: PFN hypothesis mechanism
- •
$\mathcal{R}$: Router mapping input test points to 1 of $K$ Prompters
- •
$\mathcal{T}_{k}$: The $k$-th Prompter
- •
$D_{supp}(x)$: The $B$-Nearest Neighbor training samples to a test point, $x$
- •
$D_{cluster}^{(k)}$: The $k$-th K-Means cluster of training samples
- •
$D_{prompt}^{(k)}$: The $k$-th Prompter’s training samples context
- •
$(X_{batch},D_{prompt}^{(k)})$: The $k$-th Prompter’s “prompt”
- •
$NNS(\cdot|\cdot)$: Nearest Neighbor Search Algorithm
- •
$KNN(\cdot|\cdot,\cdot)$: K-Nearest Neighbors Algorithm
- •
$KMeans(\cdot)$: K-Means Algorithm
- •
$Sample(\cdot)$: Random Sampling
- •
$N_{train}$: Full training dataset size
- •
$N_{test}$: Full testing dataset size
- •
$N_{batch}$: Batch size
- •
$\gamma$: Single hyperparameter trading off performance and efficiency.

## 11 Nonzero Overlap Proof

We prove Theorem 1 here.

First we prove $|D_{cluster}^{(k)}|\leq B\implies D_{cluster}^{(k)}\subseteq\textsc{KNN}(x_{
center}^{(k)}|D_{train},B)$:

By KMeans definition,

$d(x_{center}^{(k)},x_{train}^{(i)})<d(x_{center}^{(k)},x_{train}^{(j)})
\implies(x_{train}^{(j)}\in D_{cluster}^{(k)}\implies x_{train}^{(i)}\in D_{
cluster}^{(k)})$,

$\implies\exists\tau_{KMeans}^{(k)}:d(x_{center}^{(k)},x_{train}^{(i)})<\tau_{
KMeans}^{(k)}\implies x_{train}^{(i)}\in D_{cluster}^{(k)}$

By KNN definition,

$\textsc{KNN}(x|D_{train},B)=\{x_{train}^{(i)}:d(x_{center}^{(k)},x_{train}^{(i
)})<\tau_{KNN}(x|D_{train})\}$,

where $|\{x_{train}^{(i)}:d(x_{center}^{(k)},x_{train}^{(i)})<\tau_{KNN}(x|D_{train})
\}|=B$

Given $|D_{cluster}^{(k)}|\leq B$,

$\implies|D_{cluster}^{(k)}|\leq|\textsc{KNN}(x_{center}^{(k)}|D_{train},B)|$

$\implies\textsc{KNN}(x_{center}^{(k)}|D_{train},B)=D_{cluster}^{(k)}\cup\{x_{
train}^{(i)}:\tau_{KMeans}^{(k)}\leq d(x_{center}^{(k)},x_{train}^{(i)})<\tau_
{KNN}(x|D_{train})\}$

$\implies D_{cluster}^{(k)}\subseteq\textsc{KNN}(x_{center}^{(k)}|D_{train},B)$

Next, we prove Theorem 1:

Given $|D_{cluster}^{(k)}|\leq B$ $\forall k\in[0,...,K-1]$ and $\mathcal{R}^{*}(x_{train}^{(i)})=p:x_{train}^{(i)}\in D_{cluster}^{(k)}$ ,

$\implies G(k)=1$

$\implies x_{train}^{(i)}\in\textsc{KNN}(x_{train}^{(i)}|D_{train},B)=D_{supp}(
x_{train}^{(i)})$

$\implies x_{train}^{(i)}\in D_{cluster}^{(\mathcal{R}^{*}(x_{train}^{(i)}))}
\subseteq\textsc{KNN}(x_{center}^{(\mathcal{R}^{*}(x_{train}^{(i)}))}|D_{train
},B)=D_{prompt}(\mathcal{R}^{*}(x_{train}^{(i)}))$

$\implies x_{train}^{(i)}\in D_{prompt}(\mathcal{R}^{*}(x_{train}^{(i)}))\cap D
_{supp}(x_{train}^{(i)})$

$\implies|D_{prompt}(\mathcal{R}^{*}(x_{train}^{(i)}))\cap D_{supp}(x_{train}^{
(i)})|\geq|\{x_{train}^{(i)}\}|=1$ $\forall i\in[0,...,N_{train}-1]$ $\forall D_{train}$.

## 12 Support Set and KNN Prompting.

We provide illustrations for the support set as described in Section 3.1 in Figure 13. Note, support sets, as defined in Section 3.1, are very similar to the KNN-Prompting idea from In-Context Learning (ICL) for Large Language Models (LLMs). We point out one key difference between KNN-Prompting with TabPFN [Hollmann et al., ] and with ICL on LLMs [Liu et al., ]: LLMs support batching multiple test pairs across multiple “prompts” where each “prompt” contains the nearest neighbors to 1 test point. In contrast, TabPFN requires multiple test cases in 1 “prompt”, which necessitates techniques like MICP to route a batch of test points to the same “prompt”. As shown in Section 5.2, TabPFN’s more efficient prompts achieve better performance under same GPU constraints as ICL for LLM batched prompts. We leave application of MICP to ICL and LLMs as future work.

## 13 Time and Space Complexity Details

Router and prompter initialization takes $\mathcal{O}(tN_{train}K+(N_{train}+KB)logN_{train})$ time and $\mathcal{O}(N_{train}+KB)$ space complexity and is done once before inference. For initialization, K-Means with $t$-iterations takes $\mathcal{O}(tN_{train}K)$ time and $\mathcal{O}(K)$ space. To perform efficient nearest neighbor queries, we use the ball-tree algorithm over the training dataset and cluster centers , which takes $\mathcal{O}(N_{train}log(N_{train}))$ time and $\mathcal{O}(N_{train})$ space. Using ball-tree KNN queries, constructing each ICP support set takes $\mathcal{O}(Blog(N_{train}))$ time and $\mathcal{O}(B)$ space.

Routing takes $\mathcal{O}(log(K))$ time and $\mathcal{O}(1)$ space complexity, using efficient nearest neighbor search with ball-tree for each test sample. Router overhead is practically overcome via highly-optimized NNS implementations [Douze et al., ], which scale $K$ to the billions.

PFN transformer inference takes $\mathcal{O}(B^{2}+BN_{batch})$ time and space complexity, as MICP prompts contain at most $B$ training samples and $N_{batch}=|X_{batch}|$ testing samples. Note, unlike purely KNN-based prompting, MICP supports batched computation to further amortize PFN inference cost.

**Table 3: Datasets and algorithms considered in each Top-K subset. For fair evaluation [abvesa, ], we only consider the shared set of datasets all algorithms run on in the Top-K subsets. By considering different subsets, we evaluate MixturePFN against more or less algorithms and datasets. 20 datasets are shared by all 20 algorithms. 22 datasets are shared by Top-10 algorithms. 25 datasets are shared by Top-5 algorithms. MixturePFN achieves the best mean rank among all Top-K subsets.**
| Subset | Considered Algorithms | Considered Datasets |
| --- | --- | --- |
| ALL | MixturePFN, CatBoost, | ada-agnostic, australian, balance-scale, |
|  | TabPFN*, XGBoost, | colic, credit-approval, elevators, heart-h, |
|  | ResNet, FTTransformer, | jasmine, kc1, lymph, mfeat-fourier, mfeat-zernike, |
|  | LightGBM, SAINT, | monks-problems-2, phoneme, profb, |
|  | NODE, MLP-rtdl, | qsar-biodeg, socmob, speeddating, splice, vehicle |
|  | RandomForest, TabNet, |  |
|  | MLP, DecisionTree, |  |
|  | LinearModel, STG, |  |
|  | VIME, KNN, DANet, |  |
|  | SVM |  |
| Top-10 | MixturePFN, CatBoost, | ada-agnostic, artificial-characters, australian, |
|  | TabPFN*, XGBoost, | balance-scale, colic, credit-approval, elevators, |
|  | ResNet, LightGBM, | gesturephasesegmentationprocessed, heart-h, jasmine, |
|  | SAINT, DANet, | kc1, lymph, mfeat-fourier, mfeat-zernike, |
|  | FTTransformer, SVM, | monks-problems-2, phoneme, profb, qsar-biodeg, |
|  |  | socmob, speeddating, splice, vehicle |
| Top-5 | MixturePFN, CatBoost, | ada-agnostic, artificial-characters, australian, |
|  | TabPFN*, XGBoost, | balance-scale, colic, credit-approval, electricity, elevators, |
|  | SAINT | albert, gesturephasesegmentationprocessed, heart-h, higgs, |
|  |  | jasmine, jungle-chess-2pcs-raw-endgame-complete, kc1, |
|  |  | lymph, mfeat-fourier, mfeat-zernike, monks-problems-2, |
|  |  | phoneme, profb, qsar-biodeg, socmob, speeddating, |
|  |  | splice, vehicle |

**Table 4: MixturePFN is the Condorcet winner across 36 datasets against 19 baseline algorithms. We rank algorithms based on their accuracies.**
| Method | Condorcet Statistics |  |  |  |
| --- | --- | --- | --- | --- |
| #Votes$\uparrow$ | #Wins$\uparrow$ | #Ties | #Losses$\downarrow$ |  |
| MixturePFN | 464 | 19 | 0 | 0 |
| CatBoost | 462 | 18 | 0 | 1 |
| XGBoost | 448 | 16 | 1 | 2 |
| SAINT | 402 | 15 | 0 | 4 |
| ResNet | 367 | 14 | 2 | 3 |
| LightGBM | 360 | 12 | 0 | 7 |
| FTTransformer | 345 | 13 | 1 | 5 |
| TabPFN* | 330 | 11 | 0 | 8 |
| NODE | 304 | 11 | 0 | 8 |
| DANet | 300 | 10 | 1 | 8 |
| RandomForest | 299 | 12 | 1 | 6 |
| MLP-rtdl | 256 | 8 | 0 | 11 |
| SVM | 236 | 6 | 0 | 13 |
| TabNet | 231 | 4 | 0 | 15 |
| MLP | 229 | 7 | 0 | 12 |
| STG | 188 | 5 | 0 | 14 |
| DecisionTree | 155 | 2 | 0 | 17 |
| LinearModel | 144 | 3 | 0 | 16 |
| KNN | 112 | 0 | 0 | 19 |
| VIME | 93 | 1 | 0 | 18 |

**Table 5: MixturePFN achieves the top mean rank w.r.t. Log Likelihood across 20 datasets where all algorithms successfully run.**
| Method | All Algorithms (Log Likelihood) |  |  |  |
| --- | --- | --- | --- | --- |
| Mean $\pm$ Std | Median | Min | Max |  |
| Rank$\downarrow$ | Rank$\downarrow$ | Rank$\downarrow$ | Rank$\downarrow$ |  |
| MixturePFN | 2.350 $\pm$ 1.824 | 1.0 | 1.0 | 6.0 |
| TabPFN* | 4.550 $\pm$ 2.747 | 4.5 | 2.0 | 13.0 |
| CatBoost | 4.900 $\pm$ 4.158 | 3.0 | 1.0 | 14.0 |
| XGBoost | 5.500 $\pm$ 4.621 | 4.0 | 1.0 | 16.0 |
| SAINT | 8.300 $\pm$ 5.367 | 7.0 | 1.0 | 19.0 |
| ResNet | 8.400 $\pm$ 3.262 | 8.5 | 3.0 | 15.0 |
| FTTransformer | 8.600 $\pm$ 3.541 | 8.5 | 3.0 | 15.0 |
| DANet | 9.050 $\pm$ 3.369 | 8.5 | 3.0 | 17.0 |
| LightGBM | 9.150 $\pm$ 4.351 | 10.5 | 2.0 | 16.0 |
| MLP-rtdl | 10.800 $\pm$ 5.046 | 10.5 | 2.0 | 20.0 |
| SVM | 11.300 $\pm$ 4.766 | 11.0 | 2.0 | 19.0 |
| RandomForest | 11.600 $\pm$ 4.443 | 12.5 | 4.0 | 18.0 |
| STG | 11.900 $\pm$ 4.549 | 12.5 | 4.0 | 19.0 |
| LinearModel | 12.400 $\pm$ 4.652 | 12.5 | 5.0 | 20.0 |
| NODE | 13.350 $\pm$ 3.410 | 13.5 | 7.0 | 18.0 |
| TabNet | 13.550 $\pm$ 5.296 | 14.0 | 2.0 | 20.0 |
| MLP | 13.700 $\pm$ 3.621 | 14.0 | 4.0 | 19.0 |
| VIME | 15.350 $\pm$ 3.851 | 17.0 | 6.0 | 19.0 |
| DecisionTree | 16.800 $\pm$ 3.881 | 18.5 | 7.0 | 20.0 |
| KNN | 18.450 $\pm$ 1.936 | 19.0 | 14.0 | 20.0 |

**Table 6: MixturePFN achieves the top mean rank w.r.t. Accuracy across 20 datasets where all algorithms successfully run.**
| Method | All Algorithms (Accuracy) |  |  |  |
| --- | --- | --- | --- | --- |
| Mean $\pm$ Std | Median | Min | Max |  |
| Rank$\downarrow$ | Rank$\downarrow$ | Rank$\downarrow$ | Rank$\downarrow$ |  |
| MixturePFN | 3.950 $\pm$ 3.570 | 3.0 | 1.0 | 13.0 |
| TabPFN* | 5.500 $\pm$ 4.056 | 5.5 | 1.0 | 15.0 |
| CatBoost | 6.350 $\pm$ 5.313 | 4.0 | 1.0 | 17.0 |
| XGBoost | 7.100 $\pm$ 5.234 | 5.5 | 1.0 | 18.0 |
| ResNet | 7.600 $\pm$ 4.017 | 6.5 | 1.0 | 16.0 |
| FTTransformer | 7.900 $\pm$ 4.158 | 7.5 | 1.0 | 17.0 |
| SAINT | 7.950 $\pm$ 5.723 | 5.5 | 1.0 | 20.0 |
| LightGBM | 8.750 $\pm$ 4.918 | 8.5 | 1.0 | 17.0 |
| NODE | 8.850 $\pm$ 3.395 | 9.0 | 3.0 | 15.0 |
| MLP-rtdl | 9.150 $\pm$ 5.033 | 9.0 | 1.0 | 18.0 |
| RandomForest | 9.350 $\pm$ 4.757 | 8.5 | 4.0 | 19.0 |
| DANet | 10.000 $\pm$ 4.405 | 11.0 | 3.0 | 19.0 |
| SVM | 12.250 $\pm$ 5.476 | 14.5 | 1.0 | 19.0 |
| TabNet | 13.050 $\pm$ 4.780 | 13.5 | 2.0 | 20.0 |
| MLP | 13.100 $\pm$ 4.253 | 14.5 | 5.0 | 18.0 |
| LinearModel | 14.100 $\pm$ 3.846 | 14.0 | 8.0 | 20.0 |
| DecisionTree | 14.250 $\pm$ 4.426 | 15.0 | 3.0 | 20.0 |
| STG | 14.800 $\pm$ 4.423 | 16.0 | 4.0 | 20.0 |
| VIME | 16.950 $\pm$ 2.854 | 18.0 | 10.0 | 20.0 |
| KNN | 17.400 $\pm$ 3.470 | 19.0 | 7.0 | 20.0 |

**Table 7: MixturePFN achieves the top mean rank w.r.t. Log Likelihood across 22 datasets where all Top-10 algorithms successfully run.**
| Method | Top-10 Algorithms (Log-Likelihood) |  |  |  |
| --- | --- | --- | --- | --- |
| Mean $\pm$ Std | Median | Min | Max |  |
| Rank$\downarrow$ | Rank$\downarrow$ | Rank$\downarrow$ | Rank$\downarrow$ |  |
| MixturePFN | 2.273 $\pm$ 1.710 | 1.0 | 1.0 | 6.0 |
| CatBoost | 3.955 $\pm$ 2.688 | 3.0 | 1.0 | 10.0 |
| XGBoost | 4.000 $\pm$ 2.663 | 3.0 | 1.0 | 10.0 |
| TabPFN* | 4.045 $\pm$ 1.965 | 4.0 | 2.0 | 9.0 |
| SAINT | 6.136 $\pm$ 2.473 | 6.5 | 1.0 | 10.0 |
| LightGBM | 6.409 $\pm$ 2.839 | 7.5 | 2.0 | 10.0 |
| FTTransformer | 6.773 $\pm$ 2.235 | 7.0 | 3.0 | 10.0 |
| ResNet | 6.864 $\pm$ 1.961 | 7.0 | 3.0 | 9.0 |
| DANet | 7.045 $\pm$ 1.988 | 7.0 | 3.0 | 10.0 |
| SVM | 7.500 $\pm$ 2.482 | 8.0 | 2.0 | 10.0 |

**Table 8: MixturePFN achieves the top mean rank w.r.t. Accuracy across 22 datasets where all Top-10 algorithms successfully run.**
| Method | Top-10 Algorithms (Accuracy) |  |  |  |
| --- | --- | --- | --- | --- |
| Mean $\pm$ Std | Median | Min | Max |  |
| Rank$\downarrow$ | Rank$\downarrow$ | Rank$\downarrow$ | Rank$\downarrow$ |  |
| MixturePFN | 3.000 $\pm$ 2.256 | 2.0 | 1.0 | 9.0 |
| TabPFN* | 4.318 $\pm$ 2.703 | 4.5 | 1.0 | 9.0 |
| CatBoost | 4.591 $\pm$ 2.964 | 4.0 | 1.0 | 10.0 |
| XGBoost | 4.773 $\pm$ 2.859 | 4.0 | 1.0 | 10.0 |
| ResNet | 5.591 $\pm$ 2.103 | 5.5 | 1.0 | 9.0 |
| LightGBM | 5.727 $\pm$ 2.847 | 6.5 | 1.0 | 10.0 |
| SAINT | 5.727 $\pm$ 2.847 | 5.0 | 1.0 | 10.0 |
| FTTransformer | 5.864 $\pm$ 2.282 | 6.0 | 1.0 | 9.0 |
| DANet | 6.955 $\pm$ 2.184 | 7.5 | 3.0 | 10.0 |
| SVM | 7.773 $\pm$ 2.907 | 9.0 | 1.0 | 10.0 |

**Table 9: MixturePFN achieves the top mean rank w.r.t. Log Likelihood across 25 datasets where all Top-5 algorithms successfully run.**
| Method | Top-5 Algorithms (Log Likelihood) |  |  |  |
| --- | --- | --- | --- | --- |
| Mean $\pm$ Std | Median | Min | Max |  |
| Rank$\downarrow$ | Rank$\downarrow$ | Rank$\downarrow$ | Rank$\downarrow$ |  |
| MixturePFN | 2.000 $\pm$ 1.166 | 1.0 | 1.0 | 4.0 |
| XGBoost | 2.760 $\pm$ 1.394 | 3.0 | 1.0 | 5.0 |
| CatBoost | 2.880 $\pm$ 1.243 | 3.0 | 1.0 | 5.0 |
| TabPFN* | 3.320 $\pm$ 1.085 | 3.0 | 2.0 | 5.0 |
| SAINT | 4.040 $\pm$ 1.311 | 5.0 | 1.0 | 5.0 |

**Table 10: MixturePFN achieves the top mean rank w.r.t. Accuracy across 25 datasets where all Top-5 algorithms successfully run.**
| Method | Top-5 Algorithms (Accuracy) |  |  |  |
| --- | --- | --- | --- | --- |
| Mean $\pm$ Std | Median | Min | Max |  |
| Rank$\downarrow$ | Rank$\downarrow$ | Rank$\downarrow$ | Rank$\downarrow$ |  |
| MixturePFN | 2.360 $\pm$ 1.292 | 2.0 | 1.0 | 5.0 |
| XGBoost | 2.880 $\pm$ 1.395 | 3.0 | 1.0 | 5.0 |
| CatBoost | 2.920 $\pm$ 1.354 | 3.0 | 1.0 | 5.0 |
| TabPFN* | 3.000 $\pm$ 1.386 | 3.0 | 1.0 | 5.0 |
| SAINT | 3.680 $\pm$ 1.406 | 4.0 | 1.0 | 5.0 |

Figure: Figure 5: Pairwise comparison matrix for Condorcet voting over the log likelihood metric. Note, MixturePFN is the Condorcet winner.
Refer to caption: extracted/5619985/figures_new/condorcet_nll.png

Figure: Figure 6: Pairwise comparison matrix for Condorcet voting over the accuracy metric. Note, MixturePFN is the Condorcet winner. Please refer to Section 15 for more discussion.
Refer to caption: extracted/5619985/figures_new/condorcet_acc.png

Figure: (a) $N_{train}$ w.r.t. TabPFN*
Refer to caption: extracted/5619985/figures_new/improv_acc.png

Figure: (a) $N_{train}$ w.r.t. baselines
Refer to caption: extracted/5619985/figures_new/xgb_samp_nll.png

Figure: (a) $N_{train}$ w.r.t. baselines
Refer to caption: extracted/5619985/figures_new/xgb_samp_acc.png

Figure: Figure 10: Wilcoxon-Signed Rank Test shows MixturePFN significantly outperforms the Top-10 baselines on the 22 shared datasets, under the accuracy metric. We compute the rank across all 10 cross-validation splits.
Refer to caption: extracted/5619985/figures_new/CD_top10_acc.png

Figure: Figure 11: Wilcoxon-Signed Rank Test shows MixturePFN significantly outperforms the all baselines on the 20 shared datasets, under the log likelihood metric. We compute the rank across all 10 cross-validation splits.
Refer to caption: extracted/5619985/figures_new/CD_all_nll.png

Figure: Figure 12: Wilcoxon-Signed Rank Test shows MixturePFN significantly outperforms the all baselines on the 20 shared datasets, under the accuracy metric. We compute the rank across all 10 cross-validation splits.
Refer to caption: extracted/5619985/figures_new/CD_all_acc.png

## 14 Choosing Baseline Datasets and Algorithms

We chose our datasets from the TabZilla benchmark, which curates 36 of the hardest 176 considered datasets across 19 algorithms. As noted in Section 4.0.1, not all algorithms run on all datasets. We note which datasets are shared by all algorithms, Top-10 algorithms, and Top-5 algorithms in Table 3. We provide the number of datasets each algorithm successfully runs on in Table 11. We provide dataset names and statistics in Table 14.

The 2 datasets MixturePFN and TabPFN* fails on contain >10 classes, which is not currently supported by the pretrained TabPFN. However, as seen in Figure 15 and described in Section 14.1, we highlight that the 34 datasets MixturePFN and TabPFN* successfully run on cover the full range of dataset properties: number of features, number of samples, and dataset irregularity. Specifically, our 34 datasets include the ones with the least and most number of samples, the least and most number of features, and the least and most irregularities. Because the focus of this work is to scale TabPFN to datasets with more number of samples, we leave extending TabPFN to more number of classes as future work.

### 14.1 Baseline Datasets

We provide dataset statistics in Table 14. As shown, our considered datasets cover a wide range of dataset properties in number of features, number of samples, and std. kurtosis. As shown, MixturePFN achieves the best average performance across all datasets. Furthermore, we plot the full range of dataset properties covered, along with the 2 held out datasets from TabZilla in Figure 15, showing that the datasets successfully run on are representative of the benchmark as a whole. Refer to Section 14 for more details.

We followed the same experimental setup as TabZilla [McElfresh et al., ], which includes: imputing NaN features to its non-NaN mean and all other preprocessing is handled by each respective baseline. MixturePFN and TabPFN* follow TabZilla’s PFN preprocessing [McElfresh et al., ]: categorical features are encoded as ordinal features, outliers are dropped, features are normalized, results are ensembled across shuffling the feature ordering, and results are ensembled across power-law scaled and unscaled features.

### 14.2 Baseline Algorithms

#### 14.2.1 Prior-Fitted Network Models ( PFN )

TabPFN* is the only PFN-based baseline, which uses a pretrained 12-layer TabPFN transformer model, with embeddings size 512, hidden size 1024 in feed-forward layers, and 4-headed attention. TabPFN is pretrained on a handcrafted dataset prior consisting of randomly generated structural causal models [Hollmann et al., ]. During inference features and labels are randomly shuffled in batch size 32 then ensembled together, following the TabZilla benchmark [McElfresh et al., ]. Our work improves TabPFN’s scalability to different dataset properties, particular in number of training samples.

#### 14.2.2 Gradient-boosted Decision Tree Models ( GBDT )

CatBoost [Prokhorenkova et al., ], XGBoost [Chen and Guestrin, ], and LightGBM [Ke et al., ] are GBDT models. These models utilize boosting to construct an ensemble of small trees for evaluation. GBDTs are robust to uninformative or heavy tail features and achieve competitive performance over baselines across different dataset properties. Our work argues in-context learning is a potential competitor against GBDTs, as PFN transformers can potentially learn a better dataset prior than GBDTs.

#### 14.2.3 Deep Learning Algorithms

ResNet [Gorishniy et al., ], MLP-rtdl, TabNet [Arik and Pfister, ], MLP, STG [Yamada et al., ], VIME [Yoon et al., ], NODE [Popov et al., ], FTTransformer [Gorishniy et al., ], SVM [Cortes and Vapnik, ], DANet [Chen et al., ] and SAINT [Somepalli et al., ] are deep learning-based algorithms. In particular, ResNet is a Convolutional Neural Network designed for tabular learning. MLP-rtdl and MLP are 2 implementations of multilayer-perceptrons. SVM is a support vector machine. TabNet and STG is a neural network architecture that aims to learn GBDT-like operations in a fully differentiable manner. VIME is a gated neural network that is first traing with self supervision. MixturePFN outperforms deep learning algorithms by learning a prior that better regularizes the learning procedure. NODE is a neural network architecture that aims to imitate GBDTs while being fully differentiable and end2end. DANet is a specialized deep learning architecture for tabular data. FTTransformer is a feature encoding transformer model designed for tabular data. SAINT is a self-supervised transformer designed for tabular data.

#### 14.2.4 Simple Algorithms

RandomForest, DecisionTree, LinearModel, and KNN are all standard machine learning algorithms. We highlight, although MixturePFN is based on evaluating prompts only on KNN neighborhoods, it drastically outperforms KNN. This suggests that In-Context Learning-based models trained on a local neighborhood can outperform both complicated models trained on the entire dataset and simple models that return the average of local neighborhoods. Indeed combining KNN and Large Language Models have been highly successful [Xu et al., , Liu et al., , Guu et al., ].

**Table 11: The number of datasets each algorithm completed on across the entire 36 dataset TabZilla benchmark. Note, the 2 datasets that MixturePFN and TabPFN* [McElfresh et al., ] does not run on has too many labels, being unsupported by the pretrained TabPFN [Hollmann et al., ]. However these 2 datasets are not outliers compared to the 34 datasets that are supported. Note, TabPFN achieves the same results as TabPFN*, except only running on the 17 datasets with <3,000 features, hence we compare against the more powerful TabPFN* baseline instead.**
| Method | Number of Datasets Completed On |
| --- | --- |
| CatBoost [Prokhorenkova et al., ] | 35 |
| XGBoost [Chen and Guestrin, ] | 36 |
| MLP-rtdl [Goodfellow et al., , Gorishniy et al., ] | 36 |
| MLP [Goodfellow et al., , McElfresh et al., ] | 36 |
| ResNet [Gorishniy et al., ] | 35 |
| RandomForest [Liaw et al., ] | 35 |
| DecisionTree [Quinlan, ] | 35 |
| MixturePFN | 34 |
| TabPFN* [Hollmann et al., , McElfresh et al., ] | 34 |
| LinearModel [Cox, ] | 34 |
| TabNet [Arik and Pfister, ] | 33 |
| KNN [Cover and Hart, ] | 33 |
| LightGBM [Ke et al., ] | 32 |
| VIME [Yoon et al., ] | 32 |
| STG [Yamada et al., ] | 31 |
| NODE [Popov et al., ] | 30 |
| FTTransformer [Gorishniy et al., ] | 29 |
| SVM [Cortes and Vapnik, ] | 29 |
| SAINT [Somepalli et al., ] | 27 |
| DANet [Chen et al., ] | 27 |

Figure: Figure 13: We hypothesize only a subset of the training data, $D_{support}(x_{test}^{(i)})$, is required for effective in-context learning on test sample, $x_{test}^{(i)}$, and this subset is the $B$ nearest training samples in feature space: $D_{support}(x_{test}^{(i)})=\textsc{KNN}(x_{test}^{(i)}|D_{train},B)$.
Refer to caption: extracted/5619985/figures/approx3.png

## 15 Additional Results

### 15.1 Accuracy Results and Standard Deviations

We provide the same Condorcet experiments as the main paper but with the accuracy metric in Table 4. We provide the Condorcet matrix in Figures 5 and 6. These results support the same conclusions found in the main paper.

We provide the detailed statistics of the ranking experiments with both log-likelihood and accuracy across all, Top-10, and Top-5 subsets in Tables 5, 6, 7, 8, 9, 10. These results support the same conclusions found in the main paper.

We provide the Wilcoxon-Signed Rank Test with both the log-likelihood and accuracy metric across the all and Top-10 subsets in Figures 10,  11, and 12. These results support the same conclusions found in the main paper.

We provide experiments with even lighter hyperparameter tuning, as discussed in Section 17, we call this model, MixturePFN-lite. MixturePFN is tuned over >80% less configurations than baselines. MixturePFN-lite is tuned only on 2 hyparparameter settings. We provide the main paper results and Condorcet matrices as presented in Table 13 and Figure 16.

MixturePFN is the Condorcet winner and achieves the top mean rank across all experimental settings, with statistical significance among all and Top-10 subsets. MixturePFN’s Log-likelihood results are slightly better, because many algorithms are tied in accuracy across the benchmark. When this occurs, MixturePFN is more confident than baselines when it is correct.

### 15.2 Sensitivity Results on More Data

We provide the same sensitivity analysis conducted in the main paper but with the accuracy metric in Figure 7. These figures support the same conclusions found in the main paper.

We provide the same sensitivity analysis conducted in the main paper but across the number of features, number of labels, and feature irregularity in Figures 8 and 9. These figures support the same conclusions found in the main paper.

**Table 12: Trade-off of $\gamma$. $T*$+MICP is short for TabPFN* +MICP. Note as $\gamma$ increases, the accuracy and number of prompters increases, while TabPFN inference time remains constant. Routing costs are negligible with optimized nearest neighbor search [Douze et al., ].**
| Method | Accuracy | K | $\frac{\textbf{Time}}{\textbf{Prompt}}$ |
| --- | --- | --- | --- |
| Mean | Mean | Mean |  |
| TabPFN* | 83.42% | 1.00 | 1.24s |
| T*+MICP ($\gamma=1.0$) | 83.96% | 2.25 | 0.90s |
| T*+MICP ($\gamma=3.0$) | 84.23% | 5.25 | 1.02s |
| T*+MICP ($\gamma=5.0$) | 84.23% | 8.54 | 0.83s |

### 15.3 Efficiency Effectiveness Trade-Off

$\gamma$ effectively trade-offs efficiency and effectiveness. As mentioned in Section 3.2.2, MixturePFN uses a single hyperparameter, $\gamma$, to control the efficiency effectiveness tradeoff. we plot the average accuracy of TabPFN +MICP across $\gamma=[1.0,3.0,5.0]$, across the entire dataset. As shown in Table 12, as the hyperparameter $\gamma$ increases, MICP’s effectiveness is reliably trade-off for efficiency.

### 15.4 Timing Analysis

To study runtime, we subsample the electricity dataset into datasets of smaller sizes and then run TabPFN and MixturePFN. Note, TabPFN’s primary bottleneck is its $\mathcal{O}(N_{train}^{2})$ memory bottleneck. While this can be overcome via subsampling, as in the case of TabPFN*, we study the performance compromises of said approach in Sections 5.1 and  5.2. As seen in Figure 14, MixturePFN is scalable in both runtime and memory costs compared to TabPFN.

### 15.5 Detailed Results

We provide MixturePFN’s and TabPFN*’s accuracies across the 10-folds on all datasets in Tables 15 and 16.

Figure: Figure 14: When $N_{train}>3,000$, TabPFN runs out of GPU memory. MixturePFN scales much better in runtime than TabPFN by using MICP to construct bounded size “prompts”. Inference is slower than XGBoost, due to TabPFN’s transformer’s forward pass latency.
Refer to caption: extracted/5619985/figures/_preprocessing.png

Figure: Figure 15: Dataset properties of chosen algorithms from the TabZilla benchmark. We plot 3 dimensions of the dataset properties of all 36 dataset from TabZilla. The 2 bold points represent the held-out datasets. As shown, the 34 chosen datasets covers a wide-variety of dataset properties.
Refer to caption: extracted/5619985/figures/_datasets.png

## 16 Implementation

We implemented MICP by first preprocessing the train data into separate prompts via KNN, chunking each prompt into batches, then called TabZilla APIs to run the desired PFN model on each batch. We implemented CaPFN by bootstrapping our training dataset then running maximum likelihood loss on the bootstrapped datasets. MixturePFN’s implementation is built off the official TabZilla codebase [naszilla, ].

### 16.1 Optimizing TabPFN ’s Implementation

TabZilla only obtained TabPFN* results on 7 out of 34 benchmark datasets [McElfresh et al., ], due to memory constraints. We identified an implementation inefficiency where “prompts” are constructed with the entire test dataset, i.e. $(X_{test},D_{train})$, causing memory overflow. We optimized TabPFN*’s implementation by batching test samples, $(X_{batch}|D_{train}),X_{batch}\subseteq X_{test}$, with batch size 1024, and report results over all 26 datasets.

**Table 13: MixturePFN is the Condorcet winner across 36 datasets against 19 baseline algorithms. We rank algorithms based on their log-likelihoods.**
| Method | Condorcet Statistics |  |  |  |
| --- | --- | --- | --- | --- |
| #Votes$\uparrow$ | #Wins$\uparrow$ | #Ties | #Losses$\downarrow$ |  |
| MixturePFN-lite | 503 | 19 | 0 | 0 |
| XGBoost | 502 | 18 | 0 | 1 |
| CatBoost | 479 | 17 | 0 | 2 |
| SAINT | 404 | 16 | 0 | 3 |
| TabPFN* | 385 | 13 | 1 | 5 |
| LightGBM | 374 | 14 | 1 | 4 |
| DANet | 312 | 14 | 0 | 5 |
| FTTransformer | 295 | 12 | 0 | 7 |
| ResNet | 287 | 10 | 0 | 9 |
| SVM | 286 | 11 | 0 | 8 |
| STG | 286 | 9 | 0 | 10 |
| RandomForest | 248 | 7 | 0 | 12 |
| NODE | 244 | 7 | 0 | 12 |
| MLP-rtdl | 228 | 5 | 0 | 14 |
| TabNet | 210 | 5 | 0 | 14 |
| LinearModel | 202 | 3 | 1 | 15 |
| MLP | 193 | 5 | 1 | 13 |
| VIME | 134 | 2 | 0 | 17 |
| DecisionTree | 115 | 1 | 0 | 18 |
| KNN | 74 | 0 | 0 | 19 |

Figure: Figure 16: Pairwise comparison matrix for Condorcet voting over the log likelihood metric with lightly tuned MixturePFN. Note, MixturePFN-lite is the Condorcet winner.
Refer to caption: extracted/5619985/figures_new/condorcet_nll_lite.png

## 17 Hyperparameter Setup

As TabPFN transformers can handle up to 3,000 training samples, we set $B=3,000$. We empirically found the minimum number of iterations and batch-size required for loss convergence on the artificial-characters dataset to be 128 iterations and $N_{batch}=64$, which we set for all other datasets. During inference, we use a larger batch size, $N_{batch}=1024$, as gradients no longer need to be stored. We finetuned the model using the Adam optimizer with a learning rate of 0.001. As TabPFN transformers can handle up to 100 features, for datasets with over 100 features and TabPFN-based models, we use Maximum Relevance and Minimum Redundancy (mRMR) feature selection [Ding and Peng, ] to reduce the number of features to 100. We follow the TabZilla benchmark, setting $N_{ensemble}=16$, which shuffles features $N_{ensemble}/2$ times for both the original and applies power-law scaled features. MixturePFN’s router was implemented using the FAISS [Douze et al., ] library.

Due to the large variability in datasets in the TabZilla benchmark, we try 4 hyperparameter settings: (1) $\gamma=5.0$, (2) $\gamma=1.0$, (3) $\gamma=1.0$ but MRMR with 50 features instead of 100 features for feature count scalability, and (4) $\gamma=1.0$ but with Catboost instead of Ordinal encoding for categorical feature scalability [Hollmann et al., ]. Hyperparameters are chosen by picking the setting which maximizes performance on the validation set. Models are evaluated on the test set, which is not seen during hyperparameter tuning. In contrast to all other baselines, which are tuned across 30 hyperparameter settings, MixturePFN performs much less hyperparmeter tuning than baselines. Baseline hyperparameter settings are the same as the TabZilla [McElfresh et al., ] benchmark. Note, even if only the $\gamma$ parameter tuned (i.e. only settings (1) and (2)), MixturePFN is still much better than TabPFN, as presented in Table 13 and Figure 16. All results were collected over 10-folds following TabZilla [McElfresh et al., ] and OpenML. We tune the hyperparameters by splitting the train set of each fold into training and validation following TabZilla [McElfresh et al., ]. Ablation studies were performed modifying hyperparameter setting (1). Dataset preprocessing details can be found in Appendix 14.1.

## 18 Hardware

All experiments were conducted on an Nvidia V100 GPU and an AMD EPYC 7402 CPU. Each experiment is given a budget of 10 hours for a single dataset and algorithm.

**Table 14: Dataset statistics for valid TabZilla benchmark datasets. Ranks are computed across algorithms that run on all datasets MixturePFN runs on: MixturePFN, TabPFN*, XGBoost, MLP, and MLP-rtdl. Note this list of datasets was originally curated from 197 datasets, to contain only those difficult for all models. We list the top 2 performing algorithms based on log likelihood, following TabZilla, on each dataset. MixturePFN achieves state-of-the-art performance.**
| Dataset | Dataset Properties | Top 2 Algs. |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| #Samples | #Feats | #Lab. | Std. Kurt. | 1st | 2nd |  |
| lymph | 148 | 18 | 4 | 17.04 | XGBoost | TabPFN* |
| audiology | 226 | 69 | 24 | None | XGBoost | - |
| heart-h | 294 | 13 | 1 | None | MLP-rtdl | MixturePFN |
| colic | 368 | 26 | 1 | 4.0 | XGBoost | MixturePFN |
| monks-prob… | 601 | 6 | 1 | None | MixturePFN | MLP-rtdl |
| balance-scale | 625 | 4 | 3 | 0.02 | MixturePFN | TabPFN* |
| profb | 672 | 9 | 1 | 0.95 | MixturePFN | MLP-rtdl |
| Australian | 690 | 14 | 1 | 2.0 | XGBoost | TabPFN* |
| credit-approval | 690 | 15 | 1 | 74.77 | TabPFN* | MixturePFN |
| vehicle | 846 | 18 | 4 | 15.16 | MixturePFN | TabPFN* |
| credit-g | 1000 | 20 | 1 | 1.92 | MixturePFN | TabPFN* |
| qsar-biodeg | 1055 | 41 | 1 | 93.24 | MixturePFN | TabPFN* |
| cnae-9 | 1080 | 856 | 9 | None | MLP-rtdl | MLP |
| socmob | 1156 | 5 | 1 | None | XGBoost | TabPFN* |
| 100plants | 1599 | 64 | 100 | 17.66 | XGBoost | - |
| mfeat-fourier | 2000 | 76 | 10 | 0.64 | MixturePFN | TabPFN* |
| mfeat-zernike | 2000 | 47 | 10 | 1.42 | MixturePFN | TabPFN* |
| kc1 | 2109 | 21 | 1 | 28.34 | MixturePFN | TabPFN* |
| jasmine | 2984 | 144 | 1 | 47.6 | MixturePFN | XGBoost |
| splice | 3190 | 60 | 3 | None | MixturePFN | XGBoost |
| Bioresponse | 3751 | 1776 | 1 | 328.77 | XGBoost | MixturePFN |
| ada-agnostic | 4562 | 48 | 1 | None | XGBoost | MixturePFN |
| phoneme | 5404 | 5 | 1 | 1.23 | MixturePFN | XGBoost |
| SpeedDating | 8378 | 120 | 1 | 36.43 | MixturePFN | XGBoost |
| GesturePhase… | 9873 | 32 | 5 | 52.18 | MixturePFN | XGBoost |
| artificial-char… | 10218 | 7 | 10 | 0.63 | XGBoost | MixturePFN |
| elevators | 16599 | 18 | 1 | 2986.5 | MixturePFN | TabPFN* |
| guillermo | 20000 | 4296 | 1 | None | XGBoost | TabPFN* |
| nomao | 34465 | 118 | 1 | 1100.34 | XGBoost | MixturePFN |
| jungle-chess… | 44819 | 6 | 3 | 0.08 | MixturePFN | XGBoost |
| electricity | 45312 | 8 | 1 | 2693.51 | XGBoost | MixturePFN |
| higgs | 98050 | 28 | 1 | 15.53 | XGBoost | MLP |
| MiniBooNE | 130064 | 50 | 1 | 1686.9 | MixturePFN | XGBoost |
| albert | 425240 | 78 | 1 | 12162.65 | MixturePFN | XGBoost |
| airlines | 539383 | 7 | 1 | 2.01 | MixturePFN | XGBoost |
| poker-hand | 1025009 | 10 | 10 | 0.08 | XGBoost | MixturePFN |

**Table 15: Mean and Std. Accuracy of MixturePFN and TabPFN* on all datasets across 10-folds (part 1).**
| Dataset | Model | Mean $\pm$ Std Accuracy $\uparrow$ |
| --- | --- | --- |
| australian | TabPFN* | 0.868$\pm$0.036 |
| MixturePFN | 0.861$\pm$0.023 |  |
| bioresponse | TabPFN* | 0.791$\pm$0.018 |
| MixturePFN | 0.793$\pm$0.017 |  |
| gesturepha… | TabPFN* | 0.569$\pm$0.014 |
| MixturePFN | 0.704$\pm$0.011 |  |
| miniboone | TabPFN* | 0.927$\pm$0.003 |
| MixturePFN | 0.946$\pm$0.003 |  |
| speeddating | TabPFN* | 0.856$\pm$0.006 |
| MixturePFN | 0.887$\pm$0.011 |  |
| ada-agnostic | TabPFN* | 0.845$\pm$0.016 |
| MixturePFN | 0.842$\pm$0.013 |  |
| airlines | TabPFN* | 0.600$\pm$0.003 |
| MixturePFN | 0.857$\pm$0.005 |  |
| albert | TabPFN* | 0.638$\pm$0.005 |
| MixturePFN | 0.903$\pm$0.003 |  |
| artificial… | TabPFN* | 0.650$\pm$0.013 |
| MixturePFN | 0.717$\pm$0.008 |  |
| balance-scale | TabPFN* | 0.989$\pm$0.013 |
| MixturePFN | 0.997$\pm$0.010 |  |
| cnae-9 | TabPFN* | 0.896$\pm$0.029 |
| MixturePFN | 0.899$\pm$0.029 |  |
| colic | TabPFN* | 0.823$\pm$0.044 |
| MixturePFN | 0.807$\pm$0.067 |  |
| credit-app… | TabPFN* | 0.884$\pm$0.050 |
| MixturePFN | 0.872$\pm$0.053 |  |
| credit-g | TabPFN* | 0.729$\pm$0.028 |
| MixturePFN | 0.740$\pm$0.021 |  |
| electricity | TabPFN* | 0.812$\pm$0.005 |
| MixturePFN | 0.897$\pm$0.003 |  |
| elevators | TabPFN* | 0.900$\pm$0.006 |
| MixturePFN | 0.905$\pm$0.005 |  |
| guillermo | TabPFN* | 0.791$\pm$0.013 |
| MixturePFN | 0.799$\pm$0.018 |  |
| heart-h | TabPFN* | 0.837$\pm$0.044 |
| MixturePFN | 0.834$\pm$0.046 |  |
| higgs | TabPFN* | 0.665$\pm$0.007 |
| MixturePFN | 0.693$\pm$0.005 |  |
| jasmine | TabPFN* | 0.804$\pm$0.016 |
| MixturePFN | 0.861$\pm$0.008 |  |
| jungle-che… | TabPFN* | 0.823$\pm$0.006 |
| MixturePFN | 0.865$\pm$0.004 |  |
| kc1 | TabPFN* | 0.862$\pm$0.011 |
| MixturePFN | 0.866$\pm$0.013 |  |
| lymph | TabPFN* | 0.810$\pm$0.096 |
| MixturePFN | 0.810$\pm$0.096 |  |
| mfeat-fourier | TabPFN* | 0.828$\pm$0.025 |
| MixturePFN | 0.847$\pm$0.025 |  |
|  |  |  |

**Table 16: Mean and Std. Accuracy of MixturePFN and TabPFN* on all datasets across 10-folds (part 2).**
| Dataset | Model | Mean $\pm$ Std Accuracy $\uparrow$ |
| --- | --- | --- |
| mfeat-zernike | TabPFN* | 0.828$\pm$0.015 |
| MixturePFN | 0.846$\pm$0.024 |  |
| monks-prob… | TabPFN* | 1.000$\pm$0.000 |
| MixturePFN | 1.000$\pm$0.000 |  |
| nomao | TabPFN* | 0.953$\pm$0.003 |
| MixturePFN | 0.966$\pm$0.002 |  |
| phoneme | TabPFN* | 0.883$\pm$0.014 |
| MixturePFN | 0.902$\pm$0.015 |  |
| poker-hand | TabPFN* | 0.517$\pm$0.011 |
| MixturePFN | 0.677$\pm$0.002 |  |
| profb | TabPFN* | 0.691$\pm$0.028 |
| MixturePFN | 0.685$\pm$0.024 |  |
| qsar-biodeg | TabPFN* | 0.885$\pm$0.033 |
| MixturePFN | 0.883$\pm$0.038 |  |
| socmob | TabPFN* | 0.933$\pm$0.016 |
| MixturePFN | 0.929$\pm$0.017 |  |
| splice | TabPFN* | 0.876$\pm$0.018 |
| MixturePFN | 0.983$\pm$0.005 |  |
| vehicle | TabPFN* | 0.847$\pm$0.023 |
| MixturePFN | 0.847$\pm$0.024 |  |
|  |  |  |

## Peer Reviews

confidence: 2
contribution: 1
initial_rating: 3
presentation: 3
questions: I could not find an ablation study on the number of clusters K vs model performance, have you done these experiments?
soundness: 3
strengths: The paper is well written and proposes a justified solution to address the context length issue for in-context learning models such as TabPFN. Authors conduct extensive experiments on many real world dataset to demonstrate the effectiveness of the proposed approach and compare with leading tree-based and deep learning tabular methods.
summary: The paper proposes a mixture of experts approach for in-context learning on tabular data. Each expert in the mixture is a K-means cluster and the model routes the input instance to the closest cluster. This addresses the problem of context size in large datasets and provides a better selection of prompt instance than random sampling. To adapt the model to this type of routing authors also propose fine tuning by selecting a cluster of each training instance and maximizing the likelihood.
weaknesses: There is a very related previous work "Retrieval & Fine-Tuning for In-Context Tabular Models" by Thomas et al, which proposes both nearest neighbor retrieval to improve the prompt and fine tuning with this approach to adapt the model to the target distribution. I think the authors have to compare with this work and highlight what is novel in MixturePFN.

---

confidence: 2
contribution: 3
initial_rating: 6
presentation: 3
questions: See weaknesses.

Can categorical features simply be encoded as ordinal features? Is that not implying false relationships between unordered elements?
soundness: 3
strengths: - The MICP strategy effectively reduces memory usage, allowing the model to handle larger datasets compared to existing TabPFN
- CAPFN bootstrapping and finetuning approach appears to be an effective way to mitigate distribution shift ICL for tabular data
- Extensive benchmarks against 19 strong baselines show good performance in both mean rank and Condorcet ranking across diverse datasets
summary: The paper proposes the MixturePFN framework, which extends TabPFN for large tabular datasets by addressing the performance and scalability limitations of the number of table rows. The authors propose:
1. Mixture of In-Context Prompters (MICP), which optimizes inference by using a sparse mixture of experts to route test samples to specific "prompters" that create context-specific prompts to separate large training datasets into manageable clusters. 
2. Context-Aware Finetuning (CAPFN), which addresses distributional shift issues by specializing each prompter on its assigned
context via parameter efficient finetuning.
weaknesses: - While MIXTUREPFN improves dataset scalability, it still struggles with feature-rich datasets, potentially limiting its applicability in domains with high-dimensional data, such as patient healthcare data. I realize the authors leave this to future work, but this is an area where simple XGBoost performs quite well, and I would be curious about their thoughts on tackling this issue.

- MICP's reliance on K-Means clustering to segment data into meaningful clusters as the quality of clusters can vary significantly based on dataset properties / distance metric chosen. Poor clustering could lead to suboptimal routing and ineffective prompts for certain test samples. I'd be curious to see some ablations in this area.

- The CAPFN bootstrapping method might introduce biases or overfitting if the sampled subsets are not representative of the entire dataset. Bootstrapping from small clusters may fail to capture enough diversity, especially in cases with imbalanced classes or rare features. I'd be also curious to see how this method works with highly imbalanced labels e.g. 1\% positive.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 3
questions: 1. Can you provide a comparison with LoCalPFN [1]? If not possible, I think the comparison should be done using k neighbor samples rather than random sampling, at least for TabPFN*.

2. I see that the authors say in the limitations section that they didn't do it on a dataset with a million samples, but I'm somewhat curious about the effectiveness of MixturePFN on a dataset with a million samples, since the paper is aimed at the scale-up aspect.

3. I'm also curious about the effectiveness of MixturePFN on datasets with hundreds or thousands of features, which is very practical in the real world.

----
[1] Thomas et al., Retrieval & Fine-Tuning for In-Context Tabular Models, NeurIPS 2024
soundness: 3
strengths: 1. The idea of Mixture of Experts blending into TabPFN seems novel.

2. The effectiveness of MixturePFN is well evaluated in well-established benchmarks against a variety of baseline methods.

3. Writing is easy to follow.
summary: In this paper, the authors propose MixturePFN, an extension of Sparse Mixture of Experts to TabPFN to alleviate the context size limitations of the existing TabPFN. On the TabZilla benchmark, MixturePFN outperforms state-of-the-art tabular prediction models.
weaknesses: 1. The biggest weakness I think is that the paper is missing a comparison with LoCalPFN [1]. Since LoCalPFN also tries to make TabPFN effective even on datasets with many-shots, I think it should be mentioned in the paper.

----
[1] Thomas et al., Retrieval & Fine-Tuning for In-Context Tabular Models, NeurIPS 2024

Meta-review: (a) Summary of Scientific Claims and Findings

The paper presents MixturePFN, an enhancement of Sparse Mixture of Experts tailored for tabular Prior-Fitted Networks (TabPFNs). MixturePFN leverages specialized In-Context Learning (ICL) experts, applied to clusters of tabular data, to improve both efficiency and accuracy, particularly on large datasets. The framework consists of two key components: Mixture of In-Context Prompters (MICP): Routes test samples to appropriate clusters, enabling effective segmentation of large datasets, and Context-Aware Fine-Tuning (CAPFN): Mitigates distributional shifts by employing fine-tuning specific to each cluster.

(b) Strengths of the Paper

1. Innovative integration of Mixture of Experts with ICL, designed specifically for tabular data applications.

2. Demonstrates substantial performance improvements over state-of-the-art models, supported by extensive validation on diverse and large datasets.

(c) Weaknesses of the Paper and Missing Elements

1. Limited scalability to datasets with high-dimensional features, such as those found in healthcare.

2. CAPFN may introduce biases due to insufficient cluster diversity, which can be problematic for datasets with imbalanced classes.

(d) Decision and Rationale

The paper offers a significant and well-substantiated contribution to advancing ICL for tabular data. While there are concerns about scalability and generalization, the paper’s strengths and innovative approach outweigh these limitations.

Decision: Accept (Poster)