摘要

Transformer模型在众多自然语言处理任务中取得了广泛成功。然而，自注意力机制固有的二次时间复杂度限制了其在长文本上的应用。本文提出了一种基于二分划分（Binary Partitioning, BP）的多尺度片段细粒度到粗粒度注意力机制，构建了BP-Transformer（简称BPT）。BPT的注意力连接数为$O(k\cdot n\log (n/k))$，其中$k$为控制注意力稀疏程度的超参数。该模型在计算复杂度与模型表达能力之间实现了良好平衡。在文本分类、机器翻译和语言建模等多个任务上的实验表明，BPT在处理长文本时显著优于以往的自注意力模型。相关代码、超参数设置以及用于稀疏注意力的CUDA内核均已开源，基于PyTorch实现。