AST-MHSA : Code Summarization using Multi-Head Self-Attention

TL;DR

This paper introduces AST-MHSA, a novel transformer-based model that uses multi-head self-attention to efficiently extract semantic information from long ASTs for improved code summarization.

Contribution

The paper proposes a multi-head self-attention mechanism tailored for long ASTs, addressing computational challenges and enhancing semantic extraction in code summarization.

Findings

Improved summarization accuracy over baseline models

Reduced computational overhead compared to traditional transformer models

Effective extraction of semantic relations from long ASTs

Abstract

Code summarization aims to generate concise natural language descriptions for source code. The prevailing approaches adopt transformer-based encoder-decoder architectures, where the Abstract Syntax Tree (AST) of the source code is utilized for encoding structural information. However, ASTs are much longer than the corresponding source code, and existing methods ignore this size constraint by directly feeding the entire linearized AST into the encoders. This simplistic approach makes it challenging to extract truly valuable dependency relations from the overlong input sequence and leads to significant computational overhead due to self-attention applied to all nodes in the AST. To address this issue effectively and efficiently, we present a model, AST-MHSA that uses multi-head attention to extract the important semantic information from the AST. The model consists of two main components: an encoder and a decoder. The encoder takes as input the abstract syntax tree (AST) of the code and generates a sequence of hidden states. The decoder then takes these hidden states as input and generates a natural language summary of the code. The multi-head attention mechanism allows the model to learn different representations of the input code, which can be combined to generate a more comprehensive summary. The model is trained on a dataset of code and summaries, and the parameters of the model are optimized to minimize the loss between the generated summaries and the ground-truth summaries.