A plataforma de aprendizado de máquina PAI do Alibaba Cloud e a equipe do professor Gao Ming da East China Normal University publicaram em conjunto o modelo Transformer de atenção esparsa com reconhecimento de estrutura SASA no SIGIR2022. Este é um método de otimização do modelo Transformer para sequências de código longas e está comprometido em melhorar a efeito em cenários de código longo e desempenho. Como a complexidade do módulo de autoatenção aumenta exponencialmente com o comprimento da sequência, a maioria dos modelos de linguagem pré-treinados (PPLM) baseados em programação usa truncamento de sequência para processar sequências de código. O método SASA dispensa o cálculo de autoatenção e combina as características estruturais do código, melhorando assim o desempenho de tarefas de longa sequência e reduzindo a memória e a complexidade computacional.
Artigo: Tingting Liu, Chengyu Wang, Cen Chen, Ming Gao e Aoying Zhou. Entendendo linguagens de programação longas com atenção esparsa com reconhecimento de estrutura. SIGIR 2022
moldura do modelo
A figura a seguir mostra a estrutura geral do SASA:
Entre eles, o SASA inclui principalmente dois estágios: estágio de pré-processamento e estágio de treinamento do Sparse Transformer. Na etapa de pré-processamento, obtém-se a matriz de interação entre dois tokens, um é a matriz de frequência top-k e o outro é a matriz padrão AST. A matriz de frequência Top-k é usar o modelo de linguagem pré-treinado de código para aprender a frequência de interação de atenção entre tokens no corpus CodeSearchNet. A matriz de padrão AST é a árvore de sintaxe abstrata (Abstract Syntax Tree, AST) de o código de análise, que é obtido de acordo com a relação de conexão das árvores sintáticas.Informações de interação entre tokens. A fase de treinamento do Sparse Transformer usa o Transformer Encoder como a estrutura básica, substitui a autoatenção total pela autoatenção esparsa com reconhecimento de estrutura e executa o cálculo de atenção entre pares de token que estão em conformidade com um padrão específico, reduzindo assim a complexidade computacional.
A atenção esparsa da SASA inclui os quatro módulos a seguir:
- Atenção da janela deslizante: a autoatenção é calculada apenas entre os tokens na janela deslizante e as características do contexto local são preservadas. A complexidade computacional é ， é o comprimento da sequência, é o tamanho da janela deslizante.
- Atenção global: defina um determinado token global, esses tokens realizarão o cálculo de atenção com todos os tokens da sequência, de modo a obter a informação global da sequência, e a complexidade computacional é ， é o número de tokens globais.
- Atenção esparsa de top-k: a interação de atenção no modelo Transformer é esparsa e de cauda longa. Para cada token, apenas os tokens de top-k com a interação de atenção mais alta calculam a atenção e a complexidade é 。
- Atenção da estrutura com reconhecimento de AST: o código é diferente da sequência de linguagem natural e tem características estruturais mais fortes. Ao analisar o código em uma árvore sintática abstrata (AST), o intervalo de cálculo de atenção é determinado de acordo com a relação de conexão na árvore sintática .
Para se adaptar às características de computação paralela do hardware moderno, dividimos a sequência em vários blocos em vez de calcular em unidades de token. blocos de janela deslizante e blocos globais e blocos top-k e AST calculam a atenção e a complexidade computacional geral é ， é o tamanho do bloco.
每个sparse attention pattern 对应一个attention矩阵，以sliding window attention为例，其attention矩阵的计算为：
SASA伪代码：
实验结果
我们采用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/Eas…）中，贡献给开源社区。
参考文献
[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 para sequências mais longas. NeurIPS 2020