基于结构感知稀疏注意力理解长编程语言

基于编程的预训练语言模型（PPLM），如 CodeBERT，已在许多下游代码相关任务中取得巨大成功。由于 Transformer 中自注意力的内存和计算复杂度随序列长度呈二次增长，PPLM 通常将代码长度限制为 512。然而，实际应用中的代码通常较长，例如代码搜索，现有 PPLM 无法高效处理。为解决此问题，本文提出 SASA，一种结构感知稀疏注意力（Structure-Aware Sparse Attention）机制，降低了复杂度并提升了长代码理解任务的性能。SASA 的关键组件为 top-$k$ 稀疏注意力和基于抽象语法树（AST）的结构感知注意力。通过 top-$k$ 稀疏注意力，可以以更低计算成本获得最关键的注意力关系。由于代码结构代表了代码语句的逻辑，是对代码序列特征的补充，我们进一步将 AST 结构引入注意力。在 CodeXGLUE 任务上的大量实验表明，SASA 取得了优于对比基线的性能。