拉普拉斯头通过平滑令牌表示改进 Transformer

Transformer 通过多头注意力和残差连接更新令牌表示：$X \leftarrow X + \sum_{i} P^{(i)}XW_{V_i}W_{o_i}$，其中 $P^{(i)}$ 是第 $i$ 个注意力头的 softmax 注意力矩阵。我们提出用一部分 $P^{(i)}$ 的拉普拉斯矩阵 $I - P^{(i)}$ 替代，得到 $X \leftarrow X + \sum_{i \in \mathcal{A}} P^{(i)}XW_{V_i}W_{o_i} + \sum_{i \in \mathcal{L}} (I - P^{(i)})XW_{V_i}W_{o_i}$。我们的方案有两个动机：首先，它允许注意力头更新令牌表示的均值，而拉普拉斯头则直接控制序列内的方差。其次，如果将令牌视为带边权 $P^{(i)}$ 的图中的节点，则 $I - P^{(i)}$ 是对应的图拉普拉斯矩阵，该更新可解释为图上热扩散的一个步骤。我们展示，这一简单修改在监督学习、语言建模和自监督学习任务中均提升了性能。为 investigation 原因，我们检查了带和不带拉普拉斯头学习的令牌表示。在监督学习中，拉普拉斯头在同一序列内消弭令牌表示，并将序列均值与神经坍缩的几何对齐。在语言建模中，它们增加了共享相同下一个标记预测的令牌表示的可分离性。在自监督学习中，它们产生的令牌表示的主成分更适合分割。跨模态而言，它们也导致更快的谱衰减，表明更强的令牌平滑。总体而言，我们的发现挑战了令牌过度平滑本质上有害的普遍观点，显示出某些形式的平滑实际上是有益的。