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.