Localized TabICLv2: Scaling Tabular In-Context Learning through k-NN
Organizations: University of Cambridge, Cambridge, United Kingdom.
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 speedup in batch inference, and reaches approximately 249 median speedup in the single-query serving setting.
Figures & tables
| Acc retention (%) | Mean pred (s) | Mean speedup | Median speedup | |
|---|---|---|---|---|
| 16 | 96.93 | 6.72 | 2.77 | 3.04 |
| 32 | 97.57 | 9.19 | 1.97 | 2.14 |
| 64 | 98.40 | 15.61 | 1.14 | 1.11 |
| 128 | 98.90 | 27.22 | 0.66 | 0.66 |
| Method | AUC | Log-Loss | Accuracy |
|---|---|---|---|
| Full TabICLv2 | 0.855 (0.090) | 0.253 (0.198) | 0.885 (0.086) |
| Embed PT | 0.809 (0.107) | 0.339 (0.252) | 0.861 (0.097) |
| Embed FT | 0.827 (0.096) | 0.293 (0.223) | 0.873 (0.090) |
| Raw PT | 0.798 (0.114) | 0.365 (0.264) | 0.857 (0.098) |
| Raw FT | 0.817 (0.105) | 0.319 (0.210) | 0.872 (0.090) |
| XGBoost | 0.826 (0.105) | 0.340 (0.261) | 0.873 (0.093) |
| Total | Full (s) | Loc (s) | Speedup | ||
|---|---|---|---|---|---|
| 10,000 | 2,500 | 12,500 | 0.91 | 0.85 | 1.07 |
| 25,000 | 6,250 | 31,250 | 2.76 | 2.18 | 1.27 |
| 50,000 | 12,500 | 62,500 | 7.39 | 4.55 | 1.63 |
| 100,000 | 25,000 | 125,000 | 22.97 | 9.96 | 2.31 |
| 150,000 | 37,500 | 187,500 | 45.96 | 16.59 | 2.77 |
| 200,000 | 50,000 | 250,000 | 75.93 | 23.66 | 3.21 |
| Full TabICLv2 Latency (ms/query) | Localized TabICLv2 Latency (ms/query) | Approx. Speedup | |
| 1 | 76,000 | 59 | 249 |
| 5 | 15,000 | 12 | 250 |
| 10 | 7,500 | 6 | 248 |
| 50 | 1,500 | 1.3 | 229 |
| 100 | 750 | 0.8 | 192 |
| Baseline | Loc TabICLv2 win rate |
|---|---|
| XGBoost trained on k-NN | 37 / 38 |
| k-NN majority vote | 37 / 38 |
Appendix figures & tables7 assets
Supplementary material from the paper’s appendix.
Appendix
| Dataset | Task | Classes | Metric | Full | S2+S3 | Pretrain | Raw PT | Raw FT | XGBoost |
|---|---|---|---|---|---|---|---|---|---|
| APSFailure | binary | – | ROC AUC | 0.9949 | 0.9911 | 0.9911 | 0.9930 | 0.9929 | 0.9895 |
| Amazon_employee_access | binary | – | ROC AUC | 0.8587 | 0.8247 | 0.8231 | 0.7591 | 0.7636 | 0.8301 |
| Bank_Customer_Churn | binary | – | ROC AUC | 0.8709 | 0.8409 | 0.8115 | 0.8169 | 0.8334 | 0.8426 |
| Bioresponse | binary | – | ROC AUC | 0.8605 | 0.8082 | 0.8268 | 0.8492 | 0.8364 | 0.8702 |
| Diabetes130US | binary | – | ROC AUC | 0.6667 | 0.6390 | 0.6163 | 0.5796 | 0.5934 | 0.6412 |
| E-CommerceShipping | binary | – | ROC AUC | 0.7470 | 0.7392 | 0.7294 | 0.7280 | 0.7366 | 0.7289 |
| Dataset | Task | Full | S2+S3 | Pretrain | Raw PT | Raw FT | XGBoost |
|---|---|---|---|---|---|---|---|
| APSFailure | binary | 0.9955 | 0.9932 | 0.9938 | 0.9938 | 0.9932 | 0.9942 |
| Amazon_employee_access | binary | 0.9519 | 0.9483 | 0.9482 | 0.9431 | 0.9473 | 0.9491 |
| Bank_Customer_Churn | binary | 0.8635 | 0.8532 | 0.8238 | 0.8407 | 0.8525 | 0.8518 |
| Bioresponse | binary | 0.7883 | 0.7390 | 0.7554 | 0.7772 | 0.7608 | 0.8034 |
| Diabetes130US | binary | 0.9121 | 0.9120 | 0.8966 | 0.8636 | 0.9120 | 0.9111 |
| E-CommerceShipping | binary | 0.6838 | 0.6644 | 0.6495 | 0.6474 | 0.6655 | 0.6453 |
| Dataset | Full acc | Full pred (s) | |||||
|---|---|---|---|---|---|---|---|
| credit-card-fraud | 227845 | 0.9996 | 96.47 | 99.99% / 4.59 | 100.00% / 3.47 | 100.00% / 2.24 | 100.00% / 1.29 |
| GiveMeSomeCredit | 120000 | 0.9373 | 26.94 | 99.81% / 3.33 | 99.87% / 2.28 | 99.86% / 1.25 | 99.84% / 0.68 |
| airline_satisfaction | 103904 | 0.9628 | 22.42 | 97.47% / 3.19 | 97.55% / 2.21 | 97.75% / 1.20 | 98.09% / 0.74 |
| jannis | 66986 | 0.7674 | 13.94 | 93.93% / 2.95 | 94.21% / 2.05 | 94.40% / 1.02 | 94.49% / 0.64 |
| APSFailure | 60800 | 0.9955 | 25.74 | 99.77% / 2.99 | 99.77% / 2.67 | 99.76% / 1.64 | 99.80% / 1.08 |
| adult | 39073 | 0.8762 | 4.18 | 97.88% / 1.97 | 98.36% / 1.26 | 98.66% / 0.67 | 98.90% / 0.29 |
| Dataset ( / ) | Full (s) | Loc (s) | Speedup |
|---|---|---|---|
| credit-card-fraud (284.8k / 227.8k) | 99 | 28 | 3.50 |
| kddcup09_appetency (50.0k / 40.0k) | 63 | 24 | 2.67 |
| APSFailure (76.0k / 60.8k) | 55 | 21 | 2.59 |
| MiniBooNE (130.1k / 104.1k) | 51 | 20 | 2.52 |
| covertype (120.0k / 96.0k) | 50 | 20 | 2.50 |
| GiveMeSomeCredit (150.0k / 120.0k) | 44 | 19 | 2.29 |
| Dataset | Full total (s) | Loc total (s) | Full (ms/query) | Loc (ms/query) | Speedup | ||
|---|---|---|---|---|---|---|---|
| credit-card-fraud | 227,845 | 1 | 76.7540 | 0.0628 | 76754.02 | 62.767 | 1222.8 |
| credit-card-fraud | 227,845 | 5 | 76.7352 | 0.0633 | 15347.04 | 12.670 | 1211.3 |
| credit-card-fraud | 227,845 | 10 | 76.7377 | 0.0661 | 7673.77 | 6.608 | 1161.3 |
| credit-card-fraud | 227,845 | 50 | 76.7599 | 0.0778 | 1535.20 | 1.557 | 986.1 |
| credit-card-fraud | 227,845 | 100 | 76.8039 | 0.0903 | 768.04 | 0.903 | 850.1 |
| credit-card-fraud | 227,845 | 500 | 76.9406 | 0.2774 | 153.88 | 0.555 | 277.3 |
| Full total (s) | Loc total (s) | Full (ms/query) | Loc (ms/query) | Speedup | ||
|---|---|---|---|---|---|---|
| 10,000 | 1 | 0.7432 | 0.0524 | 743.16 | 52.377 | 14.2 |
| 10,000 | 10 | 0.7434 | 0.0524 | 74.34 | 5.245 | 14.2 |
| 10,000 | 100 | 0.7503 | 0.0629 | 7.50 | 0.629 | 11.9 |
| 25,000 | 1 | 2.2498 | 0.0538 | 2249.81 | 53.849 | 41.8 |
| 25,000 | 10 | 2.2510 | 0.0548 | 225.10 | 5.479 | 41.1 |
| 25,000 | 100 | 2.2660 | 0.0664 | 22.66 | 0.664 | 34.1 |
| Dataset | Task | Loc-FT | XGB-kNN | kNN-MV | Loc XGB-kNN | Loc kNN-MV |
|---|---|---|---|---|---|---|
| APSFailure | binary | 0.9911 | 0.9779 | 0.9704 | ✓ | ✓ |
| Amazon_employee_access | binary | 0.8247 | 0.6262 | 0.5876 | ✓ | ✓ |
| Bank_Customer_Churn | binary | 0.8409 | 0.6735 | 0.5708 | ✓ | ✓ |
| Bioresponse | binary | 0.8082 | 0.8258 | 0.7775 | ✓ | |
| Diabetes130US | binary | 0.6390 | 0.4828 | 0.4653 | ✓ | ✓ |
| E-CommerceShipping | binary | 0.7392 | 0.7356 | 0.7028 | ✓ | ✓ |
Tabular ICL
Tabular in-context learning (ICL) uses examples in a table to adapt pretrained models to prediction tasks without task-specific weight updates.
See all 108 papers