# Supervised learning pays attention

arXiv: https://arxiv.org/abs/2512.09912 · HTML: https://mlsanity.com/paper/2512.09912/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 6 of its 22 references that are in the map.
- First posted on arXiv 10 Dec 2025.

## Builds on (most cited first, 6 of 6)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Transformers learn in-context by gradient descent](https://mlsanity.com/paper/2212.07677/) (2022)
- [Quasi-Oracle Estimation of Heterogeneous Treatment Effects](https://mlsanity.com/paper/1712.04912/) (2017)
- [ranger: A Fast Implementation of Random Forests for High Dimensional Data in C++ and R](https://mlsanity.com/paper/1508.04409/) (2015)
- [In-context Learning and Gradient Descent Revisited](https://mlsanity.com/paper/2311.07772/) (2023)
- [Towards Understanding How Transformers Learn In-context Through a Representation Learning Lens](https://mlsanity.com/paper/2310.13220/) (2023)

## Built on by (most cited first, 1 of 1)

- [Interpretable AI with Local Distillation](https://mlsanity.com/paper/2608.23538/) (2026)

## Abstract (arXiv)

In-context learning with attention enables large neural networks to make context-specific predictions by selectively focusing on relevant examples. Here, we adapt this idea to supervised learning procedures such as lasso regression and gradient boosting, for tabular data. Our goals are to (1) flexibly fit personalized models for each prediction point and (2) retain model simplicity and interpretability. Our method fits a local model for each test observation by weighting the training data according to attention, a supervised similarity measure that emphasizes features and interactions that are predictive of the outcome. Attention weighting allows the method to adapt to heterogeneous data in a data-driven way, without requiring cluster or similarity pre-specification. Further, our approach is uniquely interpretable: for each test observation, we identify which features are most predictive and which training observations are most relevant. We then show how to use attention weighting for time series and spatial data, and we present a method for adapting pretrained tree-based models to distributional shift using attention-weighted residual corrections. Across real and simulated datasets, attention weighting improves predictive performance while preserving interpretability, and theory shows that attention-weighting linear models attain lower mean squared error than the standard linear model under mixture-of-models data-generating processes with known subgroup structure.