Retrieval & Fine-Tuning for In-Context Tabular Models

Tabular data is a pervasive modality spanning a wide range of domains, and the inherent diversity poses a considerable challenge for deep learning. Recent advancements using transformer-based in-context learning have shown promise on smaller and less complex datasets, but have struggled to scale to larger and more complex ones. To address this limitation, we propose a combination of retrieval and fine-tuning: we can adapt the transformer to a local subset of the data by collecting nearest neighbours, and then perform task-specific fine-tuning with this retrieved set of neighbours in context. Using TabPFN as the base model -- currently the best tabular in-context learner -- and applying our retrieval and fine-tuning scheme on top results in what we call a locally-calibrated PFN, or LoCalPFN. We conduct extensive evaluation on 95 datasets curated by TabZilla from OpenML, upon which we establish a new state-of-the-art with LoCalPFN -- even with respect to tuned tree-based models. Notably, we show a significant boost in performance compared to the base in-context model, demonstrating the efficacy of our approach and advancing the frontier of deep learning in tabular data.

Paper

References (57)

Scroll for more · 38 remaining

Similar papers

Peer review

Reviewer ZbyM5/10 · confidence 5/52024-07-12

Summary

Tabular data is an important yet understudied modality in machine learning. Following recent success of Tab transformer and TabPFN, people have made significant progress on tabular tasks. This paper is an extension of TabPFN, a previous sota tabular learning model on small scale data, and the authors combine with retrieval and finetuning. The motivation is intuitive and the idea is natural. In the retrieval part, the authors reuse the simple kNN trick and discuss in details the performance in TabPFN. In the finetuning part, the authors construct the training data by shared contexts and finetune the TabPFN model. The authors experiment on large scale datasets and proposes the state performance model LoCalPFN.

Strengths

- Tabular learning is an important topic in real life applications - The methodology of retrieval and finetuning is not dependent on ICL tabular model. - The performance gain by retriveal and finetuning is clear wrt TabPFN - Scaling experiment is provided which is important to large scale analysis

Weaknesses

- There are many limitations on the datasets requirements, e.g. number of features, number of classes, tasks - No other deep learning or transformer based models as a baseline - Sec 2.4 is not quite clear in finetuning steps, though the general ideas are appreciated.

Questions

1. L32, memory scales quadratically in the size of the dataset. Do the authors mean exactly the calculation of attention matrix? If so, it seems this memory limit s general and what the authors do is to leverage the limited context length better? 2. In Table 1, is the gain from loCalPFN w.r.t TabPFN-kNN coming from only finetuning per dataset? 3. Just make sure I understand correctly, do the authors fine tune on each dataset of 95 benchmarks? 4. How are tree methods evaluated? Any hyperparameter optimization? 5. What is the training time and inference time on each dataset and how this time scales with dataset size? 6. What is the selection criteria of 95 datasets? How are they chosen, preprocessed, etc.

Rating

5

Confidence

5

Soundness

2

Presentation

2

Contribution

3

Limitations

NA

Reviewer jqEN7/10 · confidence 4/52024-07-12

Summary

The authors extend the recently introduced TabPFN to larger and more complex datasets by fetching a relevant context for each test point using a KNN algorithm. The author evaluate two methods, TabPFN-knn, which consists of using the original TabPFN on the fetched context for each test point, and LocalPFN, which adds a finetuning step to adapt TabPFN to these new kind of local context. For this finetuning, an approximation is devised to use close points together in both the context and the query, enabling a shared context which is faster, but with results close to the result of a KNN for each point. The paper shows extensive evaluations of these two methods, demonstrating state of the art results against the original TabPFN, Gradient Boosting Trees like CatBoost, and (not pretrained) neural networks on a previously introduced benchmark. Furthermore, the authors provide several ablations, for instance showing the importance of jointly adding the knn and finetuning steps.

Strengths

Important and novel contribution, which allow to use TabPFN is many more settings. The evaluation is well done. The authors benchmark their method on a previously introduced dataset, and the selection they use is principled, showing that no cherry-picking of dataset was used. The baselines are strong (modern GBDT like catboost, modern NN like ft transformer, and different tabpfn variants). The ablations are interesting and numerous. For instance, I was wondering about TabPFN-3k-32ens-int and was happy to see it. The importance of doing the finetuning jointly with the knn is interesting. The approximation used are also ablated, like in table 12. The paper is very well introduced and written, and make for a very pleasant read.

Weaknesses

Some details are missing, for instance: - how are runtimes computed for figure 11? Which hardware? Which hyperparameters? - I'm not sure whether LocalPFN undergoes some HPO or not (I think not). If it's not the case, I think it should be made clearer, as it makes the comparison to tuned GBDT more impressive. If it does, the hyperparameter space should be provided. I think TabPFN-**1k**-32ens-int (maybe less than 32 ens if it's slow?) would be an interesting baselines as TabPFN is optimized for smaller datasets than 3K. I think more aggregation metrics should be provided (for instance mean rank, mean normalized score (z-score or other rescaled scores)..). figure 4 lacks details: which datasets are used? Are the different datasets subsampled to reach smaller sample sizes? I also don't understand why absolute mean AUC is decreasing when the dataset size is increased after a certain point in Figure 7. Is it because the list of datasets is changing? The code is not currently available.

Questions

For inference, you have to use a different context for each points to predict, right? How long / memory intensive is it? What is the runtime repartition between finetuning and inference? Do you think you can finetune with a smaller k than what you use for inference? > We use the experimental results from TabZilla [35] when they are available. I assume they are always available expect for TabPFN variants? I also think that the fact that your using the hyperparameter spaces from [35] should be said in the main text (if that's indeed the case). Which distance are you using in faiss? I'm wondering if using the euclidean distance would improve the performance of the ordinal encoded features for the knn (compared to using one-hot-encoding).

Rating

7

Confidence

4

Soundness

3

Presentation

4

Contribution

4

Limitations

The authors have adequately addressed the limitation of their work.

Authorsrebuttal2024-08-12

Indeed, we had only reported the full training+eval loop runtime for the algorithm but not the inference time specifically before. Let's get into more details. To classify $N_\text{qy}$ examples given $N_\text{ctx}$ points, TabPFN would contruct a tensor of size ($L=N_\text{ctx}+N_\text{qy}, B=1, d$) with appropriate masking. As we do not share the context across queries, with TabPFN kNN we need to have a specific context for each point and as such our input tensor is size ($L=N_\text{ctx}+1, B = N_\text{qy}, d$) where the $+1$ is the query point to classify in each batch dimension. For instance, for 512 query points and 1000 context size/neighbors, the inference speed of TabPFN would be about 0.01s. A naive TabPFN-kNN is about 1s. However using bf16 precision we can lower it to about 0.4s (0.008s for TabPFN). If we used 128 queries, the number would be 0.008s (TabPFN) vs 0.1s (TabPFN-kNN). So as we observe, we are slower than TabPFN. In the 71 datasets in the attached pdf on which the deep learning baselines were able to run according to the Tabzilla design, the largest of the 95 datasets are typically excluded (this is because Tabzilla allocate a time limit for each algorithm and dataset, as such many --slow-- DL baselines do not run on the largest datasets in the given time budget). This is why TabPFN-kNN is really fast. On fig 12, including all datasets the average speed is about 10x slower but still fast considering there is no training. The fact that TabPFN-kNN is slower (and takes a lot more memory, especially when having to perform backprop as we do when performing fine-tuning end-to-end with retrieval) is why we introduced the approximate NN context sharing for fine-tuning. All in all, TabPFN-kNN is not the algorithm with the fastest inference speed, TabPFN is much faster and so is XGBoost etc... However we believe this algorithm still has very important strengths: it scales extremely well with the dataset size (compared to TabPFN whose performance degrades strongly), so in the case of production ML models which have access to a very large dataset but only have to classify new queries at a lower rate, it is a very well suited algorithm. Furthermore, indexing on large datasets is usually much faster than retraining any ML algorithm, as such it can adapt very fast to changes in data distribution (such as covid for instance) without needing expensive retraining. Note that the evaluation runtime was never the bottleneck in our research, so this is not what we focused on, but there are other simple ways one could improve upon it. For example, we can use the clustering capabilities of faiss to split the training data into a fixed number of contexts. This way we can ensure many queries will share the same context and harness the speed of TabPFN. While it may degrade the accuracy of the NN search, we think it can be an interesting trade-off if inference speed is an issue. Thank you for the interesting comment, we will include numbers or a figure specifically concerning the inference speed for clarity. Have you received the anonymized code link from the AC?

Reviewer jqEN2024-08-12

Thank you for the very clear answer! > Have you received the anonymized code link from the AC? No

Authorsrebuttal2024-08-12

Here we share the anonymous link for our code: https://anonymous.4open.science/r/retrieve_ft-E2B0/README.md Note that we removed some files (including the notebooks) to make sure anonymity is preserved. Thus the code for the analysis figures is not available in this repo (but we included that dataframe with all results). The code should be able to run, if you have a GPU with limited memory (<48G) we recommend lowering the batch size or number of neighbors (--context_length 500 for instance). If you are familiar with the TabPFN code base you will notice that a large portion of the code has been rewritten, notably to be compatible with the pytorch transformers class. You will probably be interested in the methods/pfknn.py (TabPFN-kNN) and methods/ftknn.py (LoCalPFN) files. There are currently a lot of options available in the code (different ways of computing neighbors, different embeddings etc..) which can make the code harder to read. Let us know if you have any questions (and please do not share this code)!

Reviewer iFbB6/10 · confidence 3/52024-07-12

Summary

The paper introduces Locally-Calibrated PFN (LoCalPFN), an advanced model for tabular data that enhances the transformer-based TabPFN by incorporating retrieval and fine-tuning techniques. By using k-Nearest Neighbours (kNN) to select a local context for each data point and fine-tuning the model on this retrieved set, LoCalPFN adapts more effectively to larger and more complex datasets. Extensive evaluations on 95 datasets from TabZilla demonstrate that LoCalPFN outperforms both neural and tree-based methods, setting a new state-of-the-art in tabular data classification. The key contributions include addressing TabPFN's scaling issues, proposing an improved context usage with retrieval and fine-tuning, and showcasing superior performance through extensive experimentation and ablation studies.

Strengths

- The paper introduces Locally-Calibrated PFN (LoCalPFN), which enhances transformer-based in-context learning for tabular data by combining retrieval and fine-tuning techniques. - The research is robust, with extensive evaluations on 95 datasets from TabZilla. The authors provide comprehensive experimentation and analysis, demonstrating the effectiveness of LoCalPFN compared to strong baselines, including neural and tree-based models. - The paper is well-organized and clearly written, making the complex concepts and methods accessible.

Weaknesses

- Although LoCalPFN shows improved performance, the fine-tuning process increases computational complexity and runtime, especially with large datasets. - The paper's reliance on TabPFN as the base model restricts the generalizability of the proposed method. While LoCalPFN demonstrates significant improvements, it remains unclear whether these benefits would transfer to other in-context learning models for tabular data. - The current implementation of LoCalPFN is constrained by TabPFN’s limitations on the number of features and classes, as well as its incompatibility with regression tasks.

Questions

- Have you tested the retrieval and fine-tuning techniques used in LoCalPFN on other in-context learning models? - Given the constraints on the number of features and classes due to TabPFN, how do you plan to extend LoCalPFN to handle datasets with more features and classes?

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

2

Limitations

- The paper focuses solely on classification tasks and does not address regression tasks. This omission leaves a gap in evaluating the full potential of LoCalPFN.

Reviewer j47T2/10 · confidence 3/52024-07-14

Summary

The paper proposes LoCalPFN, a new method that improves the scaling of transformer-based in-context learning for tabular data. It uses retrieval and fine-tuning to adapt the transformer to local subsets of the data, and demonstrates state-of-the-art performance on a variety of datasets. The paper makes the following contributions: * Provides a comprehensive analysis of TabPFN and identifies key limitations in TabPFN's ability to scale with dataset size and complexity. * Proposes LoCalPFN, a novel approach that combines retrieval and fine-tuning techniques. LoCalPFN leverages these methods to enable more effective utilization of context, thereby enhancing the scalability of in-context learning for tabular data. * Demonstrates the superiority of LoCalPFN through extensive evaluation and ablation studies. These evaluations reveal that LoCalPFN outperforms baselines on a wide range of datasets

Strengths

* Novel Combination of Retrieval and Fine-tuning: The paper introduces an new approach by combining retrieval and fine-tuning for tabular data. * Extensive Evaluation: The paper presents a thorough evaluation of the proposed method (LoCalPFN) across a wide range of datasets, comparing it with numerous baselines and conducting ablation studies.

Weaknesses

The paper contains several inaccurate statements, For example, in abstract, "Recent advancements using transformer-based in-context learning have shown promise on smaller and less complex datasets, but have struggled to scale to larger and more complex ones." is inaccurate -- depending on the dataset and task definitions. There are many such examples in the paper, which make it hard to understand the scope of the paper. It raises concerns about the overall reliability of the paper's findings. The paper fails to clearly define the complexity of tabular datasets, making it difficult to assess the relevance and significance of the proposed method (LoCalPFN) in addressing the challenges of complex datasets. The paper does not address the important practical considerations of the model's cost and latency at inference time, especially retrieval is used. This omission is a significant limitation, as it hinders the understanding of the model's real-world applicability. As acknowledged in the paper, LoCalPFN has a slower runtime compared to tree-based methods, which are a popular choice for tabular data. While the paper mentions a faster variant (TabPFN-kNN), its performance is not as strong as LoCalPFN." The proposed LoCalPFN method is heavily reliant on TabPFN as the base model. This raises concerns about the generalizability of the approach to other base models and task types (e.g., regression).

Questions

1) why the dataset size would affect the model quality? is it related to the task complexity and diversity of the samples? 2) can you given an example of the query x_qy for tabular data? and the distance metrics being used in kNN?

Rating

2

Confidence

3

Soundness

2

Presentation

1

Contribution

1

Limitations

The paper does not address the important practical considerations of the model's cost and latency at inference time, especially retrieval is used. This omission is a significant limitation, as it hinders the understanding of the model's real-world applicability. As acknowledged in the paper, LoCalPFN has a slower runtime compared to tree-based methods, which are a popular choice for tabular data. While the paper mentions a faster variant (TabPFN-kNN), its performance is not as strong as LoCalPFN." The proposed LoCalPFN method is heavily reliant on TabPFN as the base model. This raises concerns about the generalizability of the approach to other base models and task types (e.g., regression).

Authorsrebuttal2024-08-12

Request for Comment

Hi, We would like to again thank you for taking the time to review our paper, but would also very much appreciate if took the opportunity to engage in discussion if time permits, considering that we have clarified some of the more negative points of your review both in the rebuttal to you and the shared rebuttal to all reviewers (that one has a PDF attached). Thanks!

Reviewer jqEN2024-08-10

Thank you for the detailed answer, and for the new tables in the pdf! The thing which is still not clear for me is the inference time. I think it would be useful to report both the training and inference time separately. In many settings people care about having very low inference time, and this is already a limitation of TabPFN (very fast training but quite slow inference time). Something like inference time / 1000 samples would be interesting to report in the paper. I'm still wondering if for TabPFN-KNN you need one forward pass per inference point (which should be much slower than TabPFN). In Figure 11 TabPFN-KNN seems really fast, so I'm wondering if I'm missing something or if you have some tricks.

Reviewer ZbyM2024-08-11

Thank the authors for the feedback. I am still at the borderline for its improvement wrt TabPFN. Happy to hear what other reviewers think

Authorsrebuttal2024-08-12

Thank you for engaging with us! We have tried to address the point of the limitations of TabPFN by explaining how to expand to more classes and provided a regression experiment. Do you have a specific limitation in mind you would like us to address? Furthermore we would like to point out that our contribution is the retrieval and fine-tuning procedure (and their efficient implementation such as the shared context approximation or the joint retrieval+fine-tuning), our paper is not concerned with modifications to the base model (TabPFN) which is definitely important, but orthogonal, future work. We believe we addressed your concerns, which notably included the deep learning and their runtime experiment, clarification on the fact that HPO is indeed used for our baselines (30 rounds), and answering specific clarifications alongside with the concerns about TabPFN's limitations which are addressed in the general message and discussed above. We hope that our additional experiments and explanations will be reflected in your evaluation of our work.

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC