Mixture of In-Context Prompters for Tabular PFNs

Authors: Derek Qiang Xu, F Olcay Cirit, Reza Asadi, Yizhou Sun, Wei Wang
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
Keywords: Prior-Fitted Networks, Tabular Learning, Sparse Mixture of Experts.

Abstract:
Recent benchmarks find In-Context Learning (ICL) outperforms both deep learning and tree-based algorithms on small tabular datasets. However, on larger datasets, ICL for tabular learning suffers in both efficiency and effectiveness. In terms of efficiency, transformers incur linear space and quadratic time complexity w.r.t. context size. In terms of effectiveness, contexts at inference encounter distribution shift compared to contexts from pretraining. We propose MixturePFN, which extends Sparse Mixture of Experts to the state-of-the-art ICL for tabular learning model. Specifically, MixturePFN finetunes a specialized ICL expert on each cluster of tabular data and routes new test samples to appropriate experts at inference. MixturePFN supports constant-size contexts by splitting large training datasets into more manageable clusters. MixturePFN addresses distribution shift by finetuning an expert on each training dataset cluster via bootstrapping. Extensive experimental results shows MixturePFN outperforms 19 baselines both in mean rank and as the Condorcet winner across 36 diverse tabular datasets under both accuracy and F1 score with statistical significance.

Decision: accept (poster)
Reviewer scores: 6, 8, 6

Reviews and discussion:
<thread>
creation date:
1730102372314

thread_type:
official_review

<thread_object>
creator: reviewer_x42j

<content>
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.

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.

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

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
</content>
</thread_object>

</thread>
<thread>
creation date:
1730681031935

thread_type:
official_review

<thread_object>
creator: reviewer_dzw6

<content>
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.

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

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.

questions:
See weaknesses.

Can categorical features simply be encoded as ordinal features? Is that not implying false relationships between unordered elements?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730687745477

thread_type:
official_review

<thread_object>
creator: reviewer_5xhm

<content>
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.

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.

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.

questions:
I could not find an ablation study on the number of clusters K vs model performance, have you done these experiments?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732223035473

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Common Response

comment:
**[LoCALPFN is a parallel work]**

**Please Read:** Thank you very much for your thorough efforts reviewing our paper and its related works! According to ICLR policy, MixturePFN and LoCALPFN are contemporaneous: *“We consider papers contemporaneous if they are published within the last four months. That means, since our full paper deadline is October 1, if a paper was published (i.e., at a peer-reviewed venue) on or after July 1, 2024, authors are not required to compare their own work to that paper.”*.

In fact, MixturePFN was the first to post on arXiv. Here is a detailed timeline:


- MixturePFN was posted on arXiv in May 2024, proposing MoE KNN and finetuning.

- LoCALPFN was posted on arXiv in June 2024, proposing KNN retrieval and finetuning.

- MixturePFN was submitted to ICLR on October 1 2024.

- LoCALPFN (camera-ready) was accepted to NeurIPS on October 30 2024. After discussing with the authors of LoCALPFN, we both agree our papers are contemporaneous.

We hope our timeline gives a better perspective on the relationship between MixturePFN and LoCALPFN. Although we are *not required to*, we still provide a design and empirical comparison with LoCALPFN for completeness in the next section.

**[Comparison with LoCALPFN]**

Compared to LoCALPFN, which finetunes a single TabPFN model for retrieval, MixturePFN finetunes a specialized expert on each subset of the training/context data, improving effectiveness. Hence, LoCALPFN is a specific instance of MixturePFN: when gamma=inf and there is only one expert. 

Because LoCALPFN does not release source code or performance on individual datasets (at the time of this rebuttal), we reimplement LoCALPFN† under our framework by setting gamma=inf (i.e. KNN retrieval) and limiting the number of experts. We compute average accuracy across several datasets: electricity, phoneme, and airlines.

```
Model Name  | max experts | gamma  | elec. Acc | phon. Acc | air. Acc
------------+-------------+--------+-----------+-----------+----------
TabPFN*     | no finetune | no KNN | 81.2%     | 88.3%     | 60.0%
LoCALPFN†   | 1           | inf    | 85.4%     | 88.7%     | 64.0%
MixturePFN† | 1           | 5      | 85.6%     | 88.5%     | 64.3%
MixturePFN† | 8           | 5      | 88.0%     | 89.1%     | 64.9%
MixturePFN  | 1024        | 5      | 89.7%     | 90.2%     | 85.7%
```

As seen above, MixturePFN outperforms both TabPFN and LoCALPFN by training a specialized expert on each subset of the training dataset. The more experts that are trained, the better the performance.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732463254162

thread_type:
official_comment

<thread_object>
creator: area_chair_j1ho

<content>
title:
Reminder: Author-Reviewer Discussion Period Closing Soon

comment:
This is a reminder that the author-reviewer discussion period will end on Nov 26 AoE.

Your engagement during this phase is critical for providing valuable feedback and clarifications. If you have any remaining questions or comments, please take a moment to participate before the deadline.

Thank you for your contributions to this important process.

AC
</content>
</thread_object>

</thread>


Metareview:
metareview:
(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.