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

自我注意力是现代变形器结构的核心数学运算，也是由于其在序列长度方面的二次复杂度而成为一个重要的计算瓶颈。在这项工作中，我们推导出标量能量函数，其梯度计算自我注意力块，从而阐明了自我注意力的理论基础，提供了一种贝叶斯解释操作，并将其与基于能量的模型（例如霍普菲尔德网络）紧密联系起来。我们的公式表明，可以通过树形约简在序列轴上实现高效并行计算。我们的算法可以在多个GPU之间并行计算注意力，使跨设备解码的速度呈指数级提高（在我们的实验中高达8倍），而且需要的通信量显著减少，并且峰值内存减少了2倍。我们的代码在此处公开可用：\url{https://github.com/Zyphra/tree_attention}。

论文旨在推导出能够计算自注意力块的标量能量函数的梯度，从而阐明自注意力的理论基础，提供一种贝叶斯解释，并将其与基于能量的模型（如Hopfield Networks）紧密联系起来。此外，论文还试图解决自注意力的计算复杂度问题。

论文的关键思路是通过推导出能够计算自注意力块的标量能量函数的梯度来解决自注意力的计算复杂度问题，并揭示了通过树形归约可以有效地并行计算序列轴上的约简。

论文的亮点包括提出了一种能够并行计算自注意力的算法，使跨设备解码的速度可以渐进性地更快（在实验中高达8倍），同时需要的通信量更少，峰值内存减少了2倍。此外，论文还提供了开源代码。

近期在该领域的相关研究包括《Attention is All You Need》、《Transformer-XL》和《Reformer: The Efficient Transformer》等。