面向长代码序列的Transformer模型优化方法，提升长代码场景性能

阿里云机器学习平台PAI与华东师范大学高明教授团队合作在SIGIR2022上发表了结构感知的稀疏注意力Transformer模型SASA，这是面向长代码序列的Transformer模型优化方法，致力于提升长代码场景下的效果和性能。由于self-attention模块的复杂度随序列长度呈次方增长，多数编程预训练语言模型（Programming-based Pretrained Language Models, PPLM）采用序列截断的方式处理代码序列。SASA方法将self-attention的计算稀疏化，同时结合了代码的结构特性，从而提升了长序列任务的性能，也降低了内存和计算复杂度。

论文：Tingting Liu, Chengyu Wang, Cen Chen, Ming Gao, and Aoying Zhou. Understanding Long Programming Languages with Structure-Aware Sparse Attention. SIGIR 2022

模型框架

下图展示了SASA的整体框架：

其中，SASA主要包含两个阶段：预处理阶段和Sparse Transformer训练阶段。在预处理阶段得到两个token之间的交互矩阵，一个是top-k frequency矩阵，一个是AST pattern矩阵。Top-k frequency矩阵是利用代码预训练语言模型在CodeSearchNet语料上学习token之间的attention交互频率，AST pattern矩阵是解析代码的抽象语法树（Abstract Syntax Tree，AST ），根据语法树的连接关系得到token之间的交互信息。Sparse Transformer训练阶段以Transformer Encoder作为基础框架，将full self-attention替换为structure-aware sparse self-attention，在符合特定模式的token pair之间进行attention计算，从而降低计算复杂度。

SASA稀疏注意力一共包括如下四个模块：

• Sliding window attention：仅在滑动窗口内的token之间计算self-attention，保留局部上下文的特征，计算复杂度为，为序列长度，是滑动窗口大小。
• Global attention：设置一定的global token，这些token将与序列中所有token进行attention计算，从而获取序列的全局信息，计算复杂度为，为global token个数。
• Top-k sparse attention：Transformer模型中的attention交互是稀疏且长尾的，对于每个token，仅与其attention交互最高的top-k个token计算attention，复杂度为。
• AST-aware structure attention：代码不同于自然语言序列，有更强的结构特性，通过将代码解析成抽象语法树（AST），然后根据语法树中的连接关系确定attention计算的范围。