Understanding Long Programming Languages with Structure-Aware Sparse Attention

TL;DR

This paper introduces SASA, a structure-aware sparse attention mechanism that enables efficient processing of long code sequences in programming language models, improving performance on code understanding tasks.

Contribution

The paper proposes a novel attention mechanism combining top-k sparse attention with AST-based structure awareness to handle long code sequences efficiently.

Findings

SASA outperforms baseline models on CodeXGLUE tasks.

SASA reduces computational complexity for long code sequences.

Incorporating AST structures improves code understanding accuracy.

Abstract

Programming-based Pre-trained Language Models (PPLMs) such as CodeBERT have achieved great success in many downstream code-related tasks. Since the memory and computational complexity of self-attention in the Transformer grow quadratically with the sequence length, PPLMs typically limit the code length to 512. However, codes in real-world applications are generally long, such as code searches, which cannot be processed efficiently by existing PPLMs. To solve this problem, in this paper, we present SASA, a Structure-Aware Sparse Attention mechanism, which reduces the complexity and improves performance for long code understanding tasks. The key components in SASA are top-$k$ sparse attention and Abstract Syntax Tree (AST)-based structure-aware attention. With top-$k$ sparse attention, the most crucial attention relation can be obtained with a lower computational cost. As the code structure represents the logic of the code statements, which is a complement to the code sequence characteristics, we further introduce AST structures into attention. Extensive experiments on CodeXGLUE tasks show that SASA achieves better performance than the competing baselines.