TuneTables: Context Optimization for Scalable Prior-Data Fitted Networks

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.

Strengths

1. The paper is well-organized and easy to follow, with the problem well-motivated. 2. The paper tackles an important limitation of PFNs, which can inspire follow-up studies and enhance the practical usability of PFNs. 3. The authors conduct extensive experiments that study the advantages and limitations of the proposed method.

Weaknesses

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

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?

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

yes

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.

Strengths

- Impressive number of experiments that show TuneTables outperforms TabPFN - Soft Prompt Tuning for TabPFN is an interesting idea

Weaknesses

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

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

Rating

6

Confidence

4

Soundness

3

Presentation

1

Contribution

3

Limitations

Yes

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.

Strengths

The contributions of the paper are important and novel: this can allow TabPFN to be used in a much bigger set of contexts. The evaluation of the new TuneTables method is well-done: the method is compared to strong baselines (CatBoost, XGBoost, recent deep learning models, but also interesting variants like finetuning TabPFN) on a standard benchmark, with several aggregation metrics. Furthermore, the results of the method are impressive. The paper is quite rich, with a lot of results and ablations. The codebase seems quite easy to install and use (though I haven't actually run anything). I like that the paper starts off by considering simple solutions in Section 2.1 (some of which they keep) before moving to more complex solutions. The paper is well written. The authors are forthcoming when talking about the weaknesses of their method, for instance the slower runtime of their method compared to GBDTs.

Weaknesses

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

> 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?

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

4

Limitations

The authors have adequately addressed limitations of their work.

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.

Strengths

- TuneTables enhances the performance of PFNs on large datasets, surpassing other state-of-the-art models like CatBoost. - Moreover, this paper provides comprehensive experiments and detailed analyses. - The presentation is good, and the methodology is sound.

Weaknesses

- 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

see above

Rating

5

Confidence

5

Soundness

3

Presentation

3

Contribution

2

Limitations

The authors clearly summarized their limitations.