Diagonalized Attention for Individualized Regression: Latent-Row Localization and Prediction

Borui Peng; Liwei Lin; Feifei Wang; Long Feng

stat.MLSep 18, 2026

Diagonalized Attention for Individualized Regression: Latent-Row Localization and Prediction

Authors: Borui Peng, Liwei Lin, Feifei Wang, Long Feng

Organizations: School of Computing & Data Science, University of Hong Kong · School of Statistics, Renmin University of China

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.

Figures & tables

Method	$s=2$		$s=3$
Method	RMSE	CRLR	RMSE	CRLR
GroupLasso	27.431 (0.241)	9.97% (0.28%)	33.401 (0.288)	9.98% (0.26%)
MDSP	11.995 (0.239)	71.58% (0.77%)	1 2.099 (0.254)	75.87% (0.62%)
Two-Step	23.220 (2.516)	41.27% (10.02%)	23.759 (3.305)	54.98% (13.19%)
1-head StdAtt	22.304 (0.215)	85.68% (0.42%)	23.239 (0.243)	8 8.45% (0.30%)
$R$ -head StdAtt	19.215 (0.201)	8 7.95% (2.34%)	21.029 (0.222)	81.78% (4.51%)

Table 2: Performance of different methods across varying numbers of active rows per sample under the Gaussian regression setting. The best and second-best results are marked by bold and underline , respectively.

Figure 1: Performance of different methods under the Gaussian-mixture classification setting with $n\in\{9000,18000\}$ and $(p,d)=(100,70)$ . Each subplot displays classification accuracy (left axis, measured by Accuracy) and ROI localization (right axis, measured by CRLR), with bar heights and error bars representing the average and standard deviation of each metric, respectively.

Figure 2: (a) Distribution of token counts across comments in the UCI online comment dataset, with an inset showing a zoomed-in view of the right-tail region; (b) Classification performance of different methods on the UCI online comment dataset.

Figure S1: Effect of the class-specific dimension on the truncated-tail spectral-overlap condition. Panel (a) reports the median projected $\bar{\gamma}_{\max}^{\ast}$ over 30 paired mean realizations, with empirical $5\%$ – $95\%$ bands, together with the exact pre-projection reference and the analytic cutoff for each $s$ . Panel (b) shows the ordered projected spectra for one paired realization at $s=5$ under the fully shared ( $r_{\mathrm{cs}}/r=0$ ), near-cutoff ( $r_{\mathrm{cs}}/r=0.8$ ), and fully class-specific ( $r_{\mathrm{cs}}/r=1$ ) constructions. The shaded region marks the smallest $d^{\ast}=120$ eigenvalues entering the truncated-tail average.

$s$	analytic cutoff for $r_{\mathrm{cs}}/r$	numerical bracket	shared ( $r_{\mathrm{cs}}=0$ )	class-specific ( $r_{\mathrm{cs}}=r$ )
3	0.505	$(0.500,0.600]$	$2.0000\ [2.0000,2.0000]$ $0/30$	$0.0305\ [0.0301,0.0313]$ $30/30$
5	0.758	$(0.750,0.800]$	$4.0000\ [4.0000,4.0000]$ $0/30$	$0.0536\ [0.0530,0.0540]$ $30/30$
10	0.898	$(0.850,0.900]$	$9.0000\ [9.0000,9.0000]$ $0/30$	$0.1137\ [0.1130,0.1151]$ $30/30$

Table S1: Path from shared to class-specific enhanced directions at

d=600

c_{\mathrm{sig}}=1

, and

c_{\epsilon}=0.01

. Endpoint entries report the median and empirical

5\%

–

95\%

interval over 30 paired mean realizations, with the number satisfying the strict condition on the second line. The numerical transition bracket ends at the first grid value from which all larger grid values satisfy the condition.

Method	$\rho=0.3$		$\rho=0.6$
Method	RMSE	CRLR	RMSE	CRLR
GroupLasso	12.710 (0.199)	9.85% (0.60%)	9.659 (0.153)	9.83% (0.60%)
MDSP	8.702 (0.111)	31.06% (1.15%)	8.697 (0.124)	31.29% (1.07%)
Two-Step	13.367 (0.228)	16.98% (9.20%)	10.199 (0.175)	13.96% (8.75%)
1-head StdAtt	12.430 (0.195)	55.73% (1.73%)	8.942 (0.151)	81.93% (1.13%)
$R$ -head StdAtt	12.375 (0.205)	54.68% (1.60%)	8.352 (0.149)	82.68% (0.96%)

Table S3: Performance of different methods across varying correlation levels under the Gaussian regression setting. The best and second-best results are marked by bold and underline , respectively.

Figure S2: Performance of different methods across varying dimensions under the Gaussian-mixture classification setting with sample size $n=9000$ . Each subplot displays classification accuracy (left axis, measured by Accuracy) and ROI localization (right axis, measured by CRLR), with bar heights and error bars representing the average and standard deviation of each metric, respectively.

Method	Accuracy	CRLR
GroupLasso	85.72%	41.34%
MDSP	63.31%	39.53%
Two-Step	93.51%	72.56%
1-head StdAtt	98.12%	86.21%
$R$ -head StdAtt	97.80%	81.74%
1-head DiagAtt	99.56%	94.60%

Table S5: Performance of different methods on the sentiment analysis example from Marion et al. (2025) . The best and second-best results are marked by bold and underline , respectively.

Explore similar work

Sparse Attention Is Matrix Approximation, Not Choosing from a Bag of Values
Oct 7, 2026Fang Wan, Xufeng Liu, Fan Li +1Sparse Attention Efficient Attention
Stony Brook University · Duke University
Ordinary Least Squares as an Attention Mechanism
Apr 13, 2025Philippe Goulet CoulombeTransformer Attention Linear Regression
Département des Sciences Économiques, Université du Québec à Montréal
In-Context Learning in Linear vs. Quadratic Attention Models: An Empirical Study on Regression Tasks
Feb 19, 2026Ayush Goel, Arjun Kohli, Sarvagya SomvanshiIn-Context Learning Transformer Attention
University of California, Berkeley

Oct 7, 2026J/K move · Enter open · S save

Sparse Attention Is Matrix Approximation, Not Choosing from a Bag of Values

Fang Wan, Xufeng Liu, Fan Li, Yi Liu

Stony Brook University · Duke University

Sparse Attention Efficient Attention Attention Mechanisms LLM Inference Acceleration

Open paper AskPDF arXiv

Large Language Models (LLMs) achieve strong performance across many domains, but their efficiency is limited by the quadratic cost of attention with respect to prompt length. Sparse attention reduces this cost by retaining only a small fraction of query-key interactions to approximate the full attention matrix. However, existing methods are trapped in a mathematically wrong view: they simply keep large scalar entries or high-mass regions of the attention matrix. This treats the attention matrix as a bag of values, ignoring that it is used as a structured matrix whose entries jointly determine the attention output through multiplication with value vectors. We argue that this is the core conceptual issue: sparse attention should be formulated as matrix approximation, not as blindly choosing the largest values from a bag of entries. Based on this view, we propose Matrix Approximation Sparse Attention (MASA). MASA replaces raw attention-mass ranking with a closed-form score that measures how much each sparse unit reduces matrix-product approximation error. As a theory-grounded plug-in correction, MASA can be added to existing sparse attention frameworks without changing their sparse kernels or budgets. Extensive experiments across multiple sparse attention methods, benchmarks, and LLM backbones show consistent accuracy gains, supporting both MASA and the matrix-approximation view of sparse attention.

Abstract

Explore similar work

Sparse Attention Is Matrix Approximation, Not Choosing from a Bag of Values

Ordinary Least Squares as an Attention Mechanism

In-Context Learning in Linear vs. Quadratic Attention Models: An Empirical Study on Regression Tasks