PaperlayerView exact revision on arXiv

Source-linked AI summary

Localized TabICLv2: Scaling Tabular In-Context Learning through k-NN

Beimnet Bekele Guta

arXiv:2608.16429v1cs.LG

TL;DR

Full-context TabICLv2 scales poorly because attention over all training rows increases cost with dataset size. Localized TabICLv2 retrieves k nearest neighbours in Stage 2 representation space, retaining 98.64% of Full TabICLv2 accuracy while delivering substantial inference speedups.

  • Problem

    Full-context TabICLv2 attends to all training rows, causing computational and memory costs to grow rapidly with training-dataset size.

  • Method

    Localized TabICLv2 retrieves the top-k training rows in learned Stage 2 representation space and performs Stage 3 inference using only this local context.

  • Results

    98.64% of Full TabICLv2 accuracy is retained on TabArena, with median speedups of 2.18× in batch inference and approximately 249× in single-query serving.

  • Takeaways & Limitations

    Localized inference makes tabular foundation models more efficient and practical for real-world deployment while preserving most predictive performance on TabArena.

  • Takeaways & Limitations

    Localization may miss useful evidence, while retrieval overhead can reduce speedups on smaller datasets or large query batches.

Abstract

from arXiv · show

Foundational models for tabular data have made significant progress in recent years, with TabICLv2 reporting state-of-the-art performance on several tabular classification tasks. However, full-context tabular ICL still suffers from attention cost that grows with the training-context size, which limits its ability to handle large datasets efficiently. Localized TabICLv2 introduces a method that reduces the inference cost of TabICLv2 by retrieving only the k nearest training neighbours for each test point, measured by similarity in the model's Stage 2 row-representation space, rather than using the full training context. This requires no architectural changes, and we show that accuracy retention can be improved through additional Stage 2 and Stage 3 fine-tuning. On TabArena classification tasks, the fine-tuned localized model retains 98.64% of Full TabICLv2 accuracy and it achieves a median 2.18$\times$ speedup in batch inference, and reaches approximately 249$\times$ median speedup in the single-query serving setting.

1. Introduction

Localized TabICLv2 addresses TabICLv2’s scalability limitations by retrieving only the most relevant training rows in the learned Stage 2 representation space. Stage 2 and Stage 3 fine-tuning improves localized inference accuracy retention while reducing inference cost.

  • Motivation: TabICLv2’s full self-attention makes each test point attend to all training rows, causing computational and memory costs to grow rapidly with dataset size.This limitation arises during Stage 3 (TFicl).
  • Method: Localized TabICLv2 retrieves the top-k most similar training rows for each test point using the learned Stage 2 representation space.It then performs standard TabICL inference using only the retrieved local context.
  • Fine-tuning: Stage 2+3 fine-tuning improves representation quality and helps maintain accuracy under localized inference.The method jointly fine-tunes the retrieval representation and inference stages.
  • Evaluation: 98.64% of Full TabICLv2 accuracy is retained on TabArena by the localized model.This result demonstrates high accuracy retention while reducing inference cost.

2. Related Work

Tabular foundation models have progressed from TabPFN to TabICL and TabICLv2, but full-context inference remains costly because each query attends to all training examples. Localized TabICLv2 addresses this limitation by combining k-NN retrieval with a jointly fine-tuned representation space and ICL component.

  • Tabular foundation models: TabPFN introduced transformer-based tabular ICL with synthetic pretraining, but remains limited to smaller datasets.It performs tabular prediction in a single forward pass without task-specific parameter updates.
  • Tabular foundation models: TabICL uses fixed-dimensional row embeddings to jointly process training and test rows, while TabICLv2 adds optimized pretraining and scalable softmax attention.TabICLv2 achieves state-of-the-art results on multiple tabular benchmarks.
  • Inference limitations: Full-context inference remains costly because the final transformer stage attends to all training examples for each query.Retrieval-based localization is motivated by this computational constraint.
  • Retrieval-based methods: LoCalPFN uses k-NN retrieval to construct localized contexts, but relies on a fixed pretrained representation space designed for TabPFN's smaller-data regime.Its retrieval step is not explicitly optimized for the downstream ICL task.
  • Proposed direction: Localized TabICLv2 combines k-NN-based localization with a jointly fine-tuned representation space and ICL component.The learned embeddings are better suited for retrieval while improving accuracy retention under localized inference.

3. Methodology

Localized TabICLv2 preserves TabICLv2’s three-stage inference pipeline while inserting cosine-similarity k-NN retrieval between Stage 2 and Stage 3 to provide each query with a localized training context. Cached Stage 1–2 representations support efficient retrieval, and additional Stage 2–3 fine-tuning addresses accuracy loss from the localized inference setting.

  • Inference pipeline: Three stages produce row representations in Stages 1–2, after which Stage 3 uses labelled support rows to predict query labels.Stages 1–2 perform column-wise embedding and row-wise interaction modeling before in-context prediction.
  • Localized retrieval: Cosine-similarity k-NN retrieval selects the closest training rows in learned Stage 2 representation space for each test datapoint.The retrieved rows form a query-specific localized context between Stage 2 and Stage 3.
  • Inference implementation: Cached Stage 1–2 row representations are reused during inference, so the principal computational reduction comes from replacing full-context attention with fixed local-context attention.Retrieval and systems overhead can prevent the theoretical context reduction from translating directly into wall-clock speedup.
  • Localized retrieval: Stage 3 receives only the retrieved k training datapoints instead of the full training set, reducing effective sequence length while preserving prediction behaviour.This changes the context supplied to the standard TabICL-style Stage 3 predictor.
  • Stage 2–3 Fine-tuning: Stage 2 and Stage 3 are fine-tuned because the original model was not trained for localized inference, which can otherwise reduce accuracy.Stage 1 remains frozen while Stage 2 is adapted for row representations and Stage 3 for the ICL predictor.

4. Results and Discussion

Across 38 TabArena classification datasets, Localized TabICLv2 preserves most Full TabICLv2 accuracy while reducing inference cost through k-nearest-neighbour retrieval. Fine-tuning improves retention, learned Stage 2 representations outperform raw-feature retrieval, and speedups are especially large for limited-query serving.

  • K-Sensitivity: As k increases, accuracy retention improves monotonically while speedup gains diminish; k = 32 yields 97.57% accuracy retention with 1.97× mean speedup.Repeated selection of training rows across queries adds GPU overhead and limits achievable speedup, motivating k = 32 for later fine-tuning.
  • TabArena Accuracy Evaluation: 98.64% of Full TabICLv2 accuracy is retained after S2+S3 fine-tuning, versus 97.31% for pretrained localized inference; fine-tuning improves 33 of 38 datasets.Fine-tuning increases accuracy from 0.8613 to 0.8731, recovering approximately 49.5% of the gap to Full TabICLv2, which remains more accurate overall.
  • TabArena Accuracy Evaluation: Embed FT achieves better AUC, log-loss, and overall accuracy than Raw FT, indicating that Stage 2 representations provide a better retrieval basis than raw features.For binary classification, the fine-tuned localized model exceeds XGBoost’s ROC AUC of 0.8271; for multiclass tasks, it outperforms XGBoost while remaining below Full TabICLv2.
  • Batch-Inference Speed Evaluation: 2.18× median speedup is achieved in batch inference across 15 large TabArena datasets, though speedup varies across datasets.On credit-card-fraud, speedup increases monotonically with Ntrain when other factors are held constant.
  • Limited-Query Latency: 249× median speedup is achieved for a single query across six datasets, with a maximum speedup of 1223×.Full TabICLv2 processes all Ntrain tokens per prediction, whereas localization uses a fixed k-sized subset, yielding approximately 59 ms versus approximately 76 s at N ≈228K.
  • ICL Head Ablation: 37/38 datasets favor Localized TabICLv2 over both XGBoost-on-k-NN and k-NN majority vote using identical retrieved contexts.Because retrieval-only baselines lack the ICL head, the result indicates that the ICL component contributes beyond local label aggregation.

5. Conclusion · A. Theoretical Analysis of Speedup

Localized TabICLv2 reduces inference cost by retrieving only k relevant training neighbours for Stage 3, retaining most predictive performance after Stage 2+3 fine-tuning. Its speedup comes from reducing the Stage 3 bottleneck, although retrieval overhead and missed evidence can limit gains.

  • 5. Conclusion: Localized TabICLv2 retrieves only the most relevant training examples through k-NN in a learned representation space.The retrieved neighbours replace the full training context supplied to the ICL head.
  • 5. Conclusion: Tabular foundation models become more efficient and more practical for realworld deployment with localization.The conclusion frames efficiency as the main practical benefit of the method.
  • 5. Conclusion: Localization may underperform when useful evidence is absent from the retrieved neighbours.This limitation follows from restricting the context to locally retrieved examples.
  • 5. Conclusion: Retrieval overhead can reduce speedups on smaller datasets or large query batches.The retrieval cost can offset the savings from using a smaller Stage 3 context.
  • 5. Conclusion: Future work includes regression evaluation, direct comparison with LoCalPFN, and study of FAISS-style retrieval backends.Other directions include adaptive k, learned retrieval metrics, and improved fine-tuning strategies.
  • A. Theoretical Analysis of Speedup: Stages 1–2 have identical computational cost for Full TabICLv2 and Localized methods because training representations are computed once and test-time Stages 1–2 are shared.The shared KV-cache mechanism means Stages 1–2 do not contribute to relative speedup.
  • A. Theoretical Analysis of Speedup: The speedup difference arises from Stage 3 and retrieval cost, with Stage 3 forming TabICLv2’s main bottleneck.The inference-cost decomposition includes shared Stages 1–2, retrieval, and localized Stage 3 computation.

A.1. Complexity Analysis

Localized TabICLv2 replaces Full TabICLv2’s long-sequence computation with independent processing of each test point using a fixed context of k retrieved neighbours. The resulting theoretical gains are strongest when Ntest ≪ Ntrain, but practical speedup is constrained by retrieval, hardware, and short-sequence efficiency.

  • Complexity setup: Ntrain denotes training rows, Ntest denotes test rows, and k denotes retrieved neighbours.These quantities define the complexity comparison.
  • Full-context complexity: Full TabICLv2 executes one forward pass over a long sequence, including training-set attention and NtrainNtest test-to-training attention.The full-context method processes the combined context in a single pass.
  • Localized complexity: Localized TabICLv2 processes each test point independently with a fixed context of size k, requiring Ntest forward passes over sequences of length k + 1.This changes the computation from one long sequence to many short sequences.
  • Theoretical speedup: When Ntest ≪ Ntrain, localized inference can achieve substantial speed improvements, with gains scaling with training-set size and test-point count.The theoretical reduction is based on the Stage 3 ICL component.
  • Practical limitations: Practical speedup is limited by retrieval, indexing, stacking, data movement, hardware, and implementation effects, especially for short sequences.FlashAttention-style kernels are optimized for large sequences; launch overhead and poor occupancy reduce effective throughput on short sequences.
  • Practical speedup: Localized TabICLv2 achieves its greatest speedup when the training set is large or the number of test queries is small.These conditions reduce the relative impact of per-query overheads.

B. TabArena Evaluation

TabArena evaluation compares Full TabICLv2, several localized variants, and XGBoost across per-dataset primary-metric and accuracy results. Binary tasks use ROC AUC, while multiclass tasks use log-loss, with higher ROC AUC and lower log-loss preferred.

  • Primary-metric results: Table 6 compares six model configurations, including Full TabICLv2, localized variants with pretrained or Stage 2+3 fine-tuning, and XGBoost.The localized configurations include both raw and non-raw variants.
  • Primary-metric results: Binary TabArena tasks report ROC AUC, where higher values are better.
  • Primary-metric results: Multiclass TabArena tasks report log-loss, where lower values are better.
  • Accuracy results: Table 7 reports per-dataset TabArena accuracy for Full TabICLv2, localized variants, and XGBoost.The compared localized variants include pretrained, raw pretrained, and Stage 2+3 fine-tuned configurations.

C. k-Sensitivity Evaluation

Table 8 evaluates sensitivity to k using two outcomes: accuracy retention and predict-time speedup, both reported relative to Full TabICLv2.

  • k-Sensitivity Evaluation: Table 8 reports k-sensitivity results for localized TabICLv2.The evaluation varies k and measures its effects using the listed outcome metrics.
  • k-Sensitivity Evaluation: Accuracy retention is one reported outcome in the k-sensitivity evaluation.Accuracy retention is expressed relative to Full TabICLv2.
  • k-Sensitivity Evaluation: Predict-time speedup is the other reported outcome, measured relative to Full TabICLv2.The table reports accuracy retention and predict-time speedup as paired results.

D. Batch-Inference Speed Evaluation · E. Main Small-Batch Latency Test · F. Controlled Scaling Test: credit-card-fraud

The section evaluates inference efficiency for Full TabICLv2 versus Localized TabICLv2 across batch inference, small-batch serving, and controlled dataset scaling. It examines prediction time, per-query latency, and speedup under varying query and training-context sizes.

  • D. Batch-Inference Speed Evaluation: Batch-inference speed is evaluated directly for the localized and full-context models.Table 9 presents the batch-inference speed evaluation.
  • D. Batch-Inference Speed Evaluation: The batch-inference comparison focuses on the efficiency difference between Full TabICLv2 and Localized TabICLv2.The supplied passage identifies this as a speed evaluation, while the broader section compares the two inference settings.
  • E. Main Small-Batch Latency Test: Small-batch serving results report total prediction time, per-query latency, and speedup across different numbers of test queries.These metrics are explicitly identified for the small-batch latency test.
  • E. Main Small-Batch Latency Test: The main small-batch latency test compares Full TabICLv2 with Localized TabICLv2.Table 10 is explicitly framed as a comparison between the two models.
  • F. Controlled Scaling Test: credit-card-fraud: The controlled scaling test varies both the number of training rows and the number of test queries.This experiment uses the credit-card-fraud dataset to study latency under changing workload dimensions.
  • F. Controlled Scaling Test: credit-card-fraud: Controlled scaling results report total prediction time, per-query latency, and speedup for Full TabICLv2 versus Localized TabICLv2.Table 11 specifies these metrics and the model comparison.

G. ICL Head Ablation

Table 12 compares Localized TabICL with Stage 2+3 fine-tuning against retrieval-only baselines. It evaluates binary tasks using ROC AUC and multiclass tasks using log-loss, with XGB-kNN and kNN-MV as baseline methods.

  • Table 12 compares Localized TabICL with Stage 2+3 fine-tuning against retrieval-only baselines.
  • Binary tasks report ROC AUC, where higher values are better.
  • Multiclass tasks report log-loss, where lower values are better.
  • XGB-kNN denotes XGBoost trained on the retrieved k-nearest-neighbour context, while kNN-MV denotes k-nearest-neighbour majority vote.
Loading 2608.16429v1…