Transformer 模型中的布朗运动

Transformer 是许多语言和视觉任务的最先进模型。在本文中，我们对其多头自注意力（MHSA）模块进行了深入分析，发现：1）每个 token 是高维特征空间中的随机变量。2）经过层归一化后，这些变量被映射到超球面上的点。3）这些 token 的更新是一种布朗运动。布朗运动具有特殊性质，其二阶项不应被忽略。因此我们提出了一种用于 MHSA 模块的新二阶优化器（一种迭代 K-FAC 算法）。简而言之：所有 token 被映射到高维超球面。缩放点积注意力 $softmax(\frac{\mathbf{Q}\mathbf{K}^T}{\sqrt{d}})$ 只是球面上随机游走的马尔可夫转移矩阵。而深度学习过程会学习合适的核函数以获得这些 token 的合适位置。MHSA 模块中的训练过程对应于一种值得进一步研究的布朗运动。