【SIGIR 2022】面向长代码序列的Transformer模型优化方法,提升长代码场景性能

本文涉及的产品
交互式建模 PAI-DSW,每月250计算时 3个月
模型训练 PAI-DLC,100CU*H 3个月
模型在线服务 PAI-EAS,A10/V100等 500元 1个月
简介: 论文主导通过引入稀疏自注意力的方式来提高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/EasyNLP)中,贡献给开源社区。

论文链接: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


了解更多精彩内容,欢迎关注我们的阿里灵杰公众号

qrcode_for_gh_0ea952036c3b_344.jpg

相关实践学习
使用PAI-EAS一键部署ChatGLM及LangChain应用
本场景中主要介绍如何使用模型在线服务(PAI-EAS)部署ChatGLM的AI-Web应用以及启动WebUI进行模型推理,并通过LangChain集成自己的业务数据。
机器学习概览及常见算法
机器学习(Machine Learning, ML)是人工智能的核心,专门研究计算机怎样模拟或实现人类的学习行为,以获取新的知识或技能,重新组织已有的知识结构使之不断改善自身的性能,它是使计算机具有智能的根本途径,其应用遍及人工智能的各个领域。 本课程将带你入门机器学习,掌握机器学习的概念和常用的算法。
相关文章
|
8月前
|
数据采集 自然语言处理
传统的序列模型CRF原理
传统的序列模型CRF原理
|
1月前
|
机器学习/深度学习 人工智能 PyTorch
Transformer模型变长序列优化:解析PyTorch上的FlashAttention2与xFormers
本文探讨了Transformer模型中变长输入序列的优化策略,旨在解决深度学习中常见的计算效率问题。文章首先介绍了批处理变长输入的技术挑战,特别是填充方法导致的资源浪费。随后,提出了多种优化技术,包括动态填充、PyTorch NestedTensors、FlashAttention2和XFormers的memory_efficient_attention。这些技术通过减少冗余计算、优化内存管理和改进计算模式,显著提升了模型的性能。实验结果显示,使用FlashAttention2和无填充策略的组合可以将步骤时间减少至323毫秒,相比未优化版本提升了约2.5倍。
72 3
Transformer模型变长序列优化:解析PyTorch上的FlashAttention2与xFormers
|
7月前
|
机器学习/深度学习 缓存
Block Transformer:通过全局到局部的语言建模加速LLM推理
Block Transformer是一种优化自回归语言模型推理效率的新架构,通过块级自注意力来平衡全局和局部依赖,提高吞吐量。模型包含嵌入器、块解码器和令牌解码器,其中块解码器处理全局依赖,令牌解码器处理局部细节。这种方法减轻了KV缓存的延迟和内存开销,尤其是在长序列处理中。实验显示,尽管Block Transformer参数量增加,但推理速度显著提升,尤其是在大块长度和优化的组件比例下,实现了性能与速度的平衡。
350 7
|
5月前
|
机器学习/深度学习
ChatGPT 等相关大模型问题之Transformer中将Attention值转换为概率如何解决
ChatGPT 等相关大模型问题之Transformer中将Attention值转换为概率如何解决
|
6月前
|
机器学习/深度学习
ICML 2024:揭示非线形Transformer在上下文学习中学习和泛化的机制
【7月更文挑战第10天】Rensselaer Polytechnic Institute和IBM的研究者探讨了非线性Transformer在上下文学习的理论基础。他们展示了Transformer如何通过注意力层聚焦相关上下文,并利用MLP层进行预测,揭示了其在不需微调情况下的泛化能力。尽管研究局限于二进制分类和单层模型,它为理解复杂模型在不同任务和领域的潜在适应性提供了新视角。[论文链接:](https://arxiv.org/pdf/2402.15607)**
56 1
|
7月前
|
机器学习/深度学习 人工智能 自然语言处理
揭秘In-Context Learning(ICL):大型语言模型如何通过上下文学习实现少样本高效推理[示例设计、ICL机制详解]
揭秘In-Context Learning(ICL):大型语言模型如何通过上下文学习实现少样本高效推理[示例设计、ICL机制详解]
揭秘In-Context Learning(ICL):大型语言模型如何通过上下文学习实现少样本高效推理[示例设计、ICL机制详解]
|
8月前
|
机器学习/深度学习 自然语言处理 并行计算
【大模型】解释自我注意力的概念及其在LLM表现中的作用
【5月更文挑战第6天】【大模型】解释自我注意力的概念及其在LLM表现中的作用
|
8月前
|
机器学习/深度学习
大模型开发: 解释批量归一化以及它在训练深度网络中的好处。
批量归一化(BN)是2015年提出的加速深度学习训练的技术,旨在解决内部协变量偏移、梯度消失/爆炸等问题。BN通过在每层神经网络的小批量数据上计算均值和方差,进行标准化处理,并添加可学习的γ和β参数,保持网络表达能力。这样能加速训练,降低超参数敏感性,对抗过拟合,简化初始化。BN通过稳定中间层输入分布,提升了模型训练效率和性能。
205 3
|
8月前
|
机器学习/深度学习 算法 数据可视化
YOLO+混合注意力机制 | YOLOv5再加4.3%才可以做对手,Transformer混合设计依旧可以卷
YOLO+混合注意力机制 | YOLOv5再加4.3%才可以做对手,Transformer混合设计依旧可以卷
252 0
|
机器学习/深度学习 人工智能 算法
目标检测模型设计准则 | YOLOv7参考的ELAN模型解读,YOLO系列模型思想的设计源头
目标检测模型设计准则 | YOLOv7参考的ELAN模型解读,YOLO系列模型思想的设计源头
953 0
目标检测模型设计准则 | YOLOv7参考的ELAN模型解读,YOLO系列模型思想的设计源头