论Softmax注意力的表达能力：循环神经网络视角

自推出以来，softmax 注意力因其在广泛任务中的表达能力和可扩展性，已成为现代 Transformer 架构的支柱。然而，softmax 注意力的主要缺点是其内存需求和计算复杂度与序列长度成二次方关系。为了避免 softmax 注意力的二次方瓶颈，人们引入了线性注意力和类似方法来替换 softmax 非线性。尽管这些线性形式的注意力是从原始的 softmax 公式推导出来的，但它们在下游任务的准确率上通常表现不佳。虽然有很强的直觉认为，在查询（query）和键（key）的内积上应用 softmax 非线性，相比其他非线性函数具有更理想的特性，但为什么会存在这种性能差异的问题，至今仍未得到解答。本项工作通过推导 softmax 注意力的循环形式，证明了线性注意力是 softmax 注意力的一种近似。利用这种形式，softmax 注意力的每个部分都可以用循环神经网络 (RNN) 的语言来描述。将 softmax 注意力描述为一种 RNN，使得我们可以对其组件进行消融分析，从而理解每个部分的重要性以及它们之间如何相互作用。通过这种方式，我们的工作有助于解释为什么 softmax 注意力比其同类方法更具表达能力。