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