关于后端:SIGIR-2022面向长代码序列的Transformer模型优化方法提升长代码场景性能

37次阅读

共计 3613 个字符,预计需要花费 10 分钟才能阅读完成。

简介:论文主导通过引入稠密自注意力的形式来进步 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 计算的范畴。

为了适应古代硬件的并行计算个性,咱们将序列划分为若干 block,而非以 token 为单位进行计算,每个 query block 与个滑动窗口 blocks 和个 global blocks 以及个 top- k 和 AST blocks 计算 attention,总体的计算复杂度为,为 block size。

每个 sparse attention pattern 对应一个 attention 矩阵,以 sliding window attention 为例,其 attention 矩阵的计算为:

ASA 伪代码:

试验后果
咱们采纳 CodeXGLUE[1]提供的四个工作数据集进行评测,别离为 code clone detection,defect detection,code search,code summarization。咱们提取其中的序列长度大于 512 的数据组成长序列数据集,试验后果如下:

从试验后果能够看出,SASA 在三个数据集上的性能显著超过所有 Baseline。其中 Roberta-base[2],CodeBERT[3],GraphCodeBERT[4]是采纳截断的形式解决长序列,这将损失一部分的上下文信息。Longformer[5]和 BigBird[6]是在自然语言解决中用于解决长序列的办法,但未思考代码的构造个性,间接迁徙到代码工作上成果不佳。

为了验证 top-k sparse attention 和 AST-aware sparse attention 模块的成果,咱们在 BigCloneBench 和 Defect Detection 数据集上做了融化试验,后果如下:

sparse attention 模块不仅对于长代码的工作性能有晋升,还能够大幅缩小显存应用,在同样的设施下,SASA 能够设置更大的 batch size,而 full self-attention 的模型则面临 out of memory 的问题,具体显存应用状况如下图:

SASA 作为一个 sparse attention 的模块,能够迁徙到基于 Transformer 的其余预训练模型上,用于解决长序列的自然语言解决工作,后续将集成到开源框架 EasyNLP(https://github.com/alibaba/Ea…)中,奉献给开源社区。

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

参考文献
[1] Shuai Lu, Daya Guo, Shuo Ren, Junjie Huang, Alexey Svyatkovskiy, Ambrosio Blanco, Colin B. Clement, Dawn Drain, Daxin Jiang, Duyu Tang, Ge Li, Lidong Zhou, Linjun Shou, Long Zhou, Michele Tufano, Ming Gong, Ming Zhou, Nan Duan, Neel Sundaresan, Shao Kun Deng, Shengyu Fu, Shujie Liu. CodeXGLUE: A Machine Learning Benchmark Dataset for Code Understanding and Generation. NeurIPS Datasets and Benchmarks 2021

[2] Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, Veselin Stoyanov. RoBERTa: A Robustly Optimized BERT Pretraining Approach. CoRR abs/1907.11692 (2019)

[3] Zhangyin Feng, Daya Guo, Duyu Tang, Nan Duan, Xiaocheng Feng, Ming Gong, Linjun Shou, Bing Qin, Ting Liu, Daxin Jiang, Ming Zhou. CodeBERT: A Pre-Trained Model for Programming and Natural Languages. EMNLP 2020

[4] Daya Guo, Shuo Ren, Shuai Lu, Zhangyin Feng, Duyu Tang, Shujie Liu, Long Zhou, Nan Duan, Alexey Svyatkovskiy, Shengyu Fu, Michele Tufano, Shao Kun Deng, Colin B. Clement, Dawn Drain, Neel Sundaresan, Jian Yin, Daxin Jiang, Ming Zhou. GraphCodeBERT: Pre-training Code Representations with Data Flow. ICLR 2021

[5] Iz Beltagy, Matthew E. Peters, Arman Cohan. Longformer: The Long-Document Transformer. CoRR abs/2004.05150 (2020)

[6] Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontañón, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, Amr Ahmed. Big Bird: Transformers for Longer Sequences. NeurIPS 2020

理解更多精彩内容,欢送关注咱们的阿里灵杰公众号

原文链接:http://click.aliyun.com/m/100…

本文为阿里云原创内容,未经容许不得转载。

正文完
 0