每 Token 恒定成本的 Softmax 注意力

我们提出了一种对 Transformer 所应用的传统注意力机制的简单修改：我们不再使用缩放点积来量化查询-键的成对相似度，而是使用指数的缩放点积的对数来量化它。我们的修改通过指数核特征映射将注意力线性化，其对应的特征函数是无限维的。我们证明，我们的修改可以表示为指数的 log-sum 组合，并具有恒定大小的潜在空间，从而能够以每 Token 恒定的时间和空间复杂度进行应用。我们实现了我们的修改，验证了其在实践中有效，并得出结论：它是传统注意力机制的一个有前景的替代方案。