PaperlayerView exact revision on arXiv

Source-linked AI summary

ARASH: Adaptive Retrieval And Shot Selection for Tabular Prediction

Samirasadat Jamalidinan, Yue Xu, Kazem Cheshmi

arXiv:2608.17856v1cs.AI

TL;DR

Tabular foundation models are costly to use, and selecting relevant few-shot examples for tabular queries remains difficult. ARASH adaptively retrieves query-specific shots using local regions, label purity, and difficulty-aware budgets, achieving comparable accuracy to full-context inference while reducing prompt length and VRAM usage by 1261.5× and 2.56×.

  • Problem

    Tabular in-context learning lacks an effective query-specific strategy for selecting relevant, locally informative shots while controlling substantial computational costs.

  • Method

    ARASH combines locality-aware clustering, difficulty-aware shot allocation, and purity-aware retrieval to select adaptive query-specific demonstrations for tabular foundation models.

  • Results

    ARASH achieves comparable accuracy to full-context inference while reducing prompt length and VRAM usage by 1261.5× and 2.56×, respectively.

  • Takeaways & Limitations

    ARASH provides a more resource-efficient way to use tabular foundation models without materially sacrificing accuracy relative to full-context prompting.

Abstract

from arXiv · show

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.

I. INTRODUCTION · II. MOTIVATION

The paper motivates query-specific shot selection for tabular in-context learning because full-context prompting is costly and locality alone can produce impure, noisy evidence. ARASH combines locality and label purity to adapt shot selection, improve low-purity predictions, and avoid model inference in highly pure regions.

  • I. INTRODUCTION: Tabular prediction supports inference from structured records with missing features or labels across domains including finance and healthcare.Neural models often struggle with tabular data’s heterogeneity and mixed modalities.
  • I. INTRODUCTION: In-context learning adapts fixed-parameter models through the content and structure of few-shot input-label demonstrations.This provides a flexible alternative for supervised prediction without changing model parameters.
  • I. INTRODUCTION: Full-context prompting aggregates the entire training set but causes substantial computational cost, with TabPFN memory scaling quadratically with the number of shots.Tabular foundation models rely on carefully constructed prompts containing relevant contextual information for each query.
  • I. INTRODUCTION: Locality-based retrieval faces an unclear, dataset-dependent k and label impurity, particularly when class imbalance makes minority examples sparse.These issues can produce noisy or unrepresentative context for in-context learning.
  • I. INTRODUCTION: ARASH selects query-specific shots by clustering local training-set regions, finding pure-label regions, and using difficulty-aware selection to determine the optimal shot count.The method is designed to retrieve the minimum number of representative shots for tabular prediction.
  • II. MOTIVATION: The motivation compares Full-context using N = 3280 demonstrations, fixed-k kNN using k = 32 nearest demonstrations, and adaptive locality-and-purity-aware ARASH with TabPFN.The comparison uses four synthetic datasets, each with two features, two labels, and 3280 demonstrations.
  • II. MOTIVATION: 103 times is the reported cost reduction from kNN, which selects local demonstrations from compact geometric regions instead of using the full training set.Full-context costs scale quadratically in computation and linearly in tokens, while locality alone remains insufficient when regions have mixed labels.
  • II. MOTIVATION: Approximately 2–6% is ARASH’s accuracy improvement in low-purity cases, while highly pure regions permit direct label estimation without invoking the tabular foundation model.Purity therefore indicates both the reliability of local evidence and whether expensive model inference can be avoided.

III. ARASH · A. Problem Definition and Overview

ARASH adaptively selects compact, query-specific demonstrations for tabular prediction by jointly using feature-space locality and label purity. It formalizes few-shot prompting and constructs shot sets intended to improve prediction accuracy while reducing demonstrations passed to the model.

  • III. ARASH: ARASH targets compact, query-specific selection of informative demonstrations for tabular prediction with TFMs and LLMs.Unlike fixed local retrieval, it jointly considers feature-space locality and label purity.
  • A. Problem Definition and Overview: ARASH computes each query’s shot set in three stages using kmin, kmax, τloc, and τpur, beginning with training-data profiling and clustering.It partitions training data into feature-space clusters, then assigns each cluster a shot budget kc using a difficulty score based on local uncertainty and label imp…
  • A. Problem Definition and Overview: Few-shot prompting conditions a TFM or large language model on task-specific demonstrations to predict a query label.For classification task T and query xq, the model predicts yq conditioned on Shotsxq.
  • A. Problem Definition and Overview: Each demonstration si = (xi, yi) pairs a tabular instance with its ground-truth label, drawn from the training set Dtr.The training set is represented as labeled tabular instances.
  • A. Problem Definition and Overview: Tabular demonstrations may contain numerical, categorical, or textual feature values.Each xi comprises feature values across these possible data types.
  • A. Problem Definition and Overview: ARASH constructs a compact demonstration set for each query to maximize prediction accuracy while reducing the number of demonstrations passed to the TFM.This defines the shot-selection objective for each query xq.

B. Step I. Locality-aware Clustering

ARASH’s first step partitions the training data into meaningful local regions for downstream shot selection. It uses dataset profiling and automatic, data-dependent clustering rather than a fixed clustering strategy.

  • Locality-aware clustering: ARASH partitions the training set into local regions that serve as the basis for Step II shot selection.The resulting clustering set is C = {c1, …, ct}.
  • Locality-aware clustering: The method dynamically selects a clustering strategy from normalized feature-space characteristics instead of applying one fixed algorithm to every dataset.The selected model creates the clustering set C.
  • Data profiling: DataProfiling computes dataset size, feature dimensionality, normalized representation, PCA-based anisotropy, the Hopkins statistic, and aggregate statistics before clustering.Profiling is applied once to the training set and its outputs guide subsequent heuristic decisions.
  • Locality scoring: The Hopkins statistic is used as the dataset-level locality score LD, with values near 0.5 indicating weak locality and larger values indicating stronger clustering tendency.LD is stored in DP and reused during query routing rather than recomputed for each query or cluster.
  • Automatic selection: The automatic selector evaluates KMeans, MiniBatchKMeans, GMM, Birch, HDBSCAN, and AgglomerativeClustering using dataset-level meta-features.The selector makes one dataset-level clustering decision based on profiling signals such as size, dimensionality, and anisotropy.

C. Step II. Difficulty-aware Shot Assignment

ARASH assigns query-specific shot budgets to local clusters using difficulty-aware analysis, allocating more demonstrations to label-conflicted regions and fewer to reliable ones. It computes difficulty from normalized label entropy and impurity, normalizes scores, maps them to feasible integer budgets, and enforces label-coverage and cluster-size constraints.

  • Difficulty-aware allocation: ARASH assigns each local cluster a query-specific shot budget instead of using a fixed number of shots for every query.The budget is the number of training demonstrations included in the query-specific shot set Shotsxq.
  • Difficulty scoring: The difficulty score dc ∈[0, 1] combines normalized label entropy Hc with label impurity 1 −Pc, where Pc is majority-class purity.Pc is the fraction of shots in cluster c belonging to its majority class.
  • Difficulty scoring: dc = 1 2 (Hc + (1 −Pc)) captures uncertainty and label conflict while avoiding ambiguity from class imbalance alone.A cluster can be imbalanced yet reliable when one class dominates, whereas impurity reflects conflicting local labels.
  • Budget mapping: ˜dc = dc−dmin dmax−dmin normalizes difficulty across clusters, with ˜dc set to 0.5 when all clusters have equal difficulty.The normalized score is then mapped to an integer shot budget by linear interpolation.
  • Budget constraints: Low-difficulty clusters receive budgets near kmin, while label-mixed clusters receive budgets near kmax, subject to kc covering distinct labels and not exceeding cluster size.These constraints preserve local label representation while keeping the selected context feasible for the model.

D. Step III. Query routing and retrieval

ARASH routes each query to a cluster and retrieves k_c demonstrations using one of four strategies selected by dataset-level locality and cluster purity. When the assigned cluster is both local and pure, it uses kNN retrieval to select nearby demonstrations with a consistent label signal.

  • Query routing: ARASH assigns each query x_q to a cluster c_q before retrieving k_c demonstrations.This routing is described in lines 26–35 of Algorithm 1.
  • Strategy selection: The retrieval strategy is selected using dataset-level locality L_D and the purity P_cq of the assigned cluster.Locality indicates whether the feature space supports meaningful neighborhood-based retrieval.
  • Retrieval strategies: In the local-and-pure case, ARASH retrieves k_cq demonstrations from the assigned cluster with kNN.This is the most favorable case because demonstrations are close to the query and likely to provide a consistent label signal.

IV. ARASH IMPLEMENTATION … V. EXPERIMENTAL RESULTS

ARASH uses a portable two-stage pipeline that separates reusable dataset preprocessing from query-time retrieval and inference, with optional direct prediction when retrieved labels agree. The experiments evaluate this retrieval policy with tabular foundation models and serialized language-model backbones, alongside comparisons with prior work.

  • IV. ARASH IMPLEMENTATION: ARASH reduces inference latency by calling the foundation model only when essential.Its two-stage design minimizes overhead by performing preprocessing once for all queries and limiting model calls during inference.
  • A. Efficient Processing and Inference: ARASH preprocesses training-data-dependent information once, then reuses the resulting artifacts during query-time inference.Preprocessing includes feature normalization, data profiling, clustering, cluster statistics, and cluster-level shot-budget allocation.
  • A. Efficient Processing and Inference: When all retrieved demonstrations share one label, ARASH assigns the query to DIRECT mode and returns that label without calling the TFM.Queries with differing retrieved labels are assigned to TFM mode and passed to the model with the selected demonstrations.
  • B. Language-model Serialization: ARASH’s portable architecture supports LLaMA 3.2, Qwen 2.5, and FLAN-T5 beyond tablenative foundation models.The evaluated models include decoder-only LLaMA 3.2 and Qwen 2.5, plus encoder–decoder FLAN-T5.
  • B. Language-model Serialization: Language-model experiments require serialization because these models consume text rather than native tabular tensors.The serialization layer preserves ARASH’s control over which demonstrations enter the prompt.
  • B. Language-model Serialization: For language-model backbones, ARASH serializes each row as a feature–value key–value representation and includes retrieved demonstrations as input–label pairs.The serialized query and demonstrations are concatenated, followed by a short task-specific instruction.
  • V. EXPERIMENTAL RESULTS: The experimental-results section evaluates ARASH, compares it with prior work, examines language-model integration, and provides a technical algorithm breakdown.The reported comparisons include pairwise win rates, where a win means higher classification accuracy on a dataset.

A. Setup … B. Tabular baselines

ARASH is evaluated on OpenML-CC18 and Combo using TALENT tabular baselines, accuracy, and Macro-F1. ARASH combined with TabDPT or TabPFN achieves performance comparable to full-context models, with ARASH+TabDPT leading average accuracy.

  • Datasets and environment:: ARASH is evaluated on a combined OpenML-CC18 and Combo benchmark, with most analyses and ablations conducted on Combo datasets.Combo provides a compact, diverse classification benchmark that limits evaluation cost.
  • Baselines:: The study compares ARASH with representative TALENT tabular-learning baselines using standardized implementations and hyperparameter settings.The baselines include deep tabular models covered by the latest TALENT release.
  • Baselines:: DCNv2 and DANets target high-order feature interactions, while TabNet, GrowNet, and TabCaps provide structured deep-learning pathways.TANGOS, PTaRL, and SwitchTab use advanced training techniques for standard neural networks.
  • Evaluation Metrics:: The evaluation reports average accuracy and Macro-F1, while ARASH thresholds are tuned on a random 10% dataset split excluded from training, evaluation, and testing.The tuning set uses seed=42 and grid-searches locality, purity, and auto-clustering thresholds.
  • B. Tabular baselines: ARASH+TabPFN and TabDPT obtain the strongest overall pairwise accuracy win-rate performance against baseline models.The comparison reports the percentage of datasets where the row model exceeds the column model in accuracy.
  • B. Tabular baselines: 0.9017 average accuracy is achieved by ARASH+TabDPT, exceeding full-context TabDPT’s 0.9009 while recording Macro-F1 values of 0.800 and 0.795, respectively.ARASH+TabDPT therefore slightly improves accuracy and Macro-F1 over full-context TabDPT.
  • B. Tabular baselines: 0.895 accuracy and 0.800 Macro-F1 are achieved by ARASH+TabPFN, closely matching full-context TabPFN’s 0.8948 accuracy and 0.799 Macro-F1.XGBoost is the strongest nonfoundation-model competitor, with 0.879 accuracy and 0.785 Macro-F1.

C. ICL Demonstration-Selection Baselines · 1) Accuracy: · 2) VRAM measurement:

ARASH is compared with kNN, DPP, random, and full-context ICL on TabPFN and TabDPT using accuracy and inference-time VRAM. It achieves stronger fixed-budget accuracy than kNN and DPP while reducing memory pressure through adaptive retrieval.

  • C. ICL Demonstration-Selection Baselines: The baseline budgets use k={8, 16, 32} demonstrations per query, compared with ARASH’s average of 24–30 demonstrations.The baselines are kNN-k, DPP-k, and Random-k.
  • 1) Accuracy:: Both kNN and DPP improve accuracy as k increases, but their best fixed-budget results remain below ARASH.This indicates that local retrieval is useful for tabular ICL, while ARASH performs better at comparable fixed budgets.
  • 1) Accuracy:: 0.9017 accuracy with 29.38 demonstrations per query is achieved by ARASH on TabDPT, versus 0.8603 for kNN-32 and 0.8489 for DPP-32.These are the reported TabDPT fixed-budget results.
  • 2) VRAM measurement:: ∆VRAM measures prediction-time peak allocation minus the post-fit baseline, isolating memory induced by retrieved context.The procedure records memory after fit(), resets CUDA peak statistics, and logs maximum allocation during predict().
  • 2) VRAM measurement:: ARASH consistently incurs lower ∆VRAM than full context on TabPFN, while kNN’s ∆VRAM increases monotonically with k.The results associate adaptive retrieval with reduced memory pressure and linear growth in context length and activation memory for fixed-k retrieval.

D. ARASH for Serialized-Table Language Models · E. Analysis: When Does ARASH Work? · 1) Representative cases.:

ARASH is evaluated for serialized-table language models and shows positive median gains over the strongest fixed retrieval baseline across three backbones. Analysis identifies distinct locality–purity regimes, with representative cases showing when compact local retrieval succeeds or fails.

  • D. ARASH for Serialized-Table Language Models: ARASH is evaluated with FLAN-T5, LLaMA-3.2, and Qwen-2.5 on serialized tabular rows using in-context demonstrations.The comparison varies demonstration-selection strategy while keeping the language-model backbone and evaluation setting aligned.
  • D. ARASH for Serialized-Table Language Models: The serialized-table evaluation compares demonstration-selection methods under a fixed budget so that performance reflects selected-demonstration quality.The supplied passage frames the comparison around the quality of selected demonstrations rather than differing prompt budgets.
  • D. ARASH for Serialized-Table Language Models: Across all three backbones, ARASH achieves positive median gains over the strongest fixed retrieval baseline.The baseline is the best-performing method among random selection, kNN-feature retrieval, and kNN-text retrieval, with gains defined as ∆ = ARASH − baseline.
  • E. Analysis: When Does ARASH Work?: 43.2%, 44.2%, 9.7%, and 2.9% of queries fall into the high-locality/high-purity, high-locality/low-purity, low-locality/high-purity, and low-locality/low-purity groups, respectively.These four synthetic regimes are described as commonly occurring in real-world datasets and motivate distinct retrieval considerations.
  • 1) Representative cases.: Jungle is difficult for purely local retrieval because routed clusters overlap, labels are mixed, and some clusters have high entropy.Nearby points can therefore carry conflicting labels, making locality alone insufficient for reliable demonstration selection.
  • 1) Representative cases.: 35,855 training rows define the Full-context setting, while fixed-budget retrieval with k = 128 falls short of its accuracy.The best DPP result reaches 0.950, and the corresponding kNN result remains below Full-context and ARASH accuracy.
  • 1) Representative cases.: WDBC is favorable for local retrieval because it has clear local structure, consistent labels, and low per-cluster label entropy.ARASH matches full-context inference at 0.982 accuracy using only k = 12 demonstrations.

2) Diagnostic association model.:

The diagnostic association model relates ARASH’s accuracy gap to dataset-level locality and cluster purity. Both diagnostics are positively associated with smaller gaps, supporting compact local retrieval when neighborhoods are meaningful and labels are consistent.

  • Diagnostic association model: For each dataset D and foundation model M, the accuracy gap measures ARASH accuracy relative to Full-context accuracy.Negative, near-zero, and positive values indicate respectively lower, comparable, and higher ARASH accuracy while using fewer demonstrations.
  • Diagnostic association model: The additive model predicts the gap from dataset-level locality LD and average cluster purity PD.LD is computed from the Hopkins statistic, while PD is the average purity of the dataset’s final clusters.
  • Results: Both coefficients are positive over n = 129 observations, linking stronger locality and higher final cluster purity to smaller ARASH–Full-context accuracy gaps.The association remains after accounting for the other diagnostic.
  • Results: The positive coefficients support Step III thresholding, which selects compact cluster-local, hybrid, diversity-aware, or global retrieval based on locality and purity.This is consistent with ARASH being most reliable when local neighborhoods provide meaningful and label-consistent evidence.

F. ARASH Implication on Inference Time · G. Ablation Study · 1) Component ablation:

ARASH reduces inference cost through cached preprocessing, adaptive context selection, and a direct path for label-consistent demonstrations. Ablations show that locality-aware clustering and difficulty-aware shot allocation improve prediction quality.

  • F. ARASH Implication on Inference Time: Cached preprocessing makes table-lookup overhead negligible, leaving local TFM inference as the primary contributor to ARASH execution time.ARASH uses a two-phase framework and direct mode; preprocessing is performed once and reused across subsequent queries.
  • F. ARASH Implication on Inference Time: 43.9× smaller context accompanies a 1.37× latency reduction: ARASH averages 13.69 demonstrations and 0.436 seconds per query versus 600.25 and 0.598 for full-context TabPFN.The ARASH timing includes amortized clustering, routing, retrieval, and local TFM inference.
  • F. ARASH Implication on Inference Time: 0.02 seconds is the upper bound for ARASH’s absolute preprocessing time per query, while local TFM inference takes 0.422 seconds.Most of the query time is attributed to local TFM inference.
  • F. ARASH Implication on Inference Time: 0.22 of queries use ARASH’s direct path, which bypasses model execution for label-consistent retrieved demonstrations and achieves 0.989 accuracy.The direct path is most frequent in high-locality/high-purity queries.
  • F. ARASH Implication on Inference Time: 0.479 of queries use the direct path in the high-locality/high-purity regime, compared with 0.156 when either diagnostic is low and 0.025 when both are low.Direct-path usage varies with local structure and purity diagnostics.
  • G. Ablation Study: Replacing difficulty-aware allocation with a fixed number of shots per cluster reduces accuracy by 2%, indicating that cluster-level difficulty guides shot allocation.The ablation study examines whether ARASH’s main components are necessary across the Combo dataset, with CREDIT used for detailed diagnostics.
  • 1) Component ablation:: Removing locality-aware clustering reduces CREDIT accuracy from 0.750 to 0.700 and Macro-F1 from 0.657 to 0.592.The result indicates that global retrieval is less effective than constructing prompts from local regions.

2) Difficulty-score ablation: · VI. RELATED WORK · VII. CONCLUSION

ARASH’s difficulty score combines normalized label entropy with cluster impurity, reflecting complementary aspects of local label uncertainty. The paper situates this inference-time retrieval approach among trainable tabular models, retrieval-based ICL, and tabular foundation-model methods.

  • 2) Difficulty-score ablation:: ARASH computes cluster difficulty as d_c = αH_c + (1 − α)I_c, combining normalized label entropy with cluster impurity.Entropy captures uncertainty across the full label distribution, while impurity captures the absence of a dominant local label.
  • 2) Difficulty-score ablation:: Entropy and impurity are related but distinct because clusters with identical majority-class fractions can distribute their remaining labels differently.
  • VI. RELATED WORK: Tree ensembles, particularly gradient-boosted decision trees such as XGBoost, dominate tabular prediction and handle heterogeneous features, missing values, small data, and non-smooth boundaries.
  • VI. RELATED WORK: Neural tabular models include attention-based architectures and methods with structured biases or self-supervision, whereas ARASH uses a frozen model for inference-time selection.The cited architectures include TabTransformer, FT-Transformer, SAINT, AutoInt, TabNet, DANets, TANGOS, and SwitchTab.
  • VI. RELATED WORK: In-context learning adapts tasks through prompt demonstrations, and retrieval-based methods select query-relevant examples because performance is sensitive to demonstration choice.Prior approaches include nearest-neighbor prompting and diversity-aware retrieval methods.
  • VI. RELATED WORK: Retrieval techniques designed for non-tabular data rely on geometric locality, while tabular data has arbitrary row order that complicates locality and purity.
  • VI. RELATED WORK: Tabular foundation-model research serializes rows for language models or develops tabular-native models such as TabPFN and TabDPT; textual prompting becomes expensive with many rows.
  • VII. CONCLUSION: ARASH combines locality-aware clustering, difficulty-aware shot selection, and retrieval to augment query context for tabular prediction with TFMs and LMs.The method was evaluated using TabPFN, TabDPT, and several language models.
Loading 2608.17856v1…