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

Jingquan Yan, Yuwei Miao, Peiran Yu, Junzhou Huang

Keywords: Attention Mechanisms, Joint Optimization, Pearson Correlation, Data Homogeneity

TL;DR: We theoretically study why Pearson correlation plateaus in attention-based regression and introduce ECA framework to sharpen and extrapolate attention, consistently boosting correlation without compromising MSE, even on homogeneous datasets.

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.

Primary Area: other topics in machine learning (i.e., none of the above)

Submission Number: 1339