LLM2D
更长的注意力跨度:通过稀疏图处理技术增加Transformer的上下文长度
Longer Attention Span: Increasing Transformer Context Length with Sparse Graph Processing Techniques
作者: Nathaniel Tomczak, Sanmukh Kuppannagari
发布日期: 2/5/2025
arXiv ID: 2502.01659

摘要

arXiv:2502.01659v1 类型: cross 摘要: 变量堆栈已经在自然语言处理和生物信息学等多个领域表现出巨大的成功。这种成功源于这些模型通过注意力机制来表示和传播序列数据中单个标记之间的成对相互作用。然而,这种操作的主要限制在于其与上下文长度呈二次关系的内存和时间复杂度——需要捕捉相互作用的序列长度。这极大地限制了这些模型能够推理的序列长度。为了减少成对相互作用的数量,使该数量相对于上下文长度不再是二次的,广泛的研究已经通过引入注意力机制的稀疏性来探索减少相互作用的方法,即通过开发稀疏注意力掩码。然而,能够实现“真稀疏性”的高效实现仍然缺乏。 在这项工作中,我们通过提出一种图计算视角的注意力机制来解决这一问题,其中标记被视为图的节点,注意力掩码决定了图的边。在此视角下,我们开发了图处理算法来实现注意力机制。从理论和实验的角度来看,我们证明了我们的算法仅执行所需的计算,即它们是最有效的工作量。我们还使用流行的注意力掩码进行了广泛的实验,以探索稀疏性对执行时间和可实现的上下文长度的影响。我们的实验证明,与当前最先进的注意力实现(如FlashAttention)相比,在长序列长度的情况下,我们的算法可以显著提高执行时间。此外,我们还展示了我们的算法能够在单张NVIDIA A100 GPU(SXM4 80GB)上实现长达1.6亿的序列长度。