Tree Attention: Topology-Aware Decoding for Long-Context Attention on GPU Clusters

本文是LLM系列文章，针对《Tree Attention: Topology-Aware Decoding for Long-Context Attention on GPU Clusters》的翻译。

摘要

自注意力是现代Transformer架构的核心数学运算，由于其序列长度的二次复杂性，也是一个重要的计算瓶颈。在这项工作中，我们推导了标量能量函数，其梯度计算了自我注意块，从而阐明了自注意力的理论基础，提供了对操作的贝叶斯解释，并将其与Hopfield网络等基于能量的模型紧密联系起来。我们的公式表明，通过树约简可以并行有效地计算序列轴上的约简。我们的算法，用于在多个GPU上并行化注意力计算，使跨设备解码的执行速度比Ring attention等替代方法快得多（在我们的实验中快8倍），同时还需要更少的通信量，减少2倍的峰值内存。我们的代码可在此处公开获取。

1 引言

2 相关工作

3 自注意力

4 自注意力作为梯度

5 树注意力

6