Diet Code is Healthy: Simplifying Programs for Pre-Trained Models of Code

Abstract

Pre-trained code representation models such as CodeBERT have demonstrated superior performance in a variety of software engineering tasks, yet they are often heavy in complexity, quadratically with the length of input sequence. Our empirical analysis on CodeBERT's attention reveals that CodeBERT pays more attention to certain types of tokens and statements such as keywords and data-relevant statements. Based on these findings, we propose Diet-CodeBERT, which aims at lightweight leverage of large pre-trained models for source code. DietCodeBERT simplifies the input program of CodeBERT with three strategies, namely, word dropout, frequency filtering, and an attention-based strategy which selects statements and tokens that receive the most attention weights during pre-training. Hence, it gives a substantial reduce in the computational cost without hampering the model performance. Results on two downstream tasks show that DietCodeBERT provides comparable results as CodeBERT with 40% less computational cost in fine-tuning and testing. CCS CONCEPTS• Computing methodologies → Natural language processing.

On the Reliability and Explainability of Automated Code Generation Approaches

“…Specifically, the higher the attention weights, the more attention that is paid by the model. Therefore, there have been many prior studies that employ the attention weights of pre-trained programming language models to explain model predictions [46,49,55]. Prior studies calculate the feature importance of each token by averaging the attention weights of all layers and heads.…”

“…Actually, pre-trained code generation models contain multiple encoder and decoder layers. Although averaging attention weights to explain encoder-based code models is widely employed by prior works [46,49,55], Wan et al [46] have presented the great variability between different layers. However, it remains unclear how to determine which attention weights are more important for model inference.…”

“…Previous research has already investigated applying explainable AI approaches, despite not being well-suited, on automated code generation. There have been many works that explore the attention mechanism to explain NMT models of code [3,32,46,55]. Zhang et al [55] employed attention mechanisms to dig into critical statements and tokens learned by pre-trained NMT models in code search and code summarization.…”

“…There have been many works that explore the attention mechanism to explain NMT models of code [3,32,46,55]. Zhang et al [55] employed attention mechanisms to dig into critical statements and tokens learned by pre-trained NMT models in code search and code summarization. Paltenghi et al [32] explored to what extent the attention weights of NMT models match the reasoning of skilled humans in code summarization.…”