Laplacian Heads Improve Transformers by Smoothing Token Representations
Yuchong Zhang, Vardan Papyan
原文摘要
Transformers update token representations through multi-head attention and residual connections: $X \leftarrow X + \sum_{i} P^{(i)}\text{Norm}(X)W_{V_i}W_{o_i}$, where $P^{(i)}$ is the softmax probability matrix in head $i$. We propose replacing some $P^{(i)}$'s with $I - P^{(i)}$, giving $X \leftarrow X + \sum_{i \in \mathcal{A}} P^{(i)}\text{Norm}(X)W_{V_i}W_{o_i} + \sum_{i \in \mathcal{L}} (I - P^{(i)})\text{Norm}(X)W_{V_i}W_{o_i}$. We call heads that use $I - P^{(i)}$ Laplacian heads. Our proposal has two motivations. First, it allows attention heads to update the mean of token representations, while Laplacian heads can directly control within-sequence variance. Second, if tokens are viewed as nodes in a graph with normalized edge weights $P^{(i)}$, then $I - P^{(i)}$ is a graph Laplacian operator, and the update can be interpreted as one step of heat diffusion on the graph. We show that this simple modification improves performance across supervised learning, causal language modeling, and vision self-supervised learning tasks. To investigate why, we examine the token representations learned with and without Laplacian heads. In supervised learning, Laplacian heads collapse token representations within the same sequence and align the sequence means with the geometry of Neural Collapse. In language modeling, they increase the separability of token representations that share the same next-token prediction. In vision self-supervised learning, they produce token representations whose top principal components better capture segmentation structure. Across modalities, they also lead to faster-decaying spectra, indicating stronger token smoothing. Overall, our findings challenge the prevailing view that token oversmoothing is inherently harmful, showing instead that certain forms of smoothing can be beneficial.
cs.LG