阿里云機(jī)器學(xué)習(xí)平臺(tái)PAI與華東師范大學(xué)高明教授團(tuán)隊(duì)合作在SIGIR2022上發(fā)表了結(jié)構(gòu)感知的稀疏注意力Transformer模型SASA，這是面向長(zhǎng)代碼序列的Transformer模型優(yōu)化方法，致力于提升長(zhǎng)代碼場(chǎng)景下的效果和性能。由于self-attention模塊的復(fù)雜度隨序列長(zhǎng)度呈次方增長(zhǎng)，多數(shù)編程預(yù)訓(xùn)練語言模型（Programming-based Pretrained Language Models, PPLM）采用序列截?cái)嗟姆绞教幚泶a序列。SASA方法將self-attention的計(jì)算稀疏化，同時(shí)結(jié)合了代碼的結(jié)構(gòu)特性，從而提升了長(zhǎng)序列任務(wù)的性能，也降低了內(nèi)存和計(jì)算復(fù)雜度。

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

模型框架

下圖展示了SASA的整體框架：

其中，SASA主要包含兩個(gè)階段：預(yù)處理階段和Sparse Transformer訓(xùn)練階段。在預(yù)處理階段得到兩個(gè)token之間的交互矩陣，一個(gè)是top-k frequency矩陣，一個(gè)是AST pattern矩陣。Top-k frequency矩陣是利用代碼預(yù)訓(xùn)練語言模型在CodeSearchNet語料上學(xué)習(xí)token之間的attention交互頻率，AST pattern矩陣是解析代碼的抽象語法樹（Abstract Syntax Tree，AST ），根據(jù)語法樹的連接關(guān)系得到token之間的交互信息。Sparse Transformer訓(xùn)練階段以Transformer Encoder作為基礎(chǔ)框架，將full self-attention替換為structure-aware sparse self-attention，在符合特定模式的token pair之間進(jìn)行attention計(jì)算，從而降低計(jì)算復(fù)雜度。

SASA稀疏注意力一共包括如下四個(gè)模塊：

• Sliding window attention：僅在滑動(dòng)窗口內(nèi)的token之間計(jì)算self-attention，保留局部上下文的特征，計(jì)算復(fù)雜度為，為序列長(zhǎng)度，是滑動(dòng)窗口大小。
• Global attention：設(shè)置一定的global token，這些token將與序列中所有token進(jìn)行attention計(jì)算，從而獲取序列的全局信息，計(jì)算復(fù)雜度為，為global token個(gè)數(shù)。
• Top-k sparse attention：Transformer模型中的attention交互是稀疏且長(zhǎng)尾的，對(duì)于每個(gè)token，僅與其attention交互最高的top-k個(gè)token計(jì)算attention，復(fù)雜度為。
• AST-aware structure attention：代碼不同于自然語言序列，有更強(qiáng)的結(jié)構(gòu)特性，通過將代碼解析成抽象語法樹（AST），然后根據(jù)語法樹中的連接關(guān)系確定attention計(jì)算的范圍。

為了適應(yīng)現(xiàn)代硬件的并行計(jì)算特性，我們將序列劃分為若干block，而非以token為單位進(jìn)行計(jì)算，每個(gè)query block與

個(gè)滑動(dòng)窗口blocks和

個(gè)global blocks以及

個(gè)top-k和AST blocks計(jì)算attention，總體的計(jì)算復(fù)雜度為

，

b為block size。

每個(gè)sparse attention pattern 對(duì)應(yīng)一個(gè)attention矩陣，以sliding window attention為例，其attention矩陣的計(jì)算為：

ASA偽代碼：

實(shí)驗(yàn)結(jié)果

我們采用CodeXGLUE[1]提供的四個(gè)任務(wù)數(shù)據(jù)集進(jìn)行評(píng)測(cè)，分別為code clone detection，defect detection，code search，code summarization。我們提取其中的序列長(zhǎng)度大于512的數(shù)據(jù)組成長(zhǎng)序列數(shù)據(jù)集，實(shí)驗(yàn)結(jié)果如下：

從實(shí)驗(yàn)結(jié)果可以看出，SASA在三個(gè)數(shù)據(jù)集上的性能明顯超過所有Baseline。其中Roberta-base[2]，CodeBERT[3]，GraphCodeBERT[4]是采用截?cái)嗟姆绞教幚黹L(zhǎng)序列，這將損失一部分的上下文信息。Longformer[5]和BigBird[6]是在自然語言處理中用于處理長(zhǎng)序列的方法，但未考慮代碼的結(jié)構(gòu)特性，直接遷移到代碼任務(wù)上效果不佳。

為了驗(yàn)證top-k sparse attention和AST-aware sparse attention模塊的效果，我們?cè)贐igCloneBench和Defect Detection數(shù)據(jù)集上做了消融實(shí)驗(yàn)，結(jié)果如下：

sparse attention模塊不僅對(duì)于長(zhǎng)代碼的任務(wù)性能有提升，還可以大幅減少顯存使用，在同樣的設(shè)備下，SASA可以設(shè)置更大的batch size，而full self-attention的模型則面臨out of memory的問題，具體顯存使用情況如下圖：

SASA作為一個(gè)sparse attention的模塊，可以遷移到基于Transformer的其他預(yù)訓(xùn)練模型上，用于處理長(zhǎng)序列的自然語言處理任務(wù)，后續(xù)將集成到開源框架EasyNLP（https://github.com/alibaba/EasyNLP）中，貢獻(xiàn)給開源社區(qū)。

論文鏈接：https://arxiv.org/abs/2205.13730

參考文獻(xiàn)

[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/1000348767/

本文為阿里云原創(chuàng)內(nèi)容，未經(jīng)允許不得轉(zhuǎn)載。