高效训练张量注意力：从三次到近线性时间

张量注意力是一种多视角注意力机制，能够捕获多个模态之间的高阶关联，克服了经典矩阵注意力的表示局限性。然而，张量注意力的 $O(n^3)$ 时间复杂度是其在 transformer 中的应用所面临的重大障碍，其中 $n$ 为输入序列长度。本文我们证明，张量注意力训练的反向梯度可在近线性时间 $n^{1+o(1)}$ 下计算，即在有界条目假设下，其前向计算具有相同的复杂度。我们提供了梯度的闭式解，并提出一种利用多项式近似方法和张量代数技术的快速计算方法。此外，我们通过硬度分析证明了我们假设的必要性和紧致性，表明稍微削弱该假设会使梯度问题在真正的亚三次时间内无法求解。我们的理论结果确立了高阶 transformer 训练的可行性，可能促进张量注意力架构的实际应用。