TuneTables: Context Optimization for Scalable Prior-Data Fitted Networks

Authors: benjamin  feuer
robin tibor schirrmeister
valeriia  cherepanova
chinmay  hegde
frank  hutter
micah  goldblum
niv  cohen
colin  white
Venue: NeurIPS_cc_2024_Conference
Published: 2024-04-23 02:48:24

Abstract: While tabular classification has traditionally relied on from-scratch training, a recent breakthrough called prior-data fitted networks (PFNs) challenges this approach. Similar to large language models, PFNs make use of pretraining and in-context learning to achieve strong performance on new tasks in a single forward pass. However, current PFNs have limitations that prohibit their widespread adoption. Notably, TabPFN achieves very strong performance on small tabular datasets but is not designed to make predictions for datasets of size larger than 1000. In this work, we overcome these limitations and substantially improve the performance of PFNs via context optimization. We introduce TuneTables, a parameter-efficient fine-tuning strategy for PFNs that compresses large datasets into a smaller learned context. We conduct extensive experiments on nineteen algorithms over 98 datasets and find that TuneTables achieves the best performance on average, outperforming boosted trees such as CatBoost, while optimizing fewer than 5\% of TabPFN's parameters. Furthermore, we show that TuneTables can be used as an interpretability tool and can even be used to mitigate biases by optimizing a fairness objective.

Main paper: 
 introduction: 
 Tabular data, or data organized into rows and columns consisting of distinct features, are the oldest and one of the most ubiquitous types of data in machine learning in practice [10, 67] . Tabular data has numerous applications across medicine [41, 71] , online advertising [33, 52, 64] , finance [7, 18] , and other areas [11, 12, 72] .Competitive classification algorithms for tabular data include gradient-boosted decision trees [16, 61] and deep neural networks [31, 42, 69] . Both approaches fit their respective models on a labeled dataset containing samples from a distribution reflecting the task at hand. A recent breakthrough, prior-data fitted networks (PFNs) [35, 55] , are are a specific type of neural process which learn to perform approximate Bayesian inference in a single forward pass using in-context learning [50] . PFNs do not require optimizing parameters or fitting a model on downstream training data, instead feeding training data into the context and conditioning on it. In particular, TabPFN achieved state-of-the-art classification on small tabular datasets [35, 51] .The in-context learning approach of PFNs parallels that of large language models (LLMs) [82] . Both approaches can be viewed as approximate Bayesian inference, whether implicitly [79] or explicitly [55] . While researchers have successfully used various context optimization strategies for enhancing LLM performance [49] , no prior work has studied context optimization strategies for PFNs. Furthermore, although TabPFN achieves very strong performance on small datasets, its limitations TabPFN TuneTables 1.Pretrain on synthetic data 2.Predict D test using D train (size <1000), as context 1.Pretrain on synthetic data 2.Prompt tune on D train 3.Predict Figure 1 : TuneTables: a novel prompt-tuning technique for prior-data fitted networks. TuneTables performs prompt tuning on a pre-trained prior-fitted network (TabPFN) to distill real-world datasets into learned embeddings, allowing for stronger performance and faster inference time than TabPFN in many cases. TuneTables also expands the capabilities of pre-trained PFNs; by way of example, we demonstrate its effectiveness for bias mitigation, and as an interpretability tool.currently prohibit its widespread adoption: it only runs on datasets whose number of training samples, number of features, and number of classes are at most 1000, 100, and 10, respectively.In this work, we perform the first investigation into context optimization strategies for PFNs, allowing us to substantially improve their performance when scaled to large datasets. Specifically, we introduce TuneTables, a novel parameter-efficient fine-tuning technique for PFNs that compresses large datasets into a smaller learned context (Figure 1 ). We conduct an extensive empirical investigation into TuneTables' performance on the TabZilla Benchmark Suite, the largest benchmark considered by recent tabular data literature [51] . Over 98 datasets and 19 algorithms, we find that TuneTables achieves the best average performance and is the best-performing method on 30 of them.Because TuneTables effectively compresses the contents of large tabular datasets into the tuned prompt, no training data is needed in the context during inference, significantly speeding up inference time (similar to works on neural processes; see Section 6) . We also show that the learned prompt can be used as a tool for interpretability. Finally, we show how to use TuneTables for multi-objective optimization, such as optimizing both accuracy and fairness simultaneously, allowing users to mitigate biased predictions of a pretrained PFNs with just a lightweight tuning procedure. We open-source our code and raw results at https://github.com/penfever/TuneTables.Our contributions. We describe our main contributions below.• We introduce TuneTables, a parameter-efficient fine-tuning technique for PFNs. TuneTables achieves the highest number of wins (30) when compared to 19 algorithms over 98 datasets, while requiring less inference time than TabPFN. • We show how to use prompt tuning for multi-objective optimization, such as optimizing both accuracy and fairness simultaneously, allowing users to mitigate biased predictions of a pretrained PFNs with just a lightweight tuning procedure. • We show that TuneTables' condensed representations can be used as an interpretability tool.• We conduct an extensive study on context optimization strategies for PFNs by performing an ablation study on TuneTables, as well as studying sketching and feature selection techniques. • In order to better manage the limitations of our method, we introduce TuneTables-medium and TuneTables-light, which achieve strong tradeoffs between precision and speed.1. PFNs only accept a fixed number of features. The current design of PFNs fixes the quantity of features at the time of pretraining. This quantity cannot be changed without retraining the PFN. 2. PFNs scale poorly with the dataset size. While PFN accuracy can improve with more real-data samples at inference time [35, 58] , the memory requirements scale with the context length, making extensions beyond a certain number of samples impractical. 3. PFNs only select from a fixed number of classes. The MLP decoder head co-trained with the PFN fixes the number of classes that can be identified at test time.A motivating example. In Figure 3 (left), we present a comparison between CatBoost and TabPFNs3000 from [51] , a version of TabPFN that, for datasets above 3000 data points, uses a random subset of 3000 data points, for fitting; this version also evaluates 30 feature subsets based on mutual information (requiring 30× more time). We ablate our feature and sample subselection strategies for TabPFNs3000 in Appendix Table 13 . Although TabPFNs3000 performs very well on datasets with fewer than 1000 datapoints and 100 features, it significantly underperforms CatBoost beyond those constraints.Approach. We propose sketching, feature selection, and fine-tuning as an attempt to remedy (1) and (2) . Then in Section 3, we describe novel prompt-tuning and class-extension strategies to create TuneTables, a robust classifier which remedies (1)-(3). 
classical sketching and feature selection: 
 Sketching. The number of samples that a PFN can handle is limited to around 3000 by conventional GPU sizes. However, in the real world, datasets are often much larger. Going forward, we refer the maximum allowable context size of the PFN as n, and to the size of the real-world dataset as N .Given a training dataset D train of size N >> n, one option is to select a representative subset of the dataset, D compact ⊆ D train , to use as the context. In general, researchers have studied a variety of data summarization techniques, often called sketching, for tabular data [56] . In the context of a pretrained PFN q θ , the goal is to find a sketching function s : R N ×d → R n×d such that E D∼p(D) [-log q θ (y test | x test , s(D train ))] (4) ≈ E D∼p(D) [-log q θ (y test | x test , D train )] .Here, we consider three sketching methods for TabPFN: random, in which we select a random subset of n datapoints from the full training set; k-means, in which we compute the n-means clustering [3] of D train and select the n centers; and CoreSet, in which we compute a core-set of size n [2] .Feature selection. In addition to the limitation on dataset size, PFNs, such as TabPFN, also impose a limitation on the number of features d. Similar to sketching, given a dataset with D >> d features, we can perform feature selection or summarization in order to adhere to the constraint (formally, finding a function that reduces the feature dimension and approximates an expression similar to Equation ( 4)). Feature selection is a critical part of tabular classification, and there are several popular feature selection methods [13, 17] . We investigate three different methods: random, in which we randomly select a set of d features; mutual information, in which we select d features with the highest mutual information of the target dataset [75] ; and principal component analysis (PCA), in which we take the d first principal components. In Section 4, we find that all sketching and feature selection methods plateau in performance well before approaching parity with GBDTs, which motivates our investigation of new scaling techniques.We conclude this section with a discussion of another potential approach for scaling PFNs: finetuning. In this approach, given a dataset of size N >> n, we use gradient descent to continue training all parameters of the PFN. However, we show in Section 4 that fine-tuning takes up considerable memory resources while still not achieving competitive accuracy on large datasets; we postulate that for most datasets, the synthetic prior of TabPFN is more robust to overfitting than the actual training data; fine-tuning overfits to the validation set. This is borne out by the observation that when fine-tuning TabPFN, validation accuracy continues to improve even as test accuracy declines. To remedy this, we present new parameter-efficient solutions for scaling PFNs in the next section. 
tunetables: 
 Motivated by the strong performance of prompt tuning, we introduce TuneTables, a new tabular classification algorithm. Using a pretrained PFN (TabPFN, in our experiments) as a base model, TuneTables overcomes the limitations of PFNs, allowing them to be applied to any tabular classification problem.For the full details of the algorithm, please refer to Appendix D. We summarize here: (a) if a dataset is small enough to run with the original zero-shot version of TabPFN, we include it in our search, as the TabPFN is already highly optimized for such datasets; The reason why we do not only use TabPFN is that the exact size at which TuneTables outperforms TabPFN is dataset-dependent (with an average transition of around 800 samples). (b) if there are more than 100 features, we perform grid search over a set of feature subselection methods and select the one which performs best on the validation set; (c) orthogonally, if there are too many labels, we fit a new decoder to a frozen TabPFN; (d) we optimize over a search space of tuned prompts, both with and without real-data context during training, fitted to the dataset; (e) we report the best-performing model according to accuracy.Prompt tuning as a scalable context for PFNs. Motivated by the limitations of sketching for large contexts, we explore soft prompt tuning as an alternative. In soft prompt tuning [46] , given a dataset D = D train ∪ D test , a parameterized matrix D p×e tune is prepended to the input embedding D n×e train , where e is the transformer embedding dimension and p is a hyperparameter-the size of the tuned prompt.The paper [46] demonstrates the effectiveness of soft prompt tuning for NLP tasks by prepending a small task-specific prompt (p ≈ 5). These task-specific learned tokens prove effective at the extremely large model scales commonly encountered in NLP. Somewhat surprisingly, we show that prompt tuning is effective at similar scales of p even for the much smaller tabular data models (see Appendix D and Appendix F). However, prompt tuning increases in effectiveness when p is larger.Figure 2 : TuneTables and state-of-the-art tabular models. A critical difference plot according to mean accuracy rank across the 98 datasets in Table 1 of [51] . Algorithms which are not significantly different (p > 0.05) are connected with a horizontal black bar. TuneTables achieves the highest mean rank of any algorithm.Soft prompt tuning for tabular data. Unlike in NLP, transformers for tabular data, including PFNs, generally accept two input embeddings; a continuous-valued embedding D trainX , and a categoricallyvalued D trainy which is passed through directly. We adapt the method of [46] by fitting the parameters of D p×e tune X to D trainX and randomly initializing D p×1 tune y with an equal number of labels from each class in D train . These synthetic datapoints are optimized on the entire labeled set, and therefore give PFNs access to a much larger training set not accessible by existing methods.We further adjust the method of [46] to allow for the possibility that D tune has learned, in essence, a distilled version of D train ; at test time, we evaluate two settings, hereafter referred to as C ('context') and N C ('no context'). In the C setting, following [46] , we have D p×e tune prepended to the input embedding D n×e train . In the N C setting, we provide only D p×e tune at test time. In Section 4, we empirically evaluate both approaches. We ablate specifics of our implementation choices in Appendix D. Unlike prompt tuning for NLP, we show that the N C setting is often competitive with, if not better than, the C setting. We also ablate this choice during training; see Appendix D for implementation specifics.Extending the number of predicted classes. TabPFN uses a pretrained transformer, with a two-layer MLP as a final classifier. The pretraining procedures limit the naïve use of TabPFN to classification tasks with at most 10 classes. Yet, in many cases, datasets of interest might contain a larger number of classes, which would require pretraining a new PFN from scratch.Following the method of last-layer retraining [44, 45] , for such datasets, we fit a PFN to new posteriors given by real-world classification datasets with more than 10 classes, freezing all weights except for those of the decoder MLP and the tuned prompt (see Appendix D for more implementation details). Even after this modification, our method remains highly parameter-efficient, optimizing fewer than 5% of TabPFN's parameters. 
experiments: 
 Algorithms and datasets used. We compare TuneTables to nineteen algorithms, including three GBDTs: CatBoost [61] , LightGBM [43] , and XGBoost [16] ; 11 neural networks: DANet [14] , FT-Transformer [31] , two MLPs [31] , NODE [60] , ResNet [31] , SAINT [69] , STG [80] , TabNet [6] , TabPFN [35] , VIME [81] , and ExcelFormer [15] ; and five baselines: Decision Tree [62] , KNN [21] , Logistic Regression [22] , Random Forest [48] , and SVM [20] . For all algorithms, we use their official implementation; see Appendix C for more details. We also compare to TabPFNs3000. We run the algorithms on the TabZilla Benchmark Suite introduced in [51] . This suite consists of 98 classification datasets from OpenML [74] with a diversity of sizes and number of features [51] . See Table 4 in Appendix C for a list of all datasets with their statistics.Experimental setup. For all algorithms other than TuneTables, we perform light hyperparameter tuning by running one default setting and 29 iterations of random search using Optuna [4] ); see Appendix C for details. Following [51] , all of the algorithms come with their default set of hyperparameters used in the official implementation, and we used all of these settings. For TuneTables, we optimize via a grid search described in Appendix D. We fit for up to 100 epochs with early stopping. For all algorithms, we report the test performance of the hyperparameter set with the best performance on the validation set and cross-validate on three train/test folds from OpenML. We conduct our experiments on an NVIDIA L4 TPU with 24GB VRAM. We summarize our results across datasets by reporting mean accuracy as well as the mean normalized accuracy (after Z-score normalization) for each algorithm.Algorithm comparison. We compute statistically significant performance differences among algorithms averaged across all datasets, as done in prior work [42] . We first use a Friedman test to check whether performance differences between all algorithms are significant [30] . We reject the null hypothesis for p < 0.05. Then we use a Wilcoxon signed-rank test to check which pairs of algorithms have significant performance differences [19] . We use a Holm-Bonferroni correction to account for multiple comparisons [36] . See Figure 2 . We find that TuneTables achieves the highest average rank, although the difference among the top three algorithms is not statistically significant. In Table 1 , we present the accuracy, rank, and Z-score of all algorithms, averaged across all datasets, and find that TuneTables performs very well on all metrics but runtime.Large datasets. One limitation of the TabZilla Benchmark Suite is that even the largest dataset in the comparison contains only 45,211 samples. This scale is modest by the standards of modern tabular problems. In order to better understand the performance of TuneTables on extremely large datasets, we curate from [51] a novel benchmark, LARGESCALETABLES , consisting of 29 datasets with up to 1 900 000 samples, and 7200 features. We curate the datasets from OpenML, omitting image classification datasets. Since TabPFN generally outperforms boosted trees on these smaller datasets, and TuneTables extends TabPFN, we also heuristically select smaller datasets to LARGESCALETABLES so as not to favor either TabPFN or boosted trees. TuneTables achieves the highest average accuracy of any algorithm, and achieves the best performance on poker-hand, a dataset of size 1 025 009. The complete results can be found in Appendix Table 8 and Figure 3 (right).In order to assess the performance of TuneTables on datasets with many classes, we curate another subset of [51] , presenting results on fifteen datasets with up to 100 classes. In appendix Table 7 we show that despite the large divergence from the PFN's pretraining, TuneTables achieves a mean rank of 2.52, ahead of CatBoost and a ResNet, second only to XGBoost, whose mean rank is 2.0.Runtime comparison. We divide our consideration of runtime into inference and training. At inference time, TuneTables is around 9x faster than TabPFNs3000, on average; see Appendix Table 10 for the per-dataset details. However, the end-to-end runtime of TuneTables is over 7x that of CatBoost and XGBoost, and also slower than TabPFNs3000, because of the increased training time.In order to better understand the trade-off between accuracy and runtime, we introduce efficient variants of our method. TuneTables-medium utilizes a more efficient adaptive sequence size (i.e., the number of real data points received at train time) which scales with the size of the dataset, validates Table 2 : TuneTables-medium and TuneTables-light are substantially faster with only a modest decrease in accuracy. We compare the average accuracy and runtime (in seconds) of three versions of TuneTables and find that the medium and light versions of the algorithm are substantially faster on large datasets; we also find that TuneTables-medium sacrifices little accuracy. on a subset of the available validation set when the validation set is large, employs lower patience for early stopping, omits a zero-shot TabPFN grid search over 30 random seeds which TuneTablesstandard uses to find more optimal feature selection subsets, and, most impactfully, omits ensembling for datasets larger than a cutoff hyperparameter (which we fix at 150 000 samples). TuneTables-light includes all of the efficiency modifications of TuneTables-medium, trains for just one epoch, and uses TabPFN zero-shot to preselect the feature selection method rather than performing a grid search using TuneTables itself. In Table 2 , we compare these lighter methods to TuneTables-standard on the LARGESCALETABLES benchmark. TuneTables-medium decreases runtime by 72% compared to TuneTables-standard, while the accuracy decreases by less than a quarter of a percent. Furthermore, the runtime of TuneTables-light is comparable to CatBoost; however, performance also degrades by 5% when going from TuneTables-medium to TuneTables-light. Still, TuneTables-light shows stronger performance than TabPFNs3000 (78.7% vs. 78.1% accuracy) while having a lower inference time.Ablations. In order to better understand the significance of the changes we introduce in our method, we separately ablate the tuned prompt size and the use of ensembles in Table 12 and Table 13 , finding that smaller datasets pair well with smaller prompts and rarely benefit from ensembling.We also compare TuneTables with and without keeping the real data as additional context (referred to as C and NC, respectively, in Section 3); see Appendix Table 12 . We find that smaller datasets cannot always be fully learned by the tuned prompt, but for larger datasets, the tuned prompt alone suffices, and in some cases even outperforms the real data. Sketching and feature selection Finally, in Appendix Table 6 we give a study on the three sketching and three feature selection techniques described in Section 2. As described earlier, TabPFNs3000's performance when relying on feature selection plateaus well before approaching parity with CatBoost, a top-performing GBDT, on seven very large datasets. 
tunetables extensions: 
 Mitigating bias with prompt tuning. Many real-world applications of machine learning involve a set of protected attributes (such as race or gender) that partition the dataset into groups, in which some have higher model performance than others. Removing the sensitive attributes does not fix the algorithmic bias, because the sensitive attributes are often non-trivially correlated with other attributes in the dataset. Due to this issue, researchers have put in significant effort into mitigating the bias of ML models, with the majority of techniques devising new training strategies [9, 53] . Figure 4 : Dataset with high accuracies from just two datapoints. Shown is a two-example prompt dataset for the breast cancer dataset [78] . Malign class example has higher values for all features than benign class. Given TabPFN's pretrained nature and considerable retraining cost, the only options for mitigating biased predictions are to run a postprocessing routine on the output predictions, which generally do not perform as well as inprocessing strategies [66] . We show how to use prompt tuning to substantially reduce the bias of predictions while also improving accuracy.We conduct experiments on four datasets widely used for research in fairness: the Adult Census Income database (with sex as the sensitive attribute) [8] , speed dating (with same race as the sensitive attribute) [83] , COMPAS (with sex as the sensitive attribute) [5] , and National Longitudinal Survey (with gender as the sensitive attribute) [73] . To quantify bias, we use demographic parity [25, 76] , which measures the difference in probability of a positive outcome among the protected and unprotected groups. Formally, given protected group G 1 , unprotected group G 0 , and protected attribute x •,a , it can be computed asP (xi,yi)∈G0 (y i = 1 | x i,a ) -P (xi,yi)∈G1 (y i = 1 | x i,a ).During prompt tuning, we employ a demographic parity regularizer that aims to minimize the difference in positive outcome probabilities between the two groups:(xi,yi)∈G0 P (y i = 1 | x i,a ) - (xi,yi)∈G1 P (y i = 1 | x i,a )We use the same experimental setup as in Section 4, except that we report one shift rather than the average of three, and TuneTables is fitted to a single prompt rather than ensembled. We compare the default TabPFN to TuneTables, fine-tuning for accuracy alone vs. accuracy and demographic parity. The latter significantly improves demographic parity, compared to the default TabPFN and TuneTables fine-tuned for accuracy, and enhances accuracy relative to the default TabPFN across all datasets but one; see Table 3 .Summarizing and understanding datasets with prompt tuning. While we have demonstrated that TuneTables scales and improves the performance of PFNs, now we show that it can also help understand the discriminative features in a dataset. Often, besides a good predictive model for a dataset, users want to gain further insights into the dataset. In prompt tuning, the tuned smaller dataset can be seen as a summary of the complete dataset that emphasizes discriminative features for the given task. As an example, in Figure 4 we show that on the Breast Cancer dataset [78] , a prompt with just two synthetic examples is enough to reach high accuracies and at the same time allows an understanding of the predictive features. For example, in the Breast Cancer dataset, the malign example has higher values for all features compared to the benign example, suggesting high feature values indicate malignancy. We give further examples in Appendix F. 
related work: 
 Neural Processes and Prior-Data Fitted Networks. Prior-data Fitted Networks (PFNs) [35, 55] are a recently-proposed paradigm for machine learning, which show that fast approximate Bayesian inference is possible by training a neural network to mimic the posterior predictive distribution (PPD) in a single forward pass using in-context learning [54, 55, 58] . PFNs were shown to yield state-of-the-art empirical performance on small tabular datasets [35, 51] . PFNs have been used in other applications, including Bayesian optimization [54] learning curve extrapolation [1] , and as foundation models for hypernetworks [57] .PFNs are neural processes (NPs) [55] . Recent advances in NP are similar in nature to our work, especially recent works that aim to scale attention-based methods to larger contexts. Feng et al. [27] propose Latent Bottlenecked Attentive Neural Processes (LBANPs), a transformer-based neural process which overcomes the quadratic complexity of transformers by encoding the context dataset into a constant number of latent vectors. Guo et al. [34] propose Versatile Neural Processes (VNP), which increases the capability of NPs to handle compex signals by using a new bottleneck encoder. [63] introduces semi-parametric inducing point networks (SPIN), which can attend to a training set at inference time with linear complexity via inducing point methods. [28] introduce Constant Memory Attention Block (CMAB), an attention block that is permutation-invariant and has constant memory complexity when computing the output, as well as Constant Memory Attentive Neural Processes (CMANPs), a NP that requires constant memory. For additional related work, see Appendix B. 
conclusions, limitations, and future work: 
 In this work, we gave the first investigation into context optimization techniques for PFNs, allowing us to substantially improve their performance when scaled to large datasets. In particular, we introduced TuneTables, which uses a novel prompt-tuning technique to achieve strong performance on large datasets. We demonstrate that TuneTables mitigates the constraints of TabPFN on the dataset size, number of features, and number of class labels. Additionally, we use prompt tuning to mitigate bias without retraining TabPFN and as an interpretability tool. We open-source our code, results, and all other materials needed to reproduce our work. As PFN models improve, the context optimization techniques explored in our work will allow researchers to further optimize and scale. For example, a next-generation TabPFN might have a longer total context length, and prompt tuning will allow us to push the dataset size even further.Limitations. No current TuneTables is on par with GBDTs in terms of both accuracy and runtime simultaneously. We therefore emphasize that while we achieve strong performance metrics, we do not claim practical superiority of our method over gradient boosting, when taking into account training time. However, given the novelty of our method, we expect future research to further improve the accuracy-runtime tradeoff of TuneTables. It also does not improve on TabPFN for small datasets (fewer than 1000 samples, 100 features and 10 classes); we postulate that this is a result of overfitting.Future work. Parameter-efficient fine-tuning can be used with PFNs to ensure high-quality, differentially private predictions, given that the pretraining is done on purely synthetic data, and prompt tuning only updates a small number of parameters. Low-rank adaptation (LoRA) [37] and quantized low-rank adaptation (QLoRA) [23] have been used successfully for large language models and would be a promising technique for parameter-efficient fine-tuning of PFNs. Designing a sparse mixture-of-experts PFN using router networks is another promising technique, due to its success in the field of large language models [39] .Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments? Answer: [Yes] Justification: Our results in Section 4 do include standard deviations, and we also include critical difference plots in order to compute the statistical significance of our results. Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)• The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 10. Broader Impacts Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? Answer: [Yes] Justification: We discuss the societal impacts in Appendix A. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA]Justification: As discussed in Appendix A, our paper does not pose such risks.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: We credit the original owners of the resources we use in the supplementary material.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. 
 
 Appendix: The goal of our work is to investigate context optimization strategies for PFNs, such as prompt-tuning and sketching, which we use to scale and improve the performance of TabPFN. We do not see any negative broader societal impacts of our work that do not already exist in other classification methods. In fact, our work may further facilitate the adoption of TabPFN, which has the benefit of being pretrained and therefore having a lower carbon footprint compared to most deep learning approaches that must be trained from scratch. Furthermore, we demonstrate that our prompt-tuning strategy makes it possible to mitigate the bias of TabPFN while only fine-tuning a small set of embedding vectors; the authors of TabPFN mentioned that it is critical to study and improve TabPFN under the lens of algorithmic fairness and other dimensions of trustworthy AI [35] . This may allow practitioners to use TabPFN in sensitive settings for the first time, in which the original TabPFN would lead to biased preditions. Overall, our hope is that our work will have a positive impact for both practitioners and researchers, by facilitating the adoption of a model with a low carbon footprint, and by providing the tools to mitigate the bias of PFNs. Likewise, we hope that the possibility to summarize datasets with prompt tuning will add to the toolbox of machine learning practitioners aiming to analyze and interpret their data better, and therefore may have a positive impact on the trustworthieness of machine learning.Tabular classification. Tabular datasets are the oldest and among the most widely used dataset types in machine learning [10, 67] . GBDTs [29] build an ensemble of decision trees, with each tree fitting the residual of the loss from the previous tree. XGBoost [16] and CatBoost [61] are two of the most widely-used and highest-performing GBDTs. Researchers have also explored many methods based on neural nets [31, 38, 42] .There is an active debate in the community on which family of methods perform best on tabular data: neural nets [6, 42, 47, 60, 65] or GBDTs [10, 31, 32, 67] , with the exception of small datasets, on which TabPFN performs the best [35, 51] . Finally, sketching and feature selection methods have been extensively studied in prior works [13, 56] .Neural Processes and Prior-Data Fitted Networks. Prior-data Fitted Networks (PFNs) [35, 55] are a recently-proposed paradigm for machine learning, which show that fast approximate Bayesian inference is possible by training a neural network to mimic the posterior predictive distribution (PPD) in a single forward pass using in-context learning [54, 55, 58] . PFNs were shown to yield state-of-the-art empirical performance on small tabular datasets [35, 51] . PFNs have been used in other applications, including Bayesian optimization [54] , forecasting [24] , and learning curve extrapolation [1] .PFNs are neural processes (NPs) [55] . Recent advances in NP are similar in nature to our work, especially recent works that aim to scale attention-based methods to larger contexts. Feng et al. [27] propose Latent Bottlenecked Attentive Neural Processes (LBANPs), a transformer-based neural process which overcomes the quadratic complexity of transformers by encoding the context dataset into a constant number of latent vectors. Guo et al. [34] propose Versatile Neural Processes (VNP), which increases the capability of NPs to handle compex signals by using a new bottleneck encoder. [63] introduces semi-parametric inducing point networks (SPIN), which can attend to a training set at inference time with linear complexity via inducing point methods. [28] introduce Constant Memory Attention Block (CMAB), an attention block that is permutation-invariant and has constant memory complexity when computing the output, as well as Constant Memory Attentive Neural Processes (CMANPs), a NP that requires constant memory.For additional related work, see Appendix B.Prompt tuning. Researchers have extensively studied prompting techniques for large language models (LLMs) [49] . In such methods, one modifies the input into a prompt, and the LLM predicts the subsequent tokens, from which it derives the output predictions. Prompt tuning techniques typically start with a pretrained LLM and use the training data of a downstream task to find the best prompt for this task. 'Hard' prompt tuning involves finding the best natural language prompts using discrete search techniques [40, 77] , while 'soft' prompt tuning optimizes a prompt directly in the embedding space of the model [46, 84] . Soft prompt tuning has also been applied to multi-modal models such as vision-language models [70, 85, 86] .Hyperparameter settings for baselines. For the baselines, we use the same hyperparameter search space as in prior work [51] (which is itself similar to other works [42] ). We choose the best configuration among the default parameter sets and the additional 29 sets, using k-fold crossvalidation with 10 folds. We follow the best configurations of [51] with exceptions; for CatBoost and XGBoost, we increase the range of tree depth by a factor of 10, since a high depth is often more appropriate for very large datasets. We find that this improves performance on large datasets, but sometimes harms performance on small datasets, and increases the runtime -see Appendix G for the raw results. For each dataset, each algorithm is allowed a maximum runtime of 10 hours for the entire search process. For ExcelFormer, we decrease the number of dimensions and depth of the default transformer, as we find that without this modification, it times out on even medium-sized datasets.Hyperparameter settings for TuneTables. Consistent with the experimental design in [51] , we perform hyperparameter optimization of TuneTables for up to 30 variants; however, we utilize a grid search rather than Optuna. We use a batch size of 1 and no gradient aggregation. The space over which we conduct grid search in TuneTables is conditioned on metadata; number of features, number of samples, number of classes. The feature and class splits are set by the limits of the particular TabPFN checkpoint we optimize. When we reach a leaf node, TuneTables-standard and TuneTables-medium conduct a grid search over a fixed range of configurations. The size of our search space is always < 30, and usually < 10. For feature-large datasets, TuneTables-light conducts additional optimization in the feature space prior to grid search. Further details on our search algorithm can be found in the repository associated with the paper. TuneTables hyperparameter settings can be found in chart form in Table 5 .In our experiments, we find that we exceed hardware limitations around 3000 samples with a batch size of 1 when using TabPFN on a GPU with 48GB of VRAM. While FlashAttention, a popular optimization, is memory-linear in sequence length, but it has not yet been implemented in the public release of TabPFN, which relies on an older version of PyTorch. Flash attention may be integrated into a future release, in which case any benefits will carry over to our method. Even so, inference time continues to be quadratic, and there are other overheads on GPU memory. Overall, there remains a considerable need for new algorithmic methods which can be scaled up to the sequence lengths required by large tabular datasets.We ablate strategies for both sketching (subsampling) and feature selection for scaling TabPFN. We consider three sketching methods (random, k-means, and CoreSet) and three feature selection methods (random, PCA, mutual information) as described in Section 2. In addition to sketching, we consider two different methods for sampling class labels: equal and proportional.For efficient coreset selection, we use a variant of Farthest Point Sampling [26] ; after selecting an initial set of n=5 random points, we compute the distance of each point in the dataset to the set of already selected points. Next, we add to the set of selected points the point whose distance to the selected points is maximal. Finally, we update the distances of all the points according to the updated selected set; and continue iteratively.We limit our algorithmic comparison to TabPFN to CatBoost, which is the overall best-performing model in [51] . We compare all combinations of sketching and feature selection with both CatBoost and TabPFN; see Table 6 . Interestingly, we find that random sketching performs just as well as the more involved algorithms, k-means and CoreSet. On the other hand, PCA significantly outperforms mutual information and random feature selection methods, when the original number of features is large. First, we present the large class table from Section 4. See Table 7 .In Table 1 , we gave aggregate statistics for TuneTables and baselines across all datasets of size less than 50 000. Now, in Table 8 , we present results for TuneTables, CatBoost, and XGBoost on all datasets individually, including datasets of size nearly two million. Note that we report the end-to-end runtime of all methods, including 30 iterations of search (for example, TabPFNs3000 takes about 30× more time than TabPFN as reported in prior work [51] ). In Figure 6 , we present a critical difference diagram similar to Figure 2 but over all 29 datasets. Across these 29 datasets, TuneTables achieves performance that is not statistically different from CatBoost or XGBoost.In Table 1 and Table 8 , we compared TuneTables to GBDTs and two other neural nets. Now, we compare TuneTables to two additional neural nets: MLP and TabTransformer [38] (for four total neural net baselines). Since transformer-based methods have a higher memory consumption, we obtained full results on 17 total datasets. See Table 9 . We find that TuneTables substantially outperforms all other neural nets on average across all datasets, achieving the lowest average rank at 1.93, and achieving the highest average Z-Score of 0.85, far above the second-place neural net's value of 0.18.Inference Time We give the full details for inference time; see Table 10 .Full results for TuneTables-medium and TuneTables-light Recall at the end of Section 4, we introduced TuneTables-medium and TuneTables-light, showing that they can substantially decrease the runtime of the standard TuneTables with just a small reduction in accuracy. While we presented summary results in Table 2 , now we present the full results in Table 11 . For descriptions of TuneTables variants, see Section 4.We report our results ablating the core components of our methods: the prompt length, its use during training and inference, the use of ensembles, and the prompt tuning procedure itself (against regular finetuning). Beginning with the prompt length, we see in Table 12 that while some dataset results are not sensitive to the prompt length, others vary significantly, and generally enjoy longer prompts.Ablating the importance of having real-data context during inference (same table) we find that it is important for smaller datasets. On the larger datasets, the variant with no real data at all is better in specific cases. However, having no such data during inference, a model would perform better not having it during training as well. Introducing real data first during inference is usually harmless, but significantly deteriorates the results on specific cases. Fine tuning the entire model is not only significantly more memory-intensive but also underperforms TuneTables in terms of accuracy.In Table 13 we show the results of our method with or without ensembles, and with or without context (in training and inference). Our ensembling strategy is beneficial in both cases. When using ensembles, the results with or without context are generally similar, although one variant may outperform the other on specific datasets.Tuned prompt implementation details. We implement our tuned prompts by prepending randomly initialized vectors of identical dimension to the TabPFN encoder to the real data points fed to TabPFN as context during training.In the CT ('context in training') setting, we continue to provide real data points as part of the training context; the quantity of real data provided for each batch is drawn with uniform probability from a random variable which ranges from zero to a fixed upper bound, which is passed to the model as a hyperparameter. Usually that upper bound is 1152, the default setting for TabPFN, unless the dataset is extremely small, in which case we select it according to the limitations of that dataset.In the NCT ('no context in training') setting, no real data is provided as part of the training context; all data in the training set is used to fit the randomly initialized prompt. There are always a fixed number of data points provided, and it is the same for every batch. Usually that number is 128, unless the dataset is extremely small, in which case we select it according to the limitations of that dataset.We note that it is also possible to evaluate models trained in either setting either with context (C) or without (NC). In our experiments, we always evaluate all models on both, and report the setting with higher validation accuracy at test time.We find that on smaller datasets, the setting with context (C) sometimes outperforms the setting without context (NC), even after the prompt has been fitted, perhaps because the tuned prompt overfits on a small amount of available training data. However, when using ensembling, for datasets with more than 3000 samples in the training set, the NC setting is as good or better than the C setting. See Table 13 . We leave further investigation of this phenomenon to future work.Loss function. Most TuneTables experiments are optimized using the cross-entropy loss between the labels assigned to the training data (P) and TuneTables's outputs (Q):EQUATIONTo reduce the effects of overfitting on datasets with fewer samples, we extend our grid search to include the choice of loss function. We find that many such datasets benefit from fitting the tuned prompt via the KL divergence loss between the PFN's outputs (P) and TuneTables (Q).x P (x) logEQUATIONCompute. We conduct all of our prompt tuning experiments on Google Cloud Platform (GCP), using a single NVIDIA L4 TPU with 24GB VRAM for each experiment.Size of tuned prompts. We ablate the size of tuned prompts in Table 12 and find that the larger tuned prompts generally perform better on datasets with more samples, while smaller tuned prompts excel on datasets with few samples. For this reason, the experimental results in Section 4 are reported on prompts of size 10 or 1000, conditioned on the number of samples in the input dataset.Our experiments run for up to 100 epochs. We employ early stopping if our key metric (Accuracy) fails to improve after a fixed number of epochs.Evaluation of tuned prompts. We validate our tuned prompts every epoch on a subset of the entire validation set if the validation set is large. At a fixed interval, we run the entire validation and test sets. After the experiment concludes, we report results from the epoch with the best key metric score.Fine-tuned setting. In the fine-tuned setting, our parameters are the same, except that we use a lower fixed learning rate of 1e -3.Ensembling over tuned prompts.While individual tuned prompts are already a substantial improvement on TabPFN for sample-large datasets, we find that these improvements sometimes compound when we ensemble over multiple tuned prompts.We draw the inspiration for our ensembling approach from the internal ensembling used in TabPFN, which averages predictions over permutations of feature indices and label indices [35] . For the results presented in Section 4, we ensemble over ten permutations of each dataset, averaging the top two ensemble member predictions (as measured by NC accuracy on the validation set in the NCT setting, or C accuracy in the CT setting).In a TuneTables ensemble, each ensemble member fits its own tuned prompt to the data. Variance in the ensemble members is introduced by differences in the random initialization of the tuned prompt, as well as permuting the order of features and labels, a la TabPFN, one time before each tuned prompt is fitted. See Table 13 for ablation studies of the effectiveness of ensembling.In this section, we show that using prompt tuning, we are also able to extend PFNs to perform well on regression datasets, further highlighting the flexibility of our approach.Adapting TuneTables to regression problems introduces new challenges, as TabPFN was not designed for use with such problems. In particular, we find that end-to-end fine-tuning outperforms prompt tuning in this setting, and that the most effective grid search for regression problems utilizes a space of PFN base models. Put another way, TuneTables-regression does not use prompt tuning, and does not always use TabPFN as its foundation model. Rather, it searches a space of foundation models to find the best performer for a particular regression problem. We search over three foundation models; TabPFN, the checkpoint released in [55] and a new PFN we train from scratch for 10 epochs on a synthetic prior of regression-specific datasets.Similar to [51] , each algorithm is tuned for each dataset by maximizing the R-squared (R2) metric. Each dataset corresponds to an OpenML task, and can be preprocessed exactly like the classification datasets used in other experiments. The 15 datasets used in these experiments are "Bank-Note-Authentication-UCI" (OpenML task 361002), "EgyptianSkulls" (5040), "Wine" (190420), "Wisconsin-breast-cancer-cytology-features" (361003), "bodyfat" (5514), "california" (361089), "chscase-foot" (5012), "cleveland" (2285), "colleges" (359942), "cpu-small" (4883), "liver-disorders" (52948), "meta" (4729), "mv" (4774), "pbc" (4850), and "veteran" (4828). Unlike prior work, we report non-normalized averages, as we believe it gives a more realistic indication of real-world performance at a glance.We use 12 deep learning and boosted tree algorithms as baselines; VIME, TabTransformer, TabNet, DANet, STG, MLP, LightGBM, CatBoost, XGBoost, NODE, DeepFM. TabPFN Zero-Shot Decision Boundary In Table 14 results and Figure 7 , we show the results of our experiments. Overall, we see that TuneTables is competitive with top-performing algorithms across a variety of metrics; mean win rank, average R 2 score, etc. Equally importantly, TuneTables is significantly stronger on several datasets than both boosted trees and deep baselines.For the experiments with prompts of only two examples, we prompt-tune one example per class for 500 epochs. The prompt is input directly into TabPFN without undergoing any preprocessing steps that the original data went through. Therefore, the tuned feature values do not correspond directly to the original values in the dataset. For the breast cancer dataset, it reaches 100% accuracy; for the diabetes dataset, it reaches 78.7% accuracy.We also show decision boundaries of TabPFN and prompt-tuned TabPFN on toy 2d classification problems using scikit-learn [59] in Fig. 8 . The prompt contains four datapoints for each class. One can see how the points are tuned to recover the decision boundaries. Due to the very low prompt length and the low dataset size, the prompt-tuned decision boundaries are slightly less accurate than TabPFN (consistent with our results for small datasets in Table 12 ).In the results in Section 4, we compare TuneTables to the GBDTs, XGBoost and CatBoost, using a search space in which the number of trees can range from 1 to 1000. [68] with high accuracies from just two samples.range from 1 to 10 000. We re-compute the critical difference plot and table results: see Figure 5 , Table 15 , and Table 16 .In Figure 5 and Table 15 , we see that TuneTables still performs favorably compared to CatBoost and XGBoost on datasets of size up to 50 000. Note that the larger search space may perform better on larger datasets, it is also more challenging to search through in 30 iterations. We present the full results on all datasets in Table 16 . [51] selected to emphasize datasets with many features or many samples, we compare CatBoost and TabPFNs3000. When both models are limited to 3000 samples, TabPFNs3000 performs better on 12 of 17 datasets where significant differences exist. When CatBoost is allowed access to the entire training data, the win rate is identical. In most cases, random sample selection is sufficient for optimal performance. Both models benefit from PCA and mutual information dimension reduction when the feature space is large. The columns labeled SKT / FTS / SMP list the best performing method for sketching, feature subsampling and label-aware sketching technique, respectively. Label-aware sketching refers to a strategy where we either sample instances proportionate to their labels, or we oversample minority classes with replacement to create a class-balanced distribution. While the choice of label-aware sketching strategy is often impactful (and we use it in TuneTables), and the choice of feature subselection method can be important for some datasets, in all but one case, no sketching method we test outperforms random sampling. Bold indicates the best-performing model(s). TuneTables can effectively handle datasets with more classes than the ones used for pretraining, which was not possible with TabPFN. For each algorithm, we compute its mean test accuracy, and mean rank in terms of accuracy. We also compute the mean Z-score, computed by normalizing the set of results on each dataset (by mean 0 std. 1), so that each dataset has the same weight, and averaging each algorithm's normalized performances. We see that TuneTables performs the best across all performance-oriented metrics. Fractional num. wins values are averaged over three splits per dataset, and reflect the presence of multi-way ties on certain splits. 11 and Table 16 , we report the summed runtime for all random seeds (30 in the case of all algorithms except for TuneTables, which generally requires fewer than 30). Since TuneTables uses TabPFN predictions on small datasets, we conservatively report its end to end runtime as the sum of TabPFNs3000 runtime and TuneTables grid search runtime. This prediction is also extremely conservative for TabPFNs3000, as 30 random seeds are unnecessary on datasets with fewer than 100 features. We report runtime in this manner to mitigate concerns that we are advantaging PFNs over boosted trees. We also show the number of samples in each dataset. All datasets above the line were included in Table 1 . Table 9 : Comparison of neural nets on LARGESCALETABLES . We compare TuneTables to other prominent deep learning methods for tabular data on the 17 datasets in LARGESCALETABLES for which all algorithms reported results. For each algorithm, we compute its different metrics of accuracy and rank. We also compute the mean Z-score, computed by normalizing the set of results on each dataset (by mean 0 std. 1), so that each dataset has the same weight, and averaging each algorithm's normalized performances. We see that TuneTables performs the best across all performance-oriented metrics. Fractional num. wins values are averaged over three splits per dataset, and reflect the presence of multi-way ties on certain splits. In this table, we compare the inference time of TabPFNs3000 (which uses up to 3000 points of real data as context, when available) and four common configurations for TuneTables. We report inference time for the entire test set, as well as per 1000 samples, estimated from the inference time for the entire test set. At inference time, TuneTables does not require ensembling to achieve high accuracy, while TabPFNs3000 does. Therefore, TuneTables is considerably faster.Dataset N. classes N. Feats N. Samples TabPFNs3000 TuneTables-pt10-C TuneTables-pt10-NC TuneTables-pt1000-C TuneTables-pt1000-NC 3000 pts, ens. per1k 3000 pts per1k 0 pts per1k 3000 pts per1k 0 pts per1k • The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material.† Equal advising. Correspondence to: bf996@nyu.edu. 


Decision: This paper represents a solid contribution to the field of tabular data analysis and thus deserves to be accepted to the main proceedings of NeurIPS. Reviewers in particular described the experimental results as extensive and mentioned the motivation was clear. There were some major concerns about writing raised by a reviewer that the authors have claimed to improved; I would like to stress that improvements in writing along those lines should be considered as a requirement for acceptance. The authors also provided another baseline at the request of a different reviewer and promised to flesh out the results fully for the main paper; I would expect this promise to be fulfilled as well.

Reviews: review: paper_summary: paper_summary: The paper introduces a novel parameter-efficient fine-tuning strategy for prior-data fitted networks (PFNs) called TuneTables. PFNs, similar to large language models, utilize pretraining and in-context learning to achieve strong performance on new tasks in a single forward pass. However, existing PFNs like TabPFN are limited to small datasets (less than 1000 samples). TuneTables addresses these limitations by compressing large datasets into a smaller learned context, thus significantly improving the performance of PFNs. Besides, this paper gave an demonstration of TuneTables as an interpretability tool and for mitigating biases by optimizing a fairness objective.
strength_weakness: strength_weakness: - Since the dataset was collected by the authors themselves, there are some aspects I am unsure about. In previous benchmarks, FT-Transformer generally outperforms SAINT, but on the authors' benchmark, it does not. Given that tabular data is diverse, it is possible to collect datasets that are more friendly to sample-interaction-based architectures like SAINT, PFN, and TableTunes.

- Excelformer, a well-known method that performs well on both small-scale and large-scale datasets, was not compared. I also suggest that the authors provide tests on Excelformer's large-scale datasets.

- The authors mentioned, "For all algorithms other than TuneTables, we perform light hyperparameter tuning by running one default setting and 29 iterations of random search using Optuna." This is clearly inconsistent with the original settings of most methods. I have concerns about the results.

- This method improves upon TabPFN. However, TabPFN can only be applied to classification tasks, neglecting the regression tasks. Additionally, as stated in Lines 82-88, there are issues like a fixed number of features and a fixed number of classes. These problems persist, even though the authors employed techniques like sketching and feature selection (not to mention that sketching and feature selection can also be applied to other methods, as done in Kaggle competitions, potentially improving their performances). All these factors undermine the significance and value of this paper.
questions: questions: see above
limitations: limitations: The authors clearly summarized their limitations.
paper_summary: paper_summary: This paper aims at removing 3 issues with the recently introduced TabPFN model (a transformer pretrained on synthetic datasets, showing great in-context performance on actual small tabular datasets), namely that it can only take as input few samples, few features, and few classes. The authors manage to solve these issues by a) for classes, training a new decoder with the rest of TabPFN frozen b) for features, searching over a few simple feature reduction methods (mutal info, PCA..) c) for samples, the authors use prompt tuning: some embeddings are trained to improve TabPFN performance while going through a dataset which can be bigger than TabPFN acceptable context length. Combining all these improvement creates a new method named TuneTables.

Experiments on both a standard tabular benchmark and a new benchmark shows that TuneTables beats all their strong baselines for tabular data prediction, even on datasets with many samples, features or classes. While TuneTables is slower than Gradient Boosting Tree algorithms, the authors provides variants with different time-performance tradeoff, and suggest ways to speed up the model's inference time.

Finally, the authors shows that TuneTables can also be used to do things which were not possible to do with TabPFN, for instance optimizing a different objective for fairness, or having access to "summary points" for a specific dataset.
strength_weakness: strength_weakness: Important details are missing, for instance the hyperparameter spaces used for the baselines and for TuneTables. This is an important part of the evaluation (which is not easy to find in the codebase), and I'm interested in how the space depends on the dataset.

Some information is hard to find. I would recommend pointing directly to the relevant results (instead of a broad section like section 4). I would also recommend a table of content for the appendix. For instance I was quite interested in understanding "In Section 4, we find that all sketching and feature selection methods plateau in performance well before approaching parity with GBDTs, which motivates our investigation of new scaling techniques.". This lead me to follow: section 2.1 --> section 4 --> find paragraph on sketching --> appendix C1 --> table 5 --> not sure how to interpret.

Relatedly, all ablations should be reported using aggregated scores, in addition to individual results. This would also make it easier to go through the many (which is good!) ablations.

I also wonder why most ablations and variants are done on the new benchmark TuneTables-hard instead of the more standard TabZilla benchmark. I think it makes sense in some cases, but for instance for TuneTables-medium and light I would be interested in seeing the performance on the TabZilla benchmark.

For these two variants, I would also be interested in seeing a complete performance and runtime comparison with baselines like GBDTs in addition to what you say in the main text. This would help the reader to better understand the time-performance pareto frontier.
questions: questions: > While PFN accuracy can improve with more real-data samples at inference time [33, 55], the memory requirements scale quadratically with the context length

Isn't it linear with FlashAttention (which is used by TabPFN if I remember correctly).

Which hyperparameter spaces are you using for each model?

You say several times that you do a grid search on TuneTables hyperparameters, using 30 steps. But it seems that the hyperparameters space is bigger than 30 possibilities. Are you doing a random search?

How were the datasets in tunestable-hard chosen? (also I wonder whether the name may be a bit misleading because it first got the impression that it contained only dataset with a lot of samples or features)

In Table 11, are the results compared for a given number of HPO steps?

I'm a bit confused about how ensembling is working for TuneTables, with some sentences giving me the impression that you're taking a tuned prompt and permuting the order / labels (à la TabPFN), and other sentences giving me the impression that you're actually training several tuned prompts ("a step constituting either a new tuned prompt fitted from scratch
686 or a new member of a tuned prompt ensemble")

how are coresets computed?
limitations: limitations: The authors have adequately addressed limitations of their work.
paper_summary: paper_summary: The paper proposes TuneTables, a method for improving the performance of PFNs on large datasets with a variable number of features and classes. TuneTables uses prompt tuning (fine-tuning) to learn a small set of parameters (context / synthetic datapoints). Depending on the qualities of the dataset, a feature subselection method or a new decoder may be used. The experiments show that TuneTables outperforms TabPFN on larger datasets. In the analysis, the paper shows that bias (specifically demographic parity) can be mitigated with a regularizer.
strength_weakness: strength_weakness: My main concern with this paper is its presentation.  

Most notably, I found it hard to read Section 4. It is particularly dense. There are a lot of Table/Figure references (in total 8), all of which are not on the same page, so a reader would need to flip back and forth between pages to find the table and figures. 
Furthermore, many of the references in Section 4 point to various resources in the Appendix. 
For example, many analyses are mentioned without the tables to back them up in the main paper. Instead, tables in the appendices are linked. The main paper should be complete. The reader should be able to draw conclusions from the provided experiments (tables and figures). In this case, the conclusions are mentioned but the Appendix is used as additional space for the figures and tables. 

Here are just some of the examples: 
 - "(see Appendix Table 7 for the results). [...] we find that TuneTables achieves [...]"
 - "In Appendix Table 6 we show that [...]"
 - "At inference time, TuneTables is significantly faster than TabPFN; see Table 9" (Table 9 is in the Appendix). 
 - "ablate the tuned prompt size and the sue of ensembles in Table 11 and Table 12, finding that smaller" (Tables 11 and 12 are in the Appendix). 

Interestingly, Section 5 and 6 are not dense and very easy to read which further makes the paper awkward to read (going from a dense Section 4 to a super not dense Section 5 and 6). One way to tackle this is to shorten Sections 5 and 6, making more space for Section 4, allowing for a small subset of Tables in the Appendix to be moved to the main paper which would make Section 4 easier to read. 

Some additional notes regarding presentation:
 - "TuneTabes-hard" is the name of a dataset. The dataset name is too close to that of the method: "TuneTables". The names should ideally be easily distinguishable
 - (Figure 3) "subset of [48]" -- Is there a dataset name for [48]? It would be easier to read if so. "[48]" is used to reference the paper but is also used as the "name" of the dataset. 
 - (Figure 3) "reduces or reverses TabPFN's limitations" -> "addresses TabPFN's limitations"
 - "One limitation of [48] Table 1" -- I'm guessing this is referring to Table 1 from [48], but this is confusing as there is also a Table 1 in this paper. 
 - "GBDTs" - this term is mentioned several times without clearly defining what it is an acronym for, i.e., Gradient Boosted Decision Trees.
 - "Motivated by the limitations of sketching for large contexts" -- To the best of my knowledge, no limitations of sketching were mentioned before this point, so it is unclear what it is being referred to
 - Section 7: "PFNs are neural processes (NPs) [52]" -- If I understand correctly, it appears that NPs and PFNs are being explored separately but are similar models if not the same. As mentioned in the paper, there have been several neural processes works that tackle similar issues with scaling transformer-based models, so it would be nice if the paper connected NPs and PFNs more. For example, clarify that PFNs are neural processes earlier. The focus of the paper on PFNs and the brief mention of "(similar to works on neural processes; see Section 7)" suggested to me that PFNs and NPs are considerably different for the majority of the paper until I read Section 7. 

Minor:
 - Typo: "subset of [48]on" -> "subset of [48] on" (missing a space)
 - Section 4 is titled "Experiments". Section 5 and Section 6 are experiment/analyses-focused, so it's more accurate for them to be subsections of Section 4.
questions: questions: See Weaknesses. 


"[TuneTables] also does not improve on TabPFN for small datasets (fewer than 1000 samples, 100 features and 10 classes); we postulate this is a result of overfitting." -- Could you clarify why you postulate it is due to overfitting? On page 4, it is mentioned that "if a dataset is small enough to run with the original zero-shot version of TabPFN". Wouldn't you just use the original zero-shot version of TabPFN in this case (without any tuning)?

"(b) if there are too many features" -- Could you clarify how you measured what constitutes "too many"?

"limited to around 3000 by conventional GPU sizes." -- Is this assuming a batch size of "1" with "3000" context tokens? Or a larger batch size?


"memory requirements scale quadratically with the context length" -- This is not necessarily the case. At inference time, attention can be computed recurrently (Rabe et al., 2021; Feng et al., 2024), resulting in transformers only needing linear memory with the context length (albeit the amount of computation in total is still quadratic).

---

"Self-attention Does Not Need $O(n^2)$ Memory", MN. Rabe, C. Staats, arXiv:2112.05682, 2021

"Attention as an RNN", L. Feng, F. Tung, H. Hajimirsadeghi, MO. Ahmed, Y. Bengio, G. Mori, arXiv:2405.13956, 2024
limitations: limitations: Yes
paper_summary: paper_summary: The paper proposes a method to compress large training sets into smaller contexts through soft prompt tuning for prior-data fitted networks. The proposed techniques relax the constraints of PFNs by allowing for an increased number of features, more in-context training examples, and a greater variety of classes. The authors show that the proposed methods achieve state-of-the-art performance across a wide range of datasets. In addition, the proposed method can also be used to optimize a fairness objective.
strength_weakness: strength_weakness: 1. The proposed methods to deal with an increased number of features, training examples, and classes are disconnected, and some of the techniques are not new. The feature selection methods are classical methods that are widely used in machine learning and data sciences. Freezing the encoder of a network and fitting a new decoder is also very common in model finetuning.
2. The proposed method shows close performance to GBDTs, while it is much more computationally expensive.
3. The presentation of the results needs to be improved. In Figure 2, the lines are overlapping, making it hard to distinguish. Figure 3 is also confusing without a clear explanation of the y-axis.
questions: questions: 1. For tabular data, do you assume the type of the attributes (numerical, categorical, etc.)?
2. Do you have comparisons with previous works on data summarization?
3. For bias mitigation, I suggest also including a baseline of the base TabPFN while debiasing the selected training data.
4. Could you explain why PFNs only support a fixed number of classes?
limitations: limitations: yes

Comments: comment: Thanks very much for the discussion, and for considering our comments!
Thank you for acknowledging our rebuttal and for taking the time to review our work!
TT-medium and TT-light: Thank you for these further suggestions. For all methods in our table, we will report the runtime per 1000 samples for the TabZilla benchmark in the revised manuscript, binning the runtimes into four groups and then averaging within each bin: (FEATURE-LARGE, SAMPLE-LARGE), (FEATURE-LARGE, SAMPLE-SMALL), (FEATURE-SMALL, SAMPLE-LARGE), (FEATURE-SMALL, SAMPLE-SMALL).

Flash Attention: we thank the reviewer for pointing this out. We will add an appropriate clarification in the revised manuscript.

We thank the reviewer for their comments and a fruitful discussion.
Thank you for acknowledging our rebuttal and for taking the time to review our work!
Thank you for your response, which clearly answers my questions!

**TT-medium and TT-light**: thank you for providing the score of these models on the TabZilla benchmark. The results are interesting though I find the average accuracy to be hard to interpret, often leading to very close score (here XGBoost, TuneTable-Medium and TuneTable-Light). For the runtime, I'm also worried that it roughly follows the runtime for the biggest dataset, which is interesting but doest give the full picture. For the revised manuscript I would suggest metrics like "Runtime per 1000 samples", or reporting runtimes in different metadata feature bins.

**Flash Attention**: the TabPFN code uses Pytorch's MultiHeadAttention, which should default to FlashAttention when possible.

I'm keeping my Accept score.
Thank you for your response. Your answer has partially resolved my issue. I have increased the score.
Thanks for the response. I will maintain my current score.
Thank you for the rebuttal. Although I am unable to verify the presentation improvements, I believe the authors' commitments and stated changes. I lean towards accepting the paper and have updated my score (4 -> 6) to reflect this.