Improving Retrieval-Augmented Code Comment Generation by Retrieving for Generation
Hanzhen Lu, Zhongxin Liu
TL;DR
This work tackles Retrieval-Augmented Comment Generation by addressing the misalignment between retrievers and generators. It introduces JointCom, a joint training framework where a dense retriever and a seq2seq generator are optimized together using a top-$k$ exemplar weighting strategy, guiding the retriever to surface exemplars that most improve generation quality. Empirical results on JCSD and PCSD show substantial improvements over state-of-the-art baselines across five metrics, complemented by human evaluations favoring JointCom for naturalness, informativeness, and usefulness. The approach demonstrates robustness and generalizability, with feasibility extended to LLMs via in-context retrieval, and is accompanied by a replication package for reproducibility.
Abstract
Code comment generation aims to generate high-quality comments from source code automatically and has been studied for years. Recent studies proposed to integrate information retrieval techniques with neural generation models to tackle this problem, i.e., Retrieval-Augmented Comment Generation (RACG) approaches, and achieved state-of-the-art results. However, the retrievers in previous work are built independently of their generators. This results in that the retrieved exemplars are not necessarily the most useful ones for generating comments, limiting the performance of existing approaches. To address this limitation, we propose a novel training strategy to enable the retriever to learn from the feedback of the generator and retrieve exemplars for generation. Specifically, during training, we use the retriever to retrieve the top-k exemplars and calculate their retrieval scores, and use the generator to calculate a generation loss for the sample based on each exemplar. By aligning high-score exemplars retrieved by the retriever with low-loss exemplars observed by the generator, the retriever can learn to retrieve exemplars that can best improve the quality of the generated comments. Based on this strategy, we propose a novel RACG approach named JOINTCOM and evaluate it on two real-world datasets, JCSD and PCSD. The experimental results demonstrate that our approach surpasses the state-of-the-art baselines by 7.3% to 30.0% in terms of five metrics on the two datasets. We also conduct a human evaluation to compare JOINTCOM with the best-performing baselines. The results indicate that JOINTCOM outperforms the baselines, producing comments that are more natural, informative, and useful.