【SIGIR 2022】面向長代碼序列的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（//github.com/alibaba/EasyNLP）中，貢獻給開源社區。
論文鏈接：//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
了解更多精彩內容，歡迎關注我們的阿里靈傑公眾號
