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

2026-08-17Machine Learning

Machine Learning
AI summary

The authors developed a more efficient way to use a model called TabICLv2 for classifying tables of data. Instead of looking at all the training data every time, their new method, Localized TabICLv2, only checks the most similar examples, saving time. They show that by doing some extra training steps, their faster method keeps almost as much accuracy as the original while being much quicker, especially when making single predictions. This makes it better suited for working with large datasets.

TabICLv2tabular datain-context learningnearest neighbor retrievalmodel fine-tuninginference speedupbatch inferencesingle-query serving
Authors
Beimnet Bekele Guta
Abstract
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.