Retrieval & Fine-Tuning for In-Context Tabular Models

Motivation

• Tabular-data heterogeneity makes performant neural networks difficult to obtain, helping tree-based methods such as XGBoost and CatBoost remain dominant.TabPFN addresses this challenge through prior-fitting over millions of possible data-generating processes, enabling new datasets to be processed in a single forward pass without task-specific training or hyperparameter tuning.Source: §1 · p. 1

Problem

• Transformer-based tabular in-context learning scales poorly because memory grows quadratically with context size, limiting performance when larger datasets cannot fit into memory.TabPFN can also underfit as dataset complexity increases, even when the entire dataset fits in context.Source: §1 · pp. 1–2

Motivation

• Adaptive local contexts let TabPFN classify complex concentric-circle patterns that it cannot solve using the entire training set as context.The figure compares full-context TabPFN with nearest-neighbour contexts and reports results averaged over 25 seeds.Source: Figure 1 · p. 2, §1 · p. 2

Approach and contributions

• LoCalPFN combines nearest-neighbour retrieval with task-specific fine-tuning, achieving state-of-the-art performance against neural and well-tuned tree-based methods across 95 TabZilla datasets.The approach builds on TabPFN and adapts its context locally for each query.Source: §1 · p. 2

Contributions

• The paper identifies scaling failures in both dataset size and complexity, proposes retrieval and fine-tuning to use context more effectively, and validates the approach through extensive experimentation and ablation.These contributions include analysis of TabPFN across synthetic and real datasets and comparisons with strong baselines.Source: §1 · p. 2

2.1 Preliminaries on In-Context Learning for Tabular Data and TabPFN

• TabPFN performs classification across tasks without task-specific training by using training examples as context, but transformer memory limits restrict usable context size.Its prior-fitting procedure trains on synthetic datasets generated from diverse neural-network processes.Source: §2.1 · pp. 2–3, §2.2 · p. 3

2.2 What Constitutes a Good Context for Tabular Data?

• Random subsampling does not reliably scale with dataset size, while a single global context cannot adapt to the differing neighborhoods of individual query points.Tabular data also lacks the natural ordering available in language, complicating context design.Source: §2.2 · p. 3
• Local kNN contexts provide a query-specific alternative because nearby training points are hypothesized to contain the most critical classification information.This approach replaces the global training context with kNN(xqy), producing the TabPFN-kNN variant; k controls the expressivity trade-off between smoothing and complex decision boundaries.Source: §2.2 · p. 3, §2.3 · pp. 3–4

2.3 Better Expressivity and Scaling Using Local Information

• Large improvements in AUC occur as synthetic dataset complexity increases when TabPFN retrieves k = 10, 30, 100, or 300 neighbors instead of using the full context.Using k = 1000 is equivalent to vanilla TabPFN, so local retrieval forms a continuum between full-dataset and local contexts.Source: §2.3 · pp. 3–4
• TabPFN-kNN benefits from larger training sets even when k is much smaller than N, whereas TabPFN cannot efficiently exploit datasets exceeding its context capacity.Figure 2 reports this trend across three real datasets, with improvements often appearing at lower context lengths.Source: §2.3 · p. 4, Figure 2 · p. 4

2.4 Efficient End-to-End Fine-Tuning With Retrieved Samples

• End-to-end fine-tuning adapts the model to each dataset, while shared local contexts reduce the computational cost of training with retrieved samples.The approximation samples training points, computes their neighborhoods, and splits shuffled neighbors into shared context and query sets; the authors report faster training without significant performance degradation.Source: §2.4 · pp. 4–5, Figure 3 · p. 5

• Existing tabular deep learners use self-supervised pre-training or retrieval, but they remain limited by context size, lack of ICL generalization, or costly dataset-specific hyperparameter tuning.Source: §3 · p. 5
• TabPFN is a transformer-based in-context learner for tabular data whose rapid prior-fitting adaptation has motivated extensions such as generation, dataset distillation, and interpretability studies.Source: §3 · p. 6
• TabPFN extensions have explored sketching and fine-tuning, yet single shared contexts limit adaptation, motivating closer coupling between retrieval and fine-tuning.Source: §3 · p. 6
• Language-model research motivates jointly retrieving and fine-tuning examples, while TabPFN avoids the data-leakage concern associated with applying text-pretrained foundation models to tables.Source: §3 · p. 6

4.1 Experimental Setup

• The evaluation uses 95 OpenML-sourced datasets filtered for TabPFN compatibility, with 10-fold cross-validation and 30 rounds of baseline hyperparameter tuning per fold using validation AUC.Datasets have at most 100 features and 10 classes, no NaN values, at least one instance per class in each split, and up to 130,064 instances.Source: §4.1 · p. 6

4.2 Main Experiments

• LoCalPFN outperforms all other baselines averaged over 95 datasets, significantly improving on TabPFN and remaining competitive with tuned tree-based methods.On medium/large datasets, LoCalPFN further improves on TabPFN-kNN and is on par with or better than the other methods.Source: §4.2 · p. 7, Table 1 · p. 7

4.3 Analysis: Scaling with Dataset Size and Complexity

• LoCalPFN scales favourably with dataset size, while TabPFN’s relative AUC drops drastically beyond 3,000 instances and LoCalPFN even outperforms XGBoost on large datasets.Figure 4 reports AUC relative to Random Forest with 95% confidence intervals.Source: §4.3 · p. 7, Figure 4 · p. 8
• Across increasing complexity quantiles, TabPFN scales poorly whereas LoCalPFN continues to outperform the other methods in higher-complexity datasets.Complexity is approximated by the difference between the best and worst AUC across the listed algorithms.Source: §4.3 · pp. 7–8

4.4 Ablation Studies

• End-to-end fine-tuning jointly with retrieval is necessary: random-context fine-tuning improves TabPFN, but adding TabPFN-kNN afterward provides no further gain.The ablation therefore supports fine-tuning the model together with retrieval rather than treating the stages independently.Source: §4.4 · p. 8, Figure 5 · p. 9
• Retrieval embeddings affect performance modestly: raw standardized features are competitive, while one-hot encodings provide some improvement mainly on smaller datasets.The compared alternatives also include TabPFN encoder outputs with periodically recomputed search indices.Source: §4.4 · p. 8, Figure 5 · p. 9
• Query-dependent local contexts outperform global-context alternatives, including ensembled TabPFN variants and context averaging, while the latter can make inference scale as O(N).Task-specific context optimization improves performance but remains inferior to LoCalPFN.Source: §4.4 · pp. 8–9, Figure 5 · p. 9
• LoCalPFN improves TabPFN-kNN across neighbour-context sizes, outperforming Random Forest with a maximum of 50 neighbours and XGBoost with a maximum of 500.TabPFN-kNN is not very sensitive to the maximum when it is at least 100.Source: §4.4 · p. 9, Figure 6 · p. 9

• LoCalPFN improves tabular classification through retrieval and fine-tuning, while TabPFN-kNN provides a no-fine-tuning variant with practical utility.Source: §5 · p. 9
• The demonstrated improvements are limited to TabPFN, so transfer to future tabular in-context base models is uncertain.The framework also inherits TabPFN’s feature and class constraints and is not easily evaluated on regression tasks.Source: §5 · p. 9
• LoCalPFN outperforms tree-based methods in the experimental study, although its runtime is slower.TabPFN-kNN runs as fast as tree-based methods on the studied datasets while retaining respectable performance.Source: §5 · p. 10

• The benchmark includes 47 small datasets.Source: Table 2 · p. 14
• A separate benchmark group contains 48 medium/large datasets.Source: Table 3 · pp. 14–15
• Another table lists 71 datasets selected for benchmarking deep learning models.Source: Table 4 · pp. 15–16

• Available TabZilla results provide the tree-based and deep-learning baseline models.Source: §A.2.1 · p. 16
• The study runs different TabPFN inference variations directly using the TabPFN repository.Source: §A.2.1 · p. 16

• TabPFN-kNN uses the minimum of 10 times the square root of dataset size and 1000 neighbours, with batched FAISS inference.Source: §A.2.2 · p. 16
• LoCalPFN and TabPFN-kNN are equivalent at inference step 0, after which LoCalPFN is trained with AdamW and approximate local contexts.Training uses 1,000 sampled query points, batch size 2, and the same neighbour count as TabPFN-kNN.Source: §A.2.2 · p. 16
• LoCalPFN is evaluated every 30 gradient steps and uses validation-set AUC for fold-specific early stopping.Source: §A.2.2 · p. 16
• The proposed experiments can run on a single NVIDIA RTX 6000 GPU Ada Generation with 995Gi RAM and a 64-core AMD CPU.Source: §A.2.2 · p. 16

A.3.1 Comparison to Deep Learning Models

• LoCalPFN outperforms all compared deep-learning models on the 71-dataset subset.The comparison uses a subset where all deep-learning methods could run, because many were too computationally expensive for the full benchmark.Source: §A.3.1 · p. 16, Table 5 · p. 17

A.3.2 Comparison with Other Metrics

• LoCalPFN matches or outperforms other models on both accuracy and F1 score.The corresponding comparisons are reported in Tables 6 and 7.Source: §A.3.2 · p. 17, Table 6 · p. 17, Table 7 · p. 17

A.4 Additional Analyses

• LoCalPFN continues to outperform on the largest and most complex datasets, where TabPFN fails to scale in both size and complexity.The scaling analysis is summarized in Figure 7.Source: Figure 7 · p. 18
• TabPFN-kNN obtains lower training and test losses but a larger train–test gap, with overfitting most evident for few neighbours and small datasets.TabPFN generally overfits less, especially on larger datasets, while TabPFN-kNN provides a different underfitting/overfitting trade-off.Source: Figure 8 · p. 19

A.5.1 Importance of Joint Retrieval and Fine-tuning

• Jointly retrieving neighbours and fine-tuning improves performance drastically, whereas applying TabPFN-kNN to an already fine-tuned model degrades overall performance.The ablation is summarized in Table 8; one-hot encoding provides a relatively moderate gain over other feature-encoding configurations.Source: Table 8 · p. 20, Table 9 · p. 20

A.5 Ablation Studies

• The ablations examine TabPFN inference methods and sensitivity to the maximum neighbour count, with detailed results reported in Tables 10 and 11.Figures 9 and 10 summarize ablations separately for small and medium/large datasets.Source: §A.5 · p. 20, Figure 9 · p. 20, Table 10 · p. 21

A.5.5 Quality of Efficient Local Context

• Approximate neighbour search achieves very similar AUC to exact search under the compared configurations, supporting the efficacy of efficient local-context retrieval.The comparison uses LoCalPFN-exact-b32 and LoCalPFN-approx-q32, with the exact variant capped by GPU memory.Source: §A.5.5 · p. 21, Table 12 · p. 21

A.5.6 Run Time Analysis

• LoCalPFN achieves the highest AUC overall but requires more runtime, whereas TabPFN-kNN runs very quickly with AUC only slightly below XGBoost.Runtime measures total training and evaluation time; other deep-learning methods take even longer than LoCalPFN.Source: §A.5.6 · pp. 21–22, Figure 11 · p. 22