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

TL;DR

The Extrapolative Correlation Attention (ECA), which incorporates novel, theoretically-motivated mechanisms to improve the PCC optimization and extrapolate beyond the convex hull, is proposed, which consistently breaks the PCC plateau.

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 (10)

• 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 4: The standard deviation of predictions $\sigma_{\hat{y}}$ increases during training to match the standard deviation of labels $\sigma_y$ under MSE loss.
• Figure 5: PCC and MSE curves on synthetic datasets with in-sample homogeneity $\tilde{\sigma}\in[0.10, 0.24, 0.42, 0.73]$.

Theorems & Definitions (37)

• Proposition 2.1: MSE Mean–std–correlation Decomposition
• Lemma 2.1: Scaling-invariance of PCC
• Remark 2.1
• Remark 2.2: The PCC Plateau
• Lemma 2.2: Softmax Aggregator Jacobian
• Theorem 2.1: Gradient of PCC w.r.t. Attention Logits
• Lemma 2.3: Gradient of MSE w.r.t. Attention Logits
• Corollary 2.1: PCC/MSE Gradient Ratio Decay
• Corollary 2.2: PCC Gradient Magnitude Bound
• Remark 2.3: The Two Bottlenecks of Softmax Attention for Correlation