Retrieval & Fine-Tuning for In-Context Tabular Models

Authors: valentin  thomas
junwei  ma
rasa  hosseinzadeh
keyvan  golestan
guangwei  yu
maksims  volkovs
anthony  caterini
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 23:00:58

Abstract: Tabular data is a pervasive modality spanning a wide range of domains, and this 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 tabular 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.

Main paper: 
 introduction: 
 Tabular data is the most pervasive modality for practical problems in data science, spanning across a wide variety of domains including finance, healthcare, and science (Benjelloun et al., 2020; Ulmer et al., 2020; Clements et al., 2020; Tang et al., 2020; Urban & Gates, 2021) . The diversity and heterogeneity of tabular data pose great challenges for deep learning approaches (Grinsztajn et al., 2022) , unlike modalities such as text and image in which neural networks can be designed to specifically exploit inductive biases underlying the data (Borisov et al., 2022) . As such, obtaining a performant neural network on a particular tabular data task often results in expensive iterations of training and hyperparameter tuning. Meanwhile, tree-based methods such as XGBoost (Chen & Guestrin, 2016) and CatBoost (Prokhorenkova et al., 2018) have proven to be more robust to the inherent challenges of tabular data, and thus have remained the dominant approach for this setting (Grinsztajn et al., 2022; Shwartz-Ziv & Armon, 2022; Borisov et al., 2022 ). Yet recently, there has been progress made with transformers and In-Context Learning (ICL): one such example is TabPFN (Hollmann et al., 2023) , which is trained using a prior-fitting procedure (Müller et al., 2022) that exposes the network to millions of possible data-generating processes, thus taking a step towards encapsulating the heterogeneity of tabular data. Such approaches differ from classical algorithms in that they process entirely new datasets in a single forward pass and obviate the need for training and hyperparameter tuning.Despite the promise of transformer-based ICL methods in the tabular setting -particularly on smaller datasetsscaling remains an issue: memory scales quadratically in the size of the context. This limits performance when the entire dataset cannot fit into memory, and contrasts with classical algorithms that tend to improve as the amount of available data increases. In addition to this, and as depicted in Figure 1 , TabPFN in particular can struggle with underfitting as dataset complexity increases, even when the entire dataset fits into the context; we observe this shortcoming in real datasets as well, and suspect this could apply to any ICL-based model for tabular data.To improve the scaling of tabular ICL methods in both dataset size and complexity, we draw on two techniques that have been incredibly successful in foundational large language models: retrieval (Lewis et al., 2020) and fine-tuning (Bommasani et al., 2021) . On the retrieval side, we use the k-Nearest Neighbours (kNN) of a given query point as the context for classification; modifying the context in this way empirically allows for both enhanced processing of larger datasets and more complex decision boundaries. We also fine-tune end-to-end for each task, using an approximate neighbour scheme to c) We observe that this approach is robust to the numbers of neighbours used (k) even when the dataset complexity increases and always performs better than vanilla TabPFN using full context (N = 1000). Each point is averaged over 25 seeds.facilitate backpropagation, and demonstrate significant performance gains beyond just kNN. We named our model Locally-Calibrated PFN -or LoCalPFN for short -to represent the addition of retrieval and fine-tuning on top of a base TabPFN model, although this idea should naturally transfer to potential future ICL-based tabular foundation models as well (van Breugel & van der Schaar, 2024) . We demonstrate that LoCalPFN is state-of-the-art against both neural approaches and well-tuned tree-based techniques across a 95-dataset benchmark from TabZilla (McElfresh et al., 2023) . We summarize our contributions below:1. Provide insights into TabPFN -the current state-ofthe-art tabular ICL transformer-based framework -and analyze how its performance scales across several axes in both synthetic and real datasets. We identify a failure to scale in both dataset size and complexity.2. Propose LoCalPFN to address the scaling failures mentioned above, using a combination of retrieval and finetuning to allow for more effective use of the context.3. Show LoCalPFN compares favourably to strong baselines on a large variety of datasets through extensive experimentation, analysis, and ablation. 
improving tabular in-context learning with retrieval and fine-tuning: 
 In this section, we describe ICL applied to tabular data -in particular TabPFN -and the limitations of such an approach. Then, we present our contributions where we treat the incontext learner as a base model on top of which retrieval and fine-tuning are applied.Our method generally applies to in-context learners, specifically for classification tasks on tabular data. While, at the time of writing the only successful model of that type is TabPFN (Hollmann et al., 2023) , we expect other such base models to be published in the future. TabPFN is trained using a prior-fitting procedure (Müller et al., 2022) where a large number of synthetic datasets are generated using randomly initialized neural networks. This approach trains an underlying transformer-based network on various generative processes designed to simulate the diverse interrelations that exist among the features of realistic tabular datasets.After the prior-fitting procedure, the learned TabPFN model ingests an entire training dataset D train ≜ {(x i train , y i train )} N i=1 consisting of feature-label pairs x i train ∈ R D and y i train ∈ {1, . . . , C} for i ∈ {1, . . . , N }, along with features of a query point x qy (potentially in a batch), and outputs a distribution over labels y qy ∈ {1, . . . , C}. Specifically, denoting the TabPFN network (outputting logits) as f , the resulting posterior predictive distribution is modelled by:EQUATIONwhere [•] denotes the vector indexing operation.Contrary to classical machine learning methods which are trained on one dataset and then evaluated on the same distribution, TabPFN has been shown to be able to perform classification on a wide range of tasks without training, thanks to its diverse prior-fitting procedure. This makes it one of the rare foundation models for tabular data. Key to this is the ICL ability of TabPFN: by using various training examples as context, analogous to how transformers on language use the preceding tokens as context, TabPFN can classify new query points in a single forward pass.The quadratic growth of memory usage with context length in transformers presents a challenge: the number of support examples we can use is limited. For instance, while TabPFN performs best on small and simple datasets, where the entire training set fits within the context, it is unclear how to best use TabPFN for large and complex datasets. Naïvely, we might consider a random subsample of the training data as context (McElfresh et al., 2023; Feuer et al., 2024) . However, Ma et al. (2024) show that this method does not scale either and observe a drop in performance as the dataset size increases.Given these limitations, it is natural to ask "What constitutes a good context for tabular data?". This topic has been thoroughly researched in natural language processing, which resulted in various techniques for prompt engineering.The situation is more complicated in the tabular domain, as there is no natural order to tabular data as opposed to the natural order of the words in language.Specifically for TabPFN, some attempts have been made to use a summary of the dataset as context, through either k-means centroids (Feuer et al., 2024) or direct prompt optimization (Feuer et al., 2024; Ma et al., 2024) . Yet in either case the flexibility of the method is limited by the use of a single context for all query points. Instead, we propose a different approach here, where we use a local context tailored to each individual point we wish to classify. For tabular data, we hypothesize that the most critical information to classify a query point x qy is contained in its vicinity.Extensive evaluations (McElfresh et al., 2023) support this fact by showing that a simple kNN classifier can rival modern deep architectures designed for tabular data, such as TabNet (Arik & Pfister, 2021) and VIME (Yoon et al., 2020) . We thus believe that using nearby points as context is a good inductive bias for tabular data classification.To do this, the first step is to replace the global context by a local context, i.e., with kNN(x qy ) as the k-nearest neighbours of the query x qy in the training data D train , we replace equation 1 byp θ (y | x qy , D train ) = exp(f θ (x qy , kNN(x qy ))[y]) C c=1 exp(f θ (x qy , kNN(x qy ))[c]).(2)Better Expressivity It is well known that in kNN regression and classification, the number of neighbours k controls the bias/variance trade-off and as such the expressivity of the model. More precisely, large k tends to "oversmooth" and suffer from high bias/underfitting, while small k enables more complex decision boundaries but can suffer from more variance/overfitting (Hastie et al., 2009) . We show that this phenomenon is still true for transformers, beyond the simple kNN classifier, in Figure 1 . We generate datasets of size N = 1000 so that it can be used as context by TabPFN without subsampling. As we increase the complexity of the dataset, measured by the number of concentric circles in this case, TabPFN fails to accurately classify (e.g., for 3 pairs of circles in (a) and more generally in (c)). Retrieving fewer samples (k = 10, 30, 100, or 300) for each query point using its k-nearest neighbours from the training data leads to large improvements in AUC over TabPFN as the complexity of the data increases ((b) and (c)). Note that k = 1000 would correspond to using all samples as context, and thus is equivalent to vanilla TabPFN. As such there is a continuum between TabPFN using the full dataset as context and our local context method using kNN, which we call TabPFN-kNN.While Figure 1 is on toy synthetic data, we believe this result remains surprising: a priori, we would expect a 25-million-parameter model (TabPFN) to be able to learn a few circles, even with just ICL. Meanwhile, we believe that using local contexts allows TabPFN to fit more complex patterns, such as the three circles of Figure 1 , in the same way that using local linear regression enables more expressive (and in that case nonlinear) decision boundaries (Cleveland & Devlin, 1988; Hastie, 2017) .Better Scaling Using a local context has another benefit: it allows our method's performance to scale with the training dataset size. In machine learning, it is generally expected that the performance of an algorithm improves as the training set size N increases, since the empirical risk converges to the expected risk (Vapnik, 2013) . However, ICL-based methods (such as TabPFN) that require subsampling when the maximum context length is smaller than N do not scale with N . TabPFN-kNN, on the other hand, can still benefit from larger training set sizes N even when the number of neighbours k is much smaller than N , as the search is performed over the whole training set. We demonstrate this fact in Figure 2 for three real datasets. While the exact patterns in the loss curves differ, we observe a similar trend across many datasets, where the benefits of using TabPFN-kNN grow as the dataset becomes larger. In Figure 10 we provide more detailed figures which include training loss.In addition to retrieval, we fine-tune the model end-to-end on each dataset to further improve performance, as is common in Retrieval-Augmented Generation (RAG) (Lewis et al., 2020) . However, naïve fine-tuning is not computationally efficient. 2 5 2 6 2 7 2 8 2 9 2 10 2 11 2 12 2 13 Dataset size text for all points, with B = 1, L ctx the training dataset size (or maximum context length if too large), and L qy = N qy the number of points to classify. Contrary to text, there is no auto-regressive attention mask: the context examples all attend to each other (blue arrows on Figure 3 ) while the queries only attend to the context and not to each other (red arrows on Figure 3 ). Therefore, the predicted classes can be computed in parallel and at a reduced memory footprint.By comparison, when using a local context with exact neighbours, the context is no longer shared, and therefore the batch dimension must be used for queries: the input has shape B = N qy , L ctx = k -the number of neighboursand L qy = 1, since the queries use distinct contexts. This is significantly less efficient than the inference performed by TabPFN, which both requires much less memory, and also allows the queries to be processed in parallel. Therefore, our main limitation is in fact the forward and backward passes when using exact neighbors, unlike most applications where retrieval is the bottleneck. As such, most common approximate kNN methods cannot address this issue. Instead, to improve computational efficiency during the end-to-end fine-tuning, we opt for a simple neighbour approximation technique wherein many queries share the same context. An d) . This allows us to efficiently trade-off accuracy of the neighbours versus computational complexity: with lower B we share contexts between many points but this comes at the cost of an approximation in the kNN search as the notion of neighborhood is not transitive, i.e., the neighbour of your neighbour might not be your neighbour. However each sequence in each batch only contains examples which are in the general vicinity of each other. In practice, we observe that this method does not lead to any significant degradation in performance while allowing much faster training. 
related work: 
 Foundational Techniques for Tabular Deep Learners Deep learning techniques have historically struggled on tabular data (Grinsztajn et al., 2022) , where inductive biases are much harder to capture architecturally (Beyazit et al., 2023) as compared to text or images. The comparative lack of progress on a large foundation model for tabular data (van Breugel & van der Schaar, 2024 ) is yet more evidence of this. However, recent approaches have successfully begun to leverage foundational ideas to improve performance. For example, Non-Parametric Transformers (Kossen et al., 2021) and SAINT (Somepalli et al., 2021) both combine row-attentive transformer-based backbones with some form of self-supervised pre-training; however, the former is limited by context size (a common theme for naïve ICL-based learners), whereas the latter is not based on ICL and thus does not as easily apply to novel datasets. Models such as RIM (Qin et al., 2021) and TabR (Gorishniy et al., 2024) on the other hand demonstrate how to effectively design tabular deep learners incorporating retrieval modules, but still require costly and brittle rounds of hyperparameter tuning to adapt to any specific dataset. Our approach is meant to target some combination of all these methods: provide ICLbased generalization capabilities, but without limitations on the context size. The retrieval mechanism within TabR itself relies on kNN, which is one of the most straightforward and widely used retrieval-based machine learning methods (Hastie et al., 2009) . In fact, kNN is still being actively studied in the literature, e.g., in Differential Nearest Neighbours Regression (DNNR) (Nader et al., 2022) , which aims to make kNN differentiable; this showcases the potential of simple methods like kNN in different forms, although DNNR tackles a separate scope from our method.TabPFN and Extensions TabPFN (Hollmann et al., 2023) is a transformer-based in-context learner that has emerged as a popular model for tabular data, demonstrating strong performance on some benchmarks (McElfresh et al., 2023) . It uses a prior-fitting process (Müller et al., 2022) allowing for rapid adaptation to new tasks. This strong ability to quickly generalize makes TabPFN somewhat of a foundation model for tabular data (van Breugel & van der Schaar, 2024), from which techniques for generation (Ma et al., 2023) and dataset distillation (Ma et al., 2024) for example can emerge -interpretability is also being studied (Rundel et al., 2024) . TuneTables (Feuer et al., 2024) attempts to use tabular sketching (Munteanu & Schwiegelshohn, 2018) to summarize the incoming dataset and more effectively scale TabPFN's context; however, much like Ma et al. (2024) , this approach is limited by the use of a single context for all datapoints, as opposed to an adaptive local context. Lastly, den Breejen et al. ( 2023) is able to show some limited improvements by fine-tuning TabPFN, which we extend here by more closely pairing the retrieval and fine-tuning aspects.The idea of pre-training a model on corpora of text prior to fine-tuning has been explored in the Natural Language Processing domain for both classification and generation tasks (Dai & Le, 2015; Howard & Ruder, 2018; Radford et al., 2018) . Later iterations refined this idea to train a model and use its in-context learning abilities for new tasks (Brown et al., 2020) . This elicited research into prompt engineering to determine what to actually put in a model's context (Nye et al., 2021; Wei et al., 2022) . Similar to prompt engineering, to better utilize the model's context, one can search for similar examples from a corpora and use them to facilitate the task; this is known as Retrieval-Augmented Generation (RAG) (Lewis et al., 2020) in the generative context. Other variants of the idea include training jointly with retrieval (Guu et al., 2020; Borgeaud et al., 2022) and augmenting the output of the model with kNN via interpolating (Khandelwal et al., 2019) . These ideas are analogous to our approach of (i) fine-tuning and retrieving jointly, and (ii) disjoint kNN and fine-tuning in our ablations, respectively. LLMs have also been directly applied to tabular data (Dinh et al., 2022; Hegselmann et al., 2023; Fang et al., 2024) however, due to the pre-training of these foundation models on large text corpora, there is the possibility of data leakage, which causes concern with evaluations (Bordt et al., 2024) . Note that this is not the case with TabPFN as it has been trained on synthetic data. 
experimental setup: 
 We evaluate our methods against competitive baselines using 95 out of the 176 datasets from TabZilla (McElfresh et al., 2023) , originally sourced from OpenML (Bischl et al., 2021) . These datasets originate from diverse sources, including academic research, competitions, government agencies, and corporations. The 95 datasets are filtered from TabZilla to meet TabPFN's architectural requirements by ensuring that each dataset has at most 100 features, at most 10 classes, does not contain NaN values, and has at least one instance per class for each split. The details of the datasets are described in Appendix A.1. We further split the datasets into two subsets: "small" datasets which contain less than 2,000 instances, and "medium/large" which contain at least 2,000 instances (up to 130,064). For each dataset, we use the splits from TabZilla with train-validation-test ratio of 80:10:10. Since TabPFN was trained with a maximum of 1,024 data points as context size, the small datasets are roughly considered in-distribution for TabPFN whereas the large datasets are considered out-of-distribution.We conduct our experiments using 10-fold cross-validation over all datasets for all methods. For all baselines, we apply 30 rounds of hyperparameter tuning as in McElfresh et al. (2023) and choose the best hyperparameters for each fold according to its validation AUC. In addition, the TabPFN baseline is reported without further ensembling or transformations, unless otherwise noted. Our methods also build on top of this same TabPFN baseline without further processing. We also compare against TabPFN with transformations in Section 4.4. More details of the baseline models can be found in Appendix A.2.1. We use the faiss (Johnson et al., 2019; Douze et al., 2024) library for efficient kNN search in our methods; this enables us to harness parallel computation to accelerate the nearest neighbour search. We evaluate our methods TabPFN-kNN and LoCalPFN against other models in the following sections. Notably, without further fine-tuning, LoCalPFN is identical to TabPFN-kNN. Details of our method are in Appendix A.2.2.Note on evaluation and the computation of proper confidence intervals: While many works evaluate tabular data methods on a small set of datasets and report confidence intervals/standard deviations for those, we choose to evaluate on a large number of datasets in order to have more meaningful results. However, this makes it harder to compute meaningful uncertainty. Agarwal et al. (2021) dealt with a related problem in reinforcement learning; we follow their lead by, for example, reporting the interquartile mean (IQM, i.e., the mean of the middle 50% of scores), and we use their library to compute 95% confidence intervals via stratified bootstrapping.As shown in Section 4.2, averaged over 95 datasets, Lo-CalPFN outperforms all other baselines, with significant improvement over TabPFN itself. Among the 47 small datasets, we found that TabPFN is in fact quite competitive with other methods, similar to what had been reported by McElfresh et al. (2023) . Nevertheless, LoCalPFN further improves the performance even in this setting and positions itself as the best method. For the 48 medium/large datasets, TabPFN underperforms the tree-based methods by a wide margin. Simply applying kNN on top of TabPFN leads to a drastic performance increase on top of TabPFN. Finally, LoCalPFN further improves on TabPFN-kNN, and either performs on par with, or outperforms, all other methods. We also measure the accuracy and F1 score over all datasets and see a similar pattern; those details can be found in Table 6 and Table 7 .Deep Learning Model Comparisons: Note that most deep learning baselines are significantly more expensive to train and tune on larger datasets, and as such, most of them could not be run on all datasets (McElfresh et al., 2023) . Nevertheless, in Table 5 we compare TabPFN-kNN and LoCalPFN to other deep learning based methods on the datasets on which the baselines have been able to run, and show an even larger improvement in performance. The datasets we used for this comparison can be found in Table 4 .In this section, we further validate that LoCalPFN addresses the scaling problems of TabPFN. We see in Figure 1 and 2 that TabPFN scales badly with both size and complexity; here, we verify this phenomenon in real datasets. While this may appear contradictory to Table 1 of McElfresh et al. (2023) , which shows TabPFN excelling on a large benchmark suite, we note that the aforementioned study mostly contained small datasets and thus it did not show the same performance drop-off observed here.Scaling with Size In Figure 5 , we report the AUC of different algorithms relative to the AUC of Random Forest for different dataset sizes. We choose relative AUC for clarity as there is no clear correlation between the maximum AUC attainable on a dataset and its size.We see that, compared to the Random Forest baseline, TabPFN's performance drops drastically when the dataset size increases beyond 3,000, indicating poor scaling with N . On the other hand, the other methods we report scale more favourably with the dataset size. We also see that LoCalPFN scales favourably compared to the Random Forest baseline, and even outperforms XGBoost for large datasets. Error bars represent the 95% confidence interval.Scaling with Complexity While in Figure 1 we could easily control the complexity of the task, there is no generally accepted measure of complexity for an arbitrary dataset. Here, we propose a simple proxy for complexity: for a given dataset, we measure the difference between the best and worst AUCs of a given set of algorithms, similarly to McElfresh et al. (2023) . The rationale is that AUC itself cannot capture complexity, as for instance learning to separate two Gaussians can be done optimally by a linear classifier, but the error rate depends on their variance. In Figure 6 , we analyze performance across different levels of this complexity measure. We first calculate the difference in AUC for each dataset using all listed methods in Section 4.2, then we divide the datasets into five quantiles on the x-axis, with increasing complexity as we move to the right; on the y-axis, we report the mean AUC relative to Random Forest across 10 folds and across the datasets in each bin. We see that TabPFN scales poorly with increasing complexity, and LoCalPFN still outperforms all other methods in the quantiles of higher complexity, demonstrating that its improvements are not just limited to "easy" datasets.In this section, we provide ablation studies on different design choices for LoCalPFN. 8 , we see that this indeed improves performance over TabPFN. However, applying TabPFN-kNN on top of a naïvely fine-tuned model does not improve performance further. Therefore, it is crucial to fine-tune the model end-to-end with the retrieval (LocalPFN).We also try different embeddings.The simplest approach is to use the raw standardized features for the nearest neighbour retrieval. In Figure 7 (centre) and Table 9 , it is shown that this simple approach is actually very competitive. We compare it to two additional approaches: using one hot encodings (when the size of the resulting vector does not exceed 100 features), and the output of the encoder layer of TabPFN. For the latter, we recompute the search index every 30 gradient steps. The results show that the former, using one hot encodings, does lead to some improvement, however mostly for smaller datasets (see Figure 11 and Figure 12 ).Why do simple embeddings work so well? While Tabular data is complex in many regards (Grinsztajn et al., 2022) , features in tabular data are often semantically meaningful.For this reason, we expect distances that decompose over individual features (i.e d(x, x ′ ) = i d i (x i , x ′ i )) to be a good inductive bias for tabular data, especially when it is normalized. This would not be the case for most natural signals.Up until now, we have mostly compared to TabPFN with a random context of size 1,000. To prove our point that using a 2023)), which we denote TabPFN-32ens, and then by increasing the context size of the ensembled TabPFN to 3,000 (TabPFN-3k-32ens). As depicted in Figure 7 (right) and detailed in Table 10 , while improving significantly upon TabPFN, these are still not competitive even with our TabPFN-kNN. As one can criticize the use of a single context to classify queries, we further experimented with a "Bayesian" view of the probability by averaging it over contextsp θ (y qy | x qy , D train ) ≜ C p θ (y qy | x qy , C)p(C | D train )dC,where C is a context obtained from the training data D train . We experimented with splitting D train into chunks of size 3,000, and averaging the probabilities over those chunks. We call this method TabPFN-32ens-3k-int (for integral) and show that, while it does improve upon the single random context, it does not outperform TabPFN-kNN. Additionally, this method is very expensive as: (i) using 3,000 context examples is GPU memory intensive, and (ii) the integral over chunks makes the inference scale as O(N ). The last method we compared to is "In-Context Distillation" (Ma et al., 2024) (TabPFN-ICD) where, similarly to Feuer et al. (2024) , the authors directly optimize the context. While this last method leads to better performance (including on larger datasets, see Figure 12 ), since it performs task-specific tuning it is more comparable to LoCalPFN, which remains superior. Left: Fine-tuning jointly with retrieval yields better performance. Centre: The choice of embeddings for retrieval does not change the performance drastically but can lead to some improvements. Right: Methods using a context that does not depend on the current query do not match the performance of methods that use a local context. We also ablate the choice of the number of neighbours used as context. This is the only hyperparameter for TabPFN-kNN and also an important hyperparameter for LoCalPFN. In practice, for the number of neighbours, we use the minimum of (i) 10 times the square root of the training set size, and (ii) a pre-defined maximum. For large datasets, the number of neighbours should roughly align with the pre-defined maximum. In Figure 8 , we vary this pre-defined maximum while observing the mean AUC on the 48 medium/large datasets. We found that TabPFN-kNN is not very sensitive to this choice as long as it is at least 100. We also see that LoCalPFN is able to improve TabPFN-kNN on all context sizes. Surprisingly, we observe that LoCalPFN is able to outperform the random forest baseline using a maximum context size of only 50, and also outperform the XGBoost baseline with maximum context size of 500. The details of the ablation can be found in Table 11 .Other Ablations We also ascertain the quality of the approximate local context in real datasets in Appendix A.5.5 and Table 12 , and we provide a runtime analysis of various methods in Appendix A.5.6 and Figure 13 . 
conclusion and limitations: 
 In this paper, we demonstrate how to use retrieval and fine-tuning to improve performance on tabular data classification tasks by introducing LoCalPFN as a version of this framework that uses TabPFN as the base model. LoCalPFN breaks new ground for neural approaches on tabular data, even showing improvements over workhorse tree-based techniques. We also provide TabPFN-kNN as a variant without fine-tuning, demonstrating its superiority over the base model and practical utility.However, despite its successes, our framework also has some limitations. The first is that we have only shown that retrieval and fine-tuning improve TabPFN, since it is the only proven ICL-based tabular model. Thus, we cannot be certain that our ideas would directly transfer to new base models, although the success of these concepts in other domains provides some evidence. It is also worth noting that the original RAG paper (Lewis et al., 2020) only initially demonstrated success on BART. Next, the reliance on TabPFN as a base model brings some limitations: besides the constraints on number of features and classes discussed in Section 4.1, we are also unable to easily test our ideas in regression tasks since TabPFN is not designed for them. Although we expect these constraints to gradually be lifted as tabular foundation models improve and increase their scope, we also note that tree-based methods are not nearly as susceptible to these issues. Going further on the comparison with tree-based methods, while we note that LoCalPFN performs better than them in our experimental study, we also point out in Appendix A.5.6 that the runtime of LoCalPFN is slower. Yet it is still faster than other deep learning approaches, and the cheaper TabPFN-kNN variant runs as fast as any tree-based method on datasets we studied, while still attaining respectable performance. Overall, we believe that the benefits of our framework far outweigh the limitations, as LoCalPFN greatly expands the capabilities of deep learning on tabular data. 
 
 Appendix: A.1. Datasets For different variations of TabPFN inference techniques, we conduct experiments directly using the TabPFN repository https://github.com/automl/TabPFN.For all TabPFN-kNN experiments, we use a fixed number of neighbours equal to the minimum of (i) 10 times the square root of the dataset size, and (ii) 1000. We find this works well in practice since it adapts to small and large datasets. In practice, we use a batch size of 512 for inference using the faiss library for speedup.For LoCalPFN experiments, we use the exact same setup as TabPFN-kNN during inference. Therefore, at step 0, LoCalPFN and TabPFN-kNN are equivalent. For training LoCalPFN, we adopt the AdamW (Loshchilov & Hutter, 2019) optimizer with a learning rate of 0.01 and weight decay of 0.01. We do not have warmup or a learning rate scheduler. For the approximate local context for training, we use the same number of neighbours as TabPFN-kNN. We use a fixed number of query points (1,000) sampled from the training set and a batch of 2. For our reported results, we also use one-hot encoding for neighbour retrieval and inference. In addition, we evaluate our model every 30 gradient steps and apply early stopping based on the validation set AUC for each fold respectively.All experiments for our proposed methods can be run on a machine with a single NVIDIA RTX 6000 GPU Ada Generation, 995Gi RAM, and AMD Ryzen Threadripper PRO 5995WX 64-Cores CPU. Additional runtime analysis can be found in Figure 13 .In addition to tree-based models, we also compare LoCalPFN and TabPFN-kNN with deep learning based methods. We use the results directly from the TabZilla repository. However, due to the fact that a lot of the deep learning baselines are very computationally expensive, many of them were not able to run on all datasets. Therefore, we propose a subset of the 95 datasets which contains 71 datasets upon which all the deep learning methods could run. The details of the 71 dataset subset can be found in Table 4 .The complete results can be found in Table 5 . We can see that LoCalPFN still outperforms all other models.Here we also compare the performance of LoCalPFN with other models using accuracy and F1 score as the metric. We can observe a similar pattern here: LoCalPFN either matches or outperforms other models on either of these metrics as well. training loss for TabPFN-kNN (crosses), TabPFN (circles) for different dataset sizes and context/number of neighbours used on four datasets. We observe generally that for low number of neighbours (dark crosses) and especially for small datasets (small crosses) there is significant overfitting (higher test loss than train loss). TabPFN tends to overfit less, especially on larger datasets, which is expected. Overall, using TapPFN-kNN results in better underfitting/overfitting trade-offs where we obtain both lower test and train losses, however the gap between them increases.Figure 11 and Figure 12 show summaries of ablations on only the small datasets, and only the medium/large datasets, respectively. In the remainder of this subsection we see tables that show even further detail on the results presented in the main text. Retrieval & Fine-Tuning for In-Context Tabular Models 


Decision: The paper provides a retrieval and finetuning method that stacks on top of the current best tabular data in-context learner: TabPFN. They show consistent improvements compared to a TabPFN + k-nn baseline. The reviewers agree the experimental design is solid and the paper is thorough. The authors also provided a solid rebuttal to all the comments provided by the reviewers. There was one very negative reviewer who did not respond and gave a low score -- aside from this reviewer the average was acrtually high (5,7,6). For this reasons + having read the paper myself and it seems experimentally solid, I vote to accept this paper.

Reviews: review: paper_summary: paper_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
strength_weakness: strength_weakness: 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: 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?
limitations: 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).
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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: 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?
limitations: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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).
limitations: limitations: The authors have adequately addressed the limitation of their work.
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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: 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.
limitations: limitations: NA

Comments: title: Request for Comment
comment: 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)!
Thank you for the very clear answer! 

> Have you received the anonymized code link from the AC?

No
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.
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?
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!
Thank the authors for the feedback. I am still at the borderline for its improvement wrt TabPFN. Happy to hear what other reviewers think
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.