Diagonalized Attention for Individualized Regression: Latent-Row Localization and Prediction
Abstract
Modern text and image representations are often matrix-valued, with rows corresponding to tokens, patches, or other local feature vectors. Predictive information is often sparse but sample-specific, making classical sparse regression methods with a common support poorly suited to this heterogeneity. This paper formalizes an individualized sparse regression framework for matrix-valued covariates in which each observation has its own rows of interest, while the associated regression effects are shared across the population. To estimate this model, we introduce a diagonalized attention mechanism that uses query--key scores to localize sample-specific signal rows and a value matrix for downstream regression. The proposed method has a parameter dimension independent of sample size and can identify rows of interest for new observations without their responses. We establish existence theorems showing that, under suitable score-separation and concentration conditions, single-head and multi-head diagonalized attention models recover the latent rows with high probability, yielding prediction risk bounds. Our theory therefore provides a statistical explanation of how attention-based scoring localizes sample-specific signals in heterogeneous matrix-valued data. Simulations demonstrate strong prediction and localization in regression and misspecified classification across varying sample sizes, dimensions, and signal cardinalities. Real sentiment analyses show improved classification accuracy and interpretable token selection.
Related
- Linked-Tucker Factorized Individualized Regression for Paired Multivariate Categorical Outcomes
- Understanding axial attention in TSFMs through single location regression
- Sharp Generalization for Nonparametric Regression in Interpolation Space by Shallow Over-Parameterized Neural Networks with Channel Attention
- An automated ranking platform for machine learning regression models for meat spoilage prediction using multi-spectral imaging and metabolic profiling
- Multi-source conformal prediction: leveraging heterogeneity via localization
BibTeX
@misc{peng2026diagonalized,
title = {{Diagonalized Attention for Individualized Regression: Latent-Row Localization and Prediction}},
author = {Borui Peng and Liwei Lin and Feifei Wang and Long Feng},
howpublished = {arXiv},
year = {2026},
month = sep,
doi = {10.48550/arxiv.2609.21320},
eprint = {2609.21320},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2609.21320},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - Diagonalized Attention for Individualized Regression: Latent-Row Localization and Prediction
AU - Borui Peng
AU - Liwei Lin
AU - Feifei Wang
AU - Long Feng
PY - 2026
DA - 2026/09/18
PB - arXiv
DO - 10.48550/arxiv.2609.21320
UR - https://arxiv.org/abs/2609.21320
AB - Modern text and image representations are often matrix-valued, with rows corresponding to tokens, patches, or other local feature vectors. Predictive information is often sparse but sample-specific, making classical sparse regression methods with a common support poorly suited to this heterogeneity. This paper formalizes an individualized sparse regression framework for matrix-valued covariates in which each observation has its own rows of interest, while the associated regression effects are shared across the population. To estimate this model, we introduce a diagonalized attention mechanism that uses query--key scores to localize sample-specific signal rows and a value matrix for downstream regression. The proposed method has a parameter dimension independent of sample size and can identify rows of interest for new observations without their responses. We establish existence theorems showing that, under suitable score-separation and concentration conditions, single-head and multi-head diagonalized attention models recover the latent rows with high probability, yielding prediction risk bounds. Our theory therefore provides a statistical explanation of how attention-based scoring localizes sample-specific signals in heterogeneous matrix-valued data. Simulations demonstrate strong prediction and localization in regression and misspecified classification across varying sample sizes, dimensions, and signal cardinalities. Real sentiment analyses show improved classification accuracy and interpretable token selection.
ER -