论文
可逆查询键耦合与注意力机制组成
Invertible Query-Key Coupling Composes with Attention Mechanisms
摘要
缩放点积注意力将其查询和键形成为独立的线性投影,因此两者在对它们进行评分的点积之前不会交互。我们研究耦合查询-密钥动态,这是一种预评分转换,在标准评分之前通过共享可逆耦合共同演化每个标记的查询和密钥。我们将其实现为真实非体积保持流风格的交替仿射图:耦合是初始化时的恒等式,为每个头添加一小部分参数,并保持 softmax 和周围架构不变。我们将耦合置于现有的注意力方法之上,而不是取代它们,并询问这种组合是否有帮助。在 WikiText-103 上,为差分注意力添加耦合在 150M 和 455M 参数上都得到了改善。在 455M 处,序列长度 512 处的增益非常显着(p=0.003,六个种子),经受了 Bonferroni 校正并在保留的测试分割上复制;它还适用于旋转嵌入训练长度 512、1024 和 2048。当将耦合添加到查询键标准化(在 150M 时显着)和多词元注意(方向)时,会出现相同的附加方向。匹配控制将增益归因于联合预评分耦合而不是增加的容量,并表明去除可逆性保证保留了 455M 增益,但比 150M 的基本方法差得多,因此可逆性是使耦合跨尺度可靠的原因。就其本身而言,耦合会降低 60M 和 150M 处的困惑度(单侧 Welch 测试,p<0.05),但增益会随着尺度的增大而变窄,并且在 455M 处并不显着。我们将结构与耦合流的表达能力联系起来,使用联想回忆研究来映射耦合在哪里有帮助以及在哪里降低敏锐检索的能力,并得出结论,耦合在与评分阶段方法的组合中最有用,而不是作为独立的变化。