面向長程式碼序列的 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，整體的計算複雜度為

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（https://github.com/alibaba/EasyNLP）中，貢獻給開源社群。

論文連結：
https://arxiv.org/abs/2205.13730

以上是面向長程式碼序列的 Transformer 模型最佳化方法，提升長程式碼場景效能的詳細內容。更多資訊請關注PHP中文網其他相關文章！