TuneTables: Context Optimization for Scalable Prior-Data Fitted Networks

TL;DR

TuneTables enhances prior-data fitted networks for tabular classification by introducing a parameter-efficient fine-tuning method that improves performance, scalability, interpretability, and bias mitigation across numerous datasets.

Contribution

The paper presents TuneTables, a novel context optimization method that significantly improves PFNs' scalability and performance while reducing parameter requirements.

Findings

Outperforms boosted trees like CatBoost on average

Optimizes fewer than 5% of TabPFN's parameters

Enables bias mitigation and interpretability tools

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 19 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. We open-source our code and raw results at https://github.com/penfever/TuneTables.