Tree Attention: Topology-aware Decoding for Long-Context Attention on GPU clusters

自注意力是现代Transformer架构的核心数学操作，由于其在序列长度上的二次复杂度，也是一个重要的计算瓶颈。在这项工作中，我们推导出标量能量函数，其梯度计算自注意力块，从而阐明自注意力的理论基础，提供了一种贝叶斯解释，并将其与基于能量的模型（如Hopfield网络）紧密联系起来。我们的公式表明，通过树形归约，可以有效地并行计算序列轴上的缩减。我们的算法可以在多个GPU之间并行计算注意力，使跨设备解码的速度比Ring Attention等替代方法快得多（在我们的实验中高达8倍），同时需要显著更少的通信量并且产生2倍的最大内存开销。我们的代码可以在此处公开获取：\url{https://github.com/Zyphra/tree_attention}。

论文旨在解决现代transformer架构中self-attention操作的计算瓶颈问题，并提出一种基于能量函数的计算方法。同时，论文还试图将self-attention与基于能量的模型（如Hopfield Networks）相联系，提供贝叶斯解释。

论文提出了一种基于能量函数的self-attention计算方法，通过树形归约实现了跨设备的并行计算，从而提高了计算效率。

论文提供了self-attention的理论基础和贝叶斯解释，同时提出了一种高效的并行计算方法。实验结果表明，该方法比Ring Attention等替代方案更快，通信开销更小，内存占用更少。论文代码已经公开。

与self-attention相关的研究包括：'Attention is All You Need'、'Transformer-XL'、'Sparse Transformer'等。