论文：刘婷婷、王成宇、陈岑、高明和周傲英。通过结构感知稀疏注意力来理解长编程语言。 2022 年西吉尔

模型框架

下图为SASA整体框架：

其中SASA主要由两个阶段组成：预处理阶段和Sparse Transformer训练阶段。在预处理阶段，获得两个token之间的交互矩阵，一个是top-k频率矩阵，另一个是AST模式矩阵。 Top-k 频率矩阵使用代码预训练语言模型来学习CodeSearchNet 语料库上标记之间的注意力交互频率。 AST模式矩阵是解析代码的抽象语法树（AST）。它是根据语法树的连接关系得到的。代币之间的交互信息。 Sparse Transformer 训练阶段使用Transformer Encoder 作为基本框架，用结构感知的稀疏自注意力替代完全自注意力，并在符合特定模式的token 对之间进行注意力计算，从而降低计算复杂度。

SASA稀疏注意力包括以下四个模块：

滑动窗口注意力：仅在滑动窗口内的标记之间计算自注意力，保留局部上下文的特征。计算复杂度为， 为序列长度， 为滑动窗口大小。全局关注：设置某些全局标记。这些token会与序列中的所有token进行注意力计算，以获得序列的全局信息。计算复杂度为，即全局token的数量。 Top-k稀疏注意力：Transformer模型中的注意力交互是稀疏且长尾的。对于每个token，仅计算注意力交互最高的top-k token。复杂性是。 AST-aware Structure Attention：代码不同于自然语言序列，具有更强的结构特征。将代码解析为抽象语法树（AST），然后根据语法树中的连接关系确定注意力计算的范围。为了适应现代硬件的并行计算特性，我们将序列分成若干块，而不是以令牌为单位进行计算。每个查询块都与

滑动窗块和

全局块和

用top-k和AST块计算attention，整体计算复杂度为

,

b 是块大小。

每个稀疏注意力模式对应一个注意力矩阵。以滑动窗口注意力为例，注意力矩阵的计算为：

ASA伪代码：

实验结果

我们使用CodeXGLUE[1]提供的四个任务数据集进行评估，即代码克隆检测、缺陷检测、代码搜索和代码摘要。我们提取序列长度大于512的数据，形成长序列数据集。实验结果如下：

从实验结果可以看出，SASA在三个数据集上显着优于所有Baselines。其中，Roberta-base[2]、CodeBERT[3]和GraphCodeBERT[4]使用截断来处理长序列，这会丢失部分上下文信息。 Longformer[5]和BigBird[6]是自然语言处理中用于处理长序列的方法，但它们没有考虑代码的结构特征，直接转移到代码任务中效果不佳。

为了验证top-k稀疏注意力和AST-aware稀疏注意力模块的效果，我们在BigCloneBench和缺陷检测数据集上进行了消融实验。结果如下：

稀疏注意力模块不仅提高了长代码任务的性能，还大大减少了显存的使用。相同设备下，SASA可以设置更大的batch size，而full self-attention模型则面临内存不足的问题。具体内存使用情况如下：

作为稀疏注意力模块，SASA可以迁移到其他基于Transformer的预训练模型来处理长序列自然语言处理任务。未来将集成到开源框架EasyNLP（https://github.com/alibaba/EasyNLP）中。对开源社区的贡献。

论文链接：https://arxiv.org/abs/2205.13730

参考文献

[1] 鲁帅、郭大亚、任硕、黄俊杰、Alexey Svyatkovskiy、Ambrosio Blanco、Colin B. Clement、Dawn Drain、蒋大新、唐杜宇、李革、周立东、寿林军、周龙、 Michele Tufano、龚明、周明、Nan Duan、Neel Sundaresan、邓少昆、付胜宇、刘书杰。 CodeXGLUE: 用于代码理解和生成的机器学习基准数据集。 NeurIPS 数据集和基准2021

[2] Yinhan Liu、Myle Ott、Naman Goyal、Jingfei Du、Mandar Joshi、Danqi Chen、Omer Levy、Mike Lewis、Luke Zettlemoyer、Veselin Stoyanov。 RoBERTa: 一种稳健优化的BERT 预训练方法。 CoRR 绝对值/1907.11692 (2019)

[3] 冯章印，郭大亚，唐杜宇，端南，冯小成，龚明，寿林军，秦兵，刘婷，蒋大新，周明。 CodeBERT: 用于编程和自然语言的预训练模型。欧洲管理国家实验室2020

[4] 郭大亚、任硕、鲁帅、冯章印、唐杜宇、刘书杰、周龙、段南、Alexey Svyatkovskiy、付胜宇、Michele Tufano、邓少昆、Colin B. Clement、Dawn Drain、Neel Sundaresan 、 建银、 蒋大新、 周明。 GraphCodeBERT: 使用数据流预训练代码表示。 ICLR 2021

[5] 伊兹·贝尔塔吉(Iz Beltagy)、马修·E·彼得斯(Matthew E. Peters)、阿曼·科汉(Arman Cohan)。 Longformer: 长文档变压器。 CoRR 绝对/2004.05150 (2020)