Breaking the Correlation Plateau: On the Optimization and Capacity Limits of Attention-Based Regressors

Jingquan Yan; Yuwei Miao; Peiran Yu; Junzhou Huang

cs.LGFeb 19, 2026

Breaking the Correlation Plateau: On the Optimization and Capacity Limits of Attention-Based Regressors

Authors: Jingquan Yan, Yuwei Miao, Peiran Yu, Junzhou Huang

Organizations: Computer Science and Engineering Department, University of Texas at Arlington

Abstract

Attention-based regression models are often trained by jointly optimizing Mean Squared Error (MSE) loss and Pearson correlation coefficient (PCC) loss, emphasizing the magnitude of errors and the order or shape of targets, respectively. A common but poorly understood phenomenon during training is the PCC plateau: PCC stops improving early in training, even as MSE continues to decrease. We provide the first rigorous theoretical analysis of this behavior, revealing fundamental limitations in both optimization dynamics and model capacity. First, in regard to the flattened PCC curve, we uncover a critical conflict where lowering MSE (magnitude matching) can paradoxically suppress the PCC gradient (shape matching). This issue is exacerbated by the softmax attention mechanism, particularly when the data to be aggregated is highly homogeneous. Second, we identify a limitation in the model capacity: we derived a PCC improvement limit for any convex aggregator (including the softmax attention), showing that the convex hull of the inputs strictly bounds the achievable PCC gain. We demonstrate that data homogeneity intensifies both limitations. Motivated by these insights, we propose the Extrapolative Correlation Attention (ECA), which incorporates novel, theoretically-motivated mechanisms to improve the PCC optimization and extrapolate beyond the convex hull. Across diverse benchmarks, including challenging homogeneous data setting, ECA consistently breaks the PCC plateau, achieving significant improvements in correlation without compromising MSE performance.

Figures & tables

Figure 1: (a) Illustration of a video-based sentiment analysis example. A sample is considered homogeneous when its within-sample dispersion $\tilde{\sigma}$ is below $\sigma_{0}$ . (b) A convex attention yields an aggregated embedding inside the convex hull of the sample’s embeddings. (C) Our ECA extrapolates beyond the hull to amplify within-sample contrasts.

Figure 2: PCC plateau example: PCC flattens at early epoch while MSE continues to decrease. The plateau becomes more noticeable as within-sample homogeneity increases, specifically, dataset A is more homogeneous than B.

Figure 3: Validation of gradient ratio (PCC/MSE) decay. The RMS ratio of PCC vs. MSE gradients (blue) is strictly constrained by the theoretical upper bound (red). The increase in prediction dispersion

\sigma_{\hat{y}}

(green) during training drives the attenuation of the PCC gradient signal relative to the MSE gradient.

Figure 5: PCC and MSE curves on synthetic datasets with in-sample homogeneity $\tilde{\sigma}\in[0.10,0.24,0.42,0.73]$ .

Method	F1 $\uparrow$	MAE $\downarrow$	PCC $\uparrow$
TFN † ( Zadeh et al., 2017 )	0.791	0.947	0.673
LMF ∗ ( Liu et al., 2018 )	0.824	0.917	0.695
EF-LSTM † ( Williams et al., 2018b )	0.785	0.949	0.669
LF-DNN † ( Williams et al., 2018a )	0.786	0.955	0.658
Graph-MFN † ( Zadeh et al., 2018 )	0.784	0.956	0.649
MulT ∗ ( Tsai et al., 2019 )	0.828	0.871	0.698

Table 2: Results on MOSI. † from THUIAR’s GitHub ( Yu et al., 2020 ; Yu et al., 2021 ) ; ∗ from ( Hazarika et al., 2020 ) ; ∗∗ reproduced from public code with provided hyper-parameters.

Method	PCC@F $\uparrow$	PCC@S $\uparrow$	PCC@M $\uparrow$	MSE $\downarrow$
EGN	0.602 $\pm$ 0.160	0.647 $\pm$ 0.164	0.629 $\pm$ 0.135	0.056 $\pm$ 0.047
EGN+ECA	0.690 $\pm$ 0.202	0.724 $\pm$ 0.191	0.716 $\pm$ 0.168	0.051 $\pm$ 0.048

Table 3: Three-fold regression PCC and MSE on 10xProteomic dataset.

Figure 6: PCC@M and MSE curves for EGN baseline and EGN+ECA on 10xProteomic dataset under each fold. Curves for PCC@S and PCC@F please refer to Figure 8 and Figure 9 .

Appendix figures & tables15 assets

Supplementary material from the paper’s appendix.

Appendix

(a) Automobile Dataset

(b) Auto MPG Dataset

(d) Garment Dataset

(e) Real Estate Dataset

(f) Servo Dataset

(g) Shares Dataset

(h) Solar Flare Dataset

Parameter	Value
DGP Settings
$D$ (Embedding dimension)	16
$K$ (Elements per sample)	10
$N_{\text{train}}$ / $N_{\text{val}}$	2000 / 300
$\sigma_{B}$	1.0
$\sigma_{\text{floor}}$ / $\sigma_{\text{label}}$	0.01 / 0.01

Appendix

Table 4: Hyperparameters for Synthetic Experiments

Figure 8: Per-fold PCC@F comparison of EGN and EGN+ECA on 10xProteomic dataset.

Figure 9: Per-fold PCC@S comparison of EGN and EGN+ECA on 10xProteomic dataset.

Param	Value
Optimizer / weight decay / lr	AdamW / $1\times 10^{-4}$ / $1\times 10^{-4}$
Batch size / epochs	64 / 100
Text / audio / vision dims	768 / 5 / 20
Projection dim (all modalities)	128
Token count / token dim	8 / 128
Seq. lengths (l/a/v)	50 / 375 / 500

Appendix

Table 5: Configuration for MOSI experiments (ALMT backbone).

Figure 20

Figure 21

Figure 22

Softmax Attention

Softmax attention weights token representations by query–key similarity and underpins Transformer retrieval, sequence modeling, and long-context computation.

See all 75 papers

Training Dynamics of Softmax Self-Attention: Fast Global Convergence via Preconditioning
Mar 2, 2026Gautam Goel, Mahdi Soltanolkotabi, Peter BartlettSelf-Attention Softmax Attention
Simons Institute, UC Berkeley · Department of Electrical and Computer Engineering, University of Southern California
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
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

Mar 2, 2026J/K move · Enter open · S save

Training Dynamics of Softmax Self-Attention: Fast Global Convergence via Preconditioning

Gautam Goel, Mahdi Soltanolkotabi, Peter Bartlett

Simons Institute, UC Berkeley · Department of Electrical and Computer Engineering, University of Southern California

Self-Attention Softmax Attention Gradient Descent Dynamics Gradient Descent

Open paper AskPDF arXiv

We study the training dynamics of gradient descent in a softmax self-attention layer trained to perform linear regression and show that a simple first-order optimization algorithm can converge to the globally optimal self-attention parameters at a geometric rate. Our analysis proceeds in two steps. First, we show that in the infinite-data limit the regression problem solved by the self-attention layer is equivalent to a nonconvex matrix factorization problem. Second, we exploit this connection to design a novel "structure-aware" variant of gradient descent which efficiently optimizes the original finite-data regression objective. Our optimization algorithm features several innovations over standard gradient descent, including a preconditioner and regularizer which help avoid spurious stationary points, and a data-dependent spectral initialization of parameters which lie near the manifold of global minima with high probability.

Abstract

Softmax Attention

Explore similar work

Training Dynamics of Softmax Self-Attention: Fast Global Convergence via Preconditioning

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

Ordinary Least Squares as an Attention Mechanism