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 σ~ is below σ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 σ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 σ~∈[0.10,0.24,0.42,0.73] .
Method
F1 ↑
MAE ↓
PCC ↑
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 ↑
PCC@S ↑
PCC@M ↑
MSE ↓
EGN
0.602 ± 0.160
0.647 ± 0.164
0.629 ± 0.135
0.056 ± 0.047
EGN+ECA
0.690 ± 0.202
0.724 ± 0.191
0.716 ± 0.168
0.051 ± 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
(c) Bike Sharing 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
Ntrain / Nval
2000 / 300
σB
1.0
σfloor / σ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×10−4 / 1×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).
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.
Gautam Goel, Mahdi Soltanolkotabi, Peter Bartlett
Simons Institute, UC Berkeley · Department of Electrical and Computer Engineering, University of Southern California
Recent work has demonstrated that transformers and linear attention models can perform in-context learning (ICL) on simple function classes, such as linear regression. In this paper, we empirically study how these two attention mechanisms differ in their ICL behavior on the canonical linear-regression task of Garg et al. We evaluate learning quality (MSE), convergence, and generalization behavior of each architecture. We also analyze how increasing model depth affects ICL performance. Our results illustrate both the similarities and limitations of linear attention relative to quadratic attention in this setting.
I show that ordinary least squares (OLS) predictions can be rewritten as the output of a restricted attention module, akin to those forming the backbone of large language models. The connection comes from viewing OLS as a similarity-based prediction rule in a learned embedding space. In this representation, least squares does not estimate coefficients per se. Instead, it selects an embedding that minimizes squared prediction error by matching training and test vectors through inner products. This maps directly onto the query-key-value structure of attention mechanisms. I then discuss extensions to dimensionality reduction, nonlinearity, and time series econometrics. Monte Carlo simulations and real-data experiments on UCI/OpenML benchmarks show that nonlinear Attention Regression performs competitively against standard machine learning baselines. In the reverse direction, I replace the attention sublayer of a transformer for tabular data with an explicit regression on polynomial features. The resulting model performs comparably to the standard transformer at a fraction of its parameter count.
Philippe Goulet Coulombe
Département des Sciences Économiques, Université du Québec à Montréal