Tram: A Token-level Retrieval-augmented Mechanism for Source Code Summarization (2024)

Abstract: Automatically generating human-readable text describing the functionality of a program is the intent of source code summarization. Although neural language models achieve significant performance in this field, they are limited by their inability to access external knowledge. To address this limitation, an emerging trend is combining neural models with external knowledge through retrieval methods. Previous methods have relied on the sentence-level retrieval paradigm on the encoder side. However, this paradigm is coarse-grained, noise-filled and cannot directly take advantage of the high-quality retrieved summary tokens on the decoder side. In this paper, we propose a fine-grained Token-level retrieval-augmented mechanism (Tram) on the decoder side rather than the encoder side to enhance the performance of neural models and produce more low-frequency tokens in generating summaries. Furthermore, to overcome the challenge of token-level retrieval in capturing contextual code semantics, we also propose integrating code semantics into individual summary tokens. The results of extensive experiments and human evaluation show that our token-level retrieval-augmented approach significantly improves performance and is more interpretable.

## introduction with software functions becoming more comprehensive and complex, it becomes a heavy burden for developers to understand software. it has been reported that nearly 90% @xcite of effort is used for maintenance, and much of this effort is spent on understanding the maintenance task and related software source codes. source code summary as a natural language is indispensable in software since humans can easily read and understand it, as shown in if ((na // 4) != (nb // 4)): @xmath0: @xmath1, end, @xmath2.is_valid) else: with the development of language models and the linguistic nature of source code, researchers explored seq2seq architecture, such as recurrent neural networks to generate summaries @xcite @xcite . soon afterward, transformer-based models @xcite @xcite were proposed, outperforming previous rnn-based models by a large margin. recently, many approaches have been proposed to leverage the structural properties of source code, such as abstract syntax tree (ast) and program dependency graph (pdg). current structure-aware methods typically either fuse structural information in a hybrid manner @xcite @xcite @xcite , or use a structured-guided way @xcite @xcite @xcite . although these methods have shown promising results, they primarily focus on leveraging the information within the code to obtain richer code representation without fully utilizing the potential of the available human-written code-summary pairs. in order to leverage external existing highquality code and the corresponding summary instances, recent works @xcite @xcite have proposed a retrieval augmented approach. their unified paradigm involves sentence-level retrieval, which uses text similarity metrics or code semantic similarity metrics to retrieve the most similar code snippet from a code repository for the given input code snippet. the retrieved code snippet and its corresponding summary are either directly concatenated with the input code snippet or semantically enhanced to augment the input code snippet on the encoder side. however, the granularity of sentence-level retrieval methods poses challenges. specifically, they can erroneously retrieve and incorporate code snippets that, while syntactically similar, are semantically distinct or those that only bear partial semantic resemblance. the unintended noise introduced through such mismatches can adversely affect the generation performance, especially for lowfrequency tokens. moreover, code summarization is essentially a generative task, the decoder autoregressively generates the summary tokens. however, previous sentence-level retrieval-augmented methods neglect to fuse the retrieved information on the decoder side, only doing so on the encoder side, which will result in the utilization pattern being indirect and insufficient. these limitations have inspired us to explore a more fine-grained and sufficient retrieval approach on the summary generation process. in order to achieve the purpose of retrieving semantic similar summary tokens on the decoder side, we first construct a datastore to store the summary tokens and corresponding representations through a pretrained base model offline. meanwhile, to overcome the challenge of not fully utilizing code semantics on the encoder side when retrieving on the decoder side, we intelligently fuse summary token representation with code token representation and ast node representation with attention weight. this approach fully considers contextual code semantics associated with summary tokens. then, at each generation step, the fused summary token representation is used to retrieve the top-k most similar tokens. as illustrated in to facilitate future research, we have made our code publicly available foot_0 . overall, the main contributions of this paper can be outlined as follows: k-nearest-neighbor machine translation. recently, non-parametric methods have been successfully applied to neural machine translation @xcite @xcite . these approaches complement advanced nmt models with external memory to alleviate the performance degradation in domain adaption. compared to these works, we have fully accounted for the code's inherent structure and have intelligently integrated code semantics into the retrieval process. additionally, we demonstrate how tram integrates with sentence-level retrieval methods. ## conclusion in this paper, we propose a novel token-level retrieval-augmented mechanism for source code summarization. by a well-designed fine-grained retrieval pattern, tram can effectively incorporate external human-written code-summary pairs on the decoder side. extensive experiments and human evaluation show that tram not only significantly improves performance but also generates more lowfrequency tokens and enhances interpretability.