Graph Conditioned Sparse-Attention for Improved Source Code Understanding

TL;DR

This paper introduces a graph-conditioned sparse-attention model for source code understanding that efficiently models long-range dependencies, achieving state-of-the-art results with linear memory and inference time growth.

Contribution

The work presents a novel sparse self-attention mechanism conditioned on graph structure, reducing computational costs while improving performance on code understanding tasks.

Findings

Achieves state-of-the-art BLEU, METEOR, ROUGE-L scores in code summarization.

Attains near state-of-the-art accuracy in variable misuse detection.

Memory and inference time grow linearly with input length.

Abstract

Transformer architectures have been successfully used in learning source code representations. The fusion between a graph representation like Abstract Syntax Tree (AST) and a source code sequence makes the use of current approaches computationally intractable for large input sequence lengths. Source code can have long-range dependencies that require larger sequence lengths to model effectively. Current approaches have a quadratic growth in computational and memory costs with respect to the sequence length. Using such models in practical scenarios is difficult. In this work, we propose the conditioning of a source code snippet with its graph modality by using the graph adjacency matrix as an attention mask for a sparse self-attention mechanism and the use of a graph diffusion mechanism to model longer-range token dependencies. Our model reaches state-of-the-art results in BLEU, METEOR, and ROUGE-L metrics for the code summarization task and near state-of-the-art accuracy in the variable misuse task. The memory use and inference time of our model have linear growth with respect to the input sequence length as compared to the quadratic growth of previous works.