ARASH: Adaptive Retrieval And Shot Selection for Tabular Prediction

Samirasadat Jamalidinan; Yue Xu; Kazem Cheshmi

cs.AIAug 18, 2026

ARASH: Adaptive Retrieval And Shot Selection for Tabular Prediction

Authors: Samirasadat Jamalidinan, Yue Xu, Kazem Cheshmi

Organizations: Department of Electrical and Computer Engineering McMaster University Hamilton, ON, Canada

Abstract

Tabular prediction is a critical task across numerous applications. The recent success of large language models has sparked various approaches for adapting them to the tabular domain. A prevalent strategy involves training or fine-tuning specialized Tabular Foundation Models (TFMs) such as TabPFN. However, TFMs require substantial computational resources, and frequent model retraining is often impractical. In-context learning (ICL), specifically, few-shot prompting, offers a resource-efficient alternative to enhance performance. Yet, identifying the most relevant rows to serve as shots remains a challenge for tabular data. This paper introduces ARASH (Adaptive, query-specific Retrieval And Shot selection), a method that improves TFM efficiency by selecting optimal shots based on local neighborhood analysis within the training set. Our results demonstrate that ARASH reduces the prompt length and memory usage of TabPFN by 1261.5 $\times$ and 2.56 $\times$ , respectively, while providing comparable accuracy.

Figures & tables

Fig. 1 : Motivation example for ARASH. Four synthetic regions are shown. Each point represents one demonstration with two features, shown on the horizontal and vertical axes. The colors in the top row denote routed local regions. The bottom row shows the same points, where colors denote class labels.

Fig. 3 : Pairwise win-rate comparison across baseline models and datasets. A win is defined as achieving higher classification accuracy on a given dataset.

Panel of Figure 3 — Fig. 4 : Accuracy-efficiency comparison between ARASH and ICL demonstration-selection baselines across (a)TabPFN and (b)TabDPT models. Lower number of demonstrations and higher accuracy are better.

Fig. 5 : Inference-only peak

\Delta

VRAM in log scale measured during inference, aggregated across multiple datasets.

Panel of Figure 5 — Fig. 6 : Accuracy improvement of ARASH over the strongest retrieval baseline across datasets for three language-model backbones.

Panel of Figure 6 — Fig. 7 : Representative analysis cases for ARASH. Top: Jungle. Bottom: WDBC. In each row, the left panel shows routed clusters in a two-dimensional PCA projection, the middle panel shows class labels in the same projection, and the right panel reports per-cluster diagnostics, including Hopkins locality, anisotropy, label entropy, and cluster size.

Predictor	Coef.	SE	$p$
Dataset-level locality $\mathrm{L}_{\mathcal{D}}$	0.088	0.027	0.0014
Final purity $P_{\mathcal{D}}$	0.090	0.019	$2.7{\times}10^{-6}$

TABLE I: Diagnostic regression for the ARASH accuracy gap

\Delta_{\mathcal{D},\mathcal{M}}=\mathrm{Acc}^{\mathrm{ARASH}}_{\mathcal{D},\mathcal{M}}-\mathrm{Acc}^{\mathrm{Full}}_{\mathcal{D},\mathcal{M}}

. Robust standard errors are reported over

n=129

observations.

Method	Accuracy	Macro-F1
ARASH	0.750	0.657
ARASH Without Step I (No clustering)	0.700	0.592
ARASH Without Step II (Fixed shots per cluster)	0.720	0.627
ARASH Without Step III (Random retrieval)	0.645	0.495

TABLE II : The effect of the three steps of ARASH + TabpFN on the accuracy of the CREDIT dataset

Variant	Accuracy	Macro-F1
ARASH with Entropy only score	$0.710\pm 0.024$	$0.630\pm 0.038$
ARASH with Impurity only score	$0.720\pm 0.020$	$0.640\pm 0.026$
ARASH with Entropy + impurity	$0.753\pm 0.010$	$0.671\pm 0.029$

TABLE III : The effect of difficulty-score component on ARASH + TabPFN accuracy for the CREDIT dataset

Tabular Prediction

Tabular prediction uses structured rows and columns to forecast labels, values, risks, and outcomes from data.

See all 112 papers

TabSieve: Explicit In-Table Evidence Selection for Tabular Prediction
Feb 12, 2026Yongyao Wang, Ziqi Miao, Lu Yang +4In-Context Example Selection Reinforcement Learning
Gaoling School of Artificial Intelligence, Renmin University of China · Shanghai AI Laboratory
TabFM: A Zero-Shot Foundation Model for Tabular Data
Sep 29, 2026Weihao Kong, Erez Louidor Ilan, Shuxin Nie +6Tabular Foundation Models Tabular Prediction
Google Research
FTFM: Factorized Tabular Foundation Models
Oct 18, 2026Josif GrabockaTabular Foundation Models Tabular ICL
University of Technology Nuremberg

Feb 12, 2026J/K move · Enter open · S save

TabSieve: Explicit In-Table Evidence Selection for Tabular Prediction

Yongyao Wang, Ziqi Miao, Lu Yang, Haonan Jia+3

Gaoling School of Artificial Intelligence, Renmin University of China · Shanghai AI Laboratory

In-Context Example Selection Reinforcement Learning In-Context Learning Tabular ICL Tabular Prediction

Open paper AskPDF arXiv

Tabular prediction can benefit from in-table rows as few-shot evidence, yet existing tabular models typically perform instance-wise inference and LLM-based prompting is often brittle. Models do not consistently leverage relevant rows, and noisy context can degrade performance. To address this challenge, we propose TabSieve, a select-then-predict framework that makes evidence usage explicit and auditable. Given a table and a query row, TabSieve first selects a small set of informative rows as evidence and then predicts the missing target conditioned on the selected evidence. To enable this capability, we construct TabSieve-SFT-40K by synthesizing high-quality reasoning trajectories from 331 real tables using a strong teacher model with strict filtering. Furthermore, we introduce TAB-GRPO, a reinforcement learning recipe that jointly optimizes evidence selection and prediction correctness with separate rewards, and stabilizes mixed regression and classification training via dynamic task-advantage balancing. Experiments on a held-out benchmark of 75 classification and 52 regression tables show that TabSieve consistently improves performance across shot budgets, with average gains of 2.92% on classification and 4.45% on regression over the second-best baseline. Further analysis indicates that TabSieve concentrates more attention on the selected evidence, which improves robustness to noisy context.

Abstract

Tabular Prediction

Explore similar work

TabSieve: Explicit In-Table Evidence Selection for Tabular Prediction

TabFM: A Zero-Shot Foundation Model for Tabular Data

FTFM: Factorized Tabular Foundation Models