The Convex Corridor: Learning Attention–Laplacian Mixtures Helps Pretraining

Attention computes input-dependent weighted averages of token vectors through the softmax probability matrix $P$. Recently proposed Laplacian heads instead use $I-P$, where $I$ is the identity matrix, to compute the deviation of each token vector from its corresponding average. Although using both Laplacian and attention heads in transformers has been shown to improve performance on many tasks, how to best combine the two operators is still unclear. We address this problem by parameterizing each head's operator as $\delta (\gamma I+P)$, where $\gamma ,\delta \in \mathbb{R}$ are learnable parameters. This modification consistently leads to lower validation loss in causal language modelling pretraining across tested model scales. To understand why, we examine the evolution of $\gamma$ and $\delta$ throughout training, and find that most heads beyond the first layer learn $\gamma \in (-1,0)$, an interval we call the convex corridor. When $\gamma$ falls within this corridor, $\delta (\gamma I+P)=\delta (|\gamma |(P-I)+(1-|\gamma |)P)$, so each head's operator becomes a scaled convex combination of attention $P$ and the negative Laplacian $P-I$, with $\gamma$ controlling the weights of the two operators and $\delta$ controlling the size of the head's output. This pattern emerges across different initializations and alternative parameterizations of the operator such as ${\alpha }_{0}I+{\alpha }_{1}P$ and $\alpha P+\beta (I-P)$, which also improve pretraining. We perform a series of controlled experiments to show that $\gamma$ being in the convex corridor contributes to the improvements in pretraining: restricting $\gamma \ge 0$ increases loss, while fixing $\gamma$ to be some value within the convex corridor already improves upon standard attention even with $\delta$ fixed to 1. Together, these findings suggest that attention should be viewed as one endpoint of the convex corridor, and that exploring the corridor’s interior is a promising direction for future transformer architectures.