阿里雲機器學習平臺PAI與華東師範大學論文入選SIGIR 2022

阿里雲機器學習平臺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，保留區域性上下文的特徵，計算複雜度為 O（n×w）， n為序列長度， w是滑動視窗大小。
• Global attention：設定一定的global token，這些token將與序列中所有token進行attention計算，從而獲取序列的全域性資訊，計算複雜度為 O（n×g）， g為global token個數。
• Top-k sparse attention：Transformer模型中的attention互動是稀疏且長尾的，對於每個token，僅與其attention互動最高的top-k個token計算attention，複雜度為 O（n×k）。
• AST-aware structure attention：程式碼不同於自然語言序列，有更強的結構特性，透過將程式碼解析成抽象語法樹（AST），然後根據語法樹中的連線關係確定attention計算的範圍。

為了適應現代硬體的平行計算特性，我們將序列劃分為若干block，而非以token為單位進行計算，每個query block與 w個滑動視窗blocks和 g個global blocks以及 k個top-k和AST blocks計算attention，總體的計算複雜度為 O（n（w+g+k）b）， b為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（ ）中，貢獻給開源社群。

論文連結：

參考文獻

[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