In-Context Learning of Soft Nearest Neighbor Classifiers for Intelligible Tabular Machine Learning (2025)

Abstract: With in-context learning foundation models like TabPFN excelling on small supervised tabular learning tasks, it has been argued that “boosted trees are not the best default choice when working with data in tables”. However, such foundation models are inherently black-box models that do not provide interpretable predictions. We introduce a novel learning task to train ICL models to act as a nearest neighbor algorithm, which enables intelligible inference and does not decrease performance empirically.

## introduction in-context learning (icl) yields state-of-the-art models for small supervised tabular learning tasks, exemplified by tabpfn @xcite . tabpfn is trained to solve supervised tabular learning tasks via in-context learning directly, meaning that at the inference time, the model is effectively fitted to the task without any weight updates. while such a model shows impressive performance, its inference mechanism is not interpretable, and users have to rely on model-agnostic explainability methods @xcite . this is in contrast to recent requirements for more transparent, interpretable, and intelligible models @xcite . 2 k-nearest neighbor (knn) algorithms, a complementary research direction, recently reappeared in tabular state-of-the-art methods, such as modern-nca @xcite and tabr @xcite . knn-based methods make predictions based on the similarity between a query and training samples, thus offering transparent, exampledriven inference. however, the performance of knn is highly dependent on a similarity function and the choice of hyperparameter k, which are both dataset-specific, rendering this approach inappropriate for foundation models (fms) working across many different datasets. the generalization to the soft-nearest neighbor method @xcite bases its predictions on the weighted sum of the labels of all training samples in the dataset, yielding accurate predictions while still being human-interpretable. the modernnca extension @xcite demonstrates that learning the similarity function via a neural network can further boost the performance. in our work, we aim to obtain intelligible, stateof-the-art, off-the-shelf models and study "how can we leverage nearest neighbor methods to make icl more intelligible?" more precisely, we propose a novel training task for tabular icl models, inspired by continuous nearest neighbor methods @xcite and rag @xcite (see figure 1 ). our contributions are the following: the following section discusses related literature on icl for tabular data, nearest neighbor methods in deep learning and intelligible deep learning methods. after introducing and evaluating our method in section 3 and section 4, we further discuss how our method relates to kernel learning in section 5 and conclude with future work and limitations in section 6. ## related work in-context learning for tabular data. one of the successful paradigms for training tabular deep learning (dl) is icl, where a model is trained on many datasets to make predictions for a test (query) set conditioned on the train (support) set. interestingly, the icl regime in the model competes with usual, in-weight learning, and has a transient nature @xcite . early works on icl for tabular data were developed for specific tasks @xcite . later work demonstrated that training these models using purely synthetic data can achieve strong performance @xcite @xcite , but general pre-training on natural data is also possible @xcite . icl models perform well and primarily differ in the data used for pre-training, e.g., real or synthetic data, and architectural design, e.g., cell-based attention (den breejen and yun, 2025), yielding continuous performance improvements over time @xcite @xcite . we leverage this model class and propose a new training task. few works argue that icl models such as tabpfn learn an efficient kernel @xcite , and we will discuss this connection in more detail in section 5. in concurrent work to make tabpfn invariant to class order, @xcite also noted this connection. their resulting model leverages a technique similar to ours but further processes a combination of labels with a non-linear module, because the main emphasis of their work is performance rather than intelligibility. finally, a complementary research direction leverages the icl capability of llms for tabular data, instead of training fms on tabular data @xcite . while they perform well for small datasets, they are computationally expensive, not robust to table manipulations, and inherently struggle with large tables @xcite . development of nearest neighbor algorithms (nna) in deep learning. nnas are used extensively in deep learning models and mostly build on nearest component analysis (nca, @xcite , also known as soft-nn, to allow for back-propagation. in nca, the label for an unseen test sample is predicted by taking a weighted average of all available training samples. the followup work nonlinear nca (nnca) @xcite extends nca to operate on features extracted with a neural network. the work of @xcite uses an nna for few-shot learning and a bi-lstm to capture global context. @xcite generalized this to a differentiable knn selection rule, outputting a set of neighbors, rather than their average. @xcite study explainability of soft-nn methods for image classification from the perspective of the kernel methods. recently, @xcite proved that a 1-nn can be learned in-context with a onelayer transformer. our model continues this line of research and is the first to explicitly combine nnas with icl, by training an icl embedder that captures global context and produces features for a soft-nn. applications of nnas in deep learning. the use of nnas can be broadly categorized into two groups: those where nnas serve as the core model and those where they enhance the performance of a downstream model. retrieval-augmented generation (rag) is a prominent method that improves the performance of an llm by enriching the context with relevant information from an external knowledge base @xcite . nnas are also employed for scaling prompt size in llms @xcite , and context localization in tabular icl models, helping to relax the support set size limitations @xcite @xcite . examples of the models with nna as a core algorithm include the extended version of nnca, modernnca @xcite , which is inspired by rag. our method softknn-icl also falls within the category of models using nna at its core. intelligibility in deep learning. a model's decisions can be made intelligible either by designing the model to be explainable from the outset (intrinsic interpretability) or by applying post hoc explanation methods after training, which often entail a computational overhead or require a separate dataset. basic dl models like mlp, resnet, or transformer are not intrinsically interpretable and require post-hoc explanation methods (molnar, 2025) like shap @xcite or lime @xcite . however, by combining neural networks with explainable methods like gam @xcite , it is possible to leverage the complex features of dl models while maintaining intrinsic explainability. a more exotic approach is to train a deep learning meta-model that predicts the optimal parameters of an explainable model @xcite , which, however, are constrained in size. our model also combines an icl transformer with an nna model, which is considered intelligible if the features of the sample are/or can be made interpretable, e.g., by dimensionality reduction @xcite . ## methodology we are interested in supervised tabular classification, which is the task to predict test labels y @xmath0: @xmath1, y s ), where x @xmath2: @xmath3. here, we focus on icl approaches, which means a pre-trained model f θ is "fitted" on the data set during the inference without weight updates, in contrast to the classical in-weight learning approach. to disambiguate the terminology, when talking about inference, we refer to the test set as query and the training set as support. we introduce a novel learning task that implements a nearest neighbor method. knn is the most popular nearest neighbor method and operates by assigning each query point a label y j based on the majority vote of its k closest neighbors in the support set. this can be written down using an indicator function 1 n (x) := {1 if @xmath4, else 0}, and defining a neighborhood n j := n (x q [j], x s , k) as a function returning a set of nearest neighbors according to a similarity function, most commonly based on the euclidean distance. then, the predicted label is: now, the prediction is the weighted average of all labels in the support set, similar to nadaraya-watson kernel regression @xcite , with the main difference that we do not explicitly condition the similarity function on the distance between inputs. we parametrize the similarity function by introducing an embedding function f θ mapping the raw data to a latent space using information from the support and query sets: a(f θ (x s , y s , x q )) → a s , a @xmath5, 1] m×n . this allows learning a similarity function based on the given learning task, and we will explain later how to use a standard transformer architecture for this. assuming similarity scores are normalized wrt. we propose implementing the similarity function a by taking the transformed embeddings of the query and the support set and a merged kv t matrix of a transformer layer @xmath6: and e @xmath7. thus, a q := a s [q] is a corresponding row of the "attention matrix" representing attention values from query x q to the support samples x s . the merged kv t matrix follows work on learning knn via icl with linear transformers @xcite . this means that we train our transformer model, for a given query point, to attend to similar points in the support set and to make predictions by weighting the labels of all points in the support set based on these similarity (attention) values. for training our model, we use the cross-entropy loss l(ŷ q , y q ) = ce(ŷ q , y q ). we refer to this model as softknn-icl and display its structure in figure 2 . we also experimented with the alternative, potentially more straightforward, implementation which outputs the 1-d logit per token by setting @xmath8, a(e s , e q ) := softmax(e s [•, 1]). however, this version results in inferior convergence and requires advanced pre-training schedules, so we do not consider it further. implementation and hardware details. our implementation is based on the repository of den @xcite , and we will release our code upon acceptance. 3 following other works in the field, e.g., @xcite and den breejen et al. (2024), the model is trained using synthetic data only. concretely, we use the tabforest prior as introduced by den @xcite , which is a mix of the original tabpfn prior @xcite and the forest prior @xcite . in practice, we add information from the label in x s as part of the input token, following the standard tabpfn methodology. optimization is performed using adam @xcite with learning rate of 4e-5. we employ cosine annealing @xcite with linear warmup (10 epochs with 8192 datasets) for learning rate scheduling. softknn-icl is trained using three v100 gpus on 24.6m synthetic datasets. ## experimental evaluation we divide the evaluation of our model into two parts. first, we perform a study using toy problems to analyze the decision boundaries of our model. second, we compare our model against competitor models on standard benchmark datasets. ## conclusion and future work we have demonstrated that a (soft) knn learning task for icl models leads to competitive performance compared to the standard learning task. the resulting softknn-icl is closely related to kernel and metric learning and can be used as a drop-in replacement for tasks requiring intelligibility. additionally, by using softknn-icl, we overcome two limitations of traditional knn methods: (1) the need to tune the number of neighbors, k, and the need to define a neighborhood (similarity) function manually. we hope this spurs research into interpretability methods targeted at instance-based learning methods, and that the in-context learning of a soft neighborhood is a valuable basis for distance learning, potentially even beyond tabular tasks. furthermore, we deem future work along the following directions particularly interesting for tabular machine learning.