LEANCODE: Understanding Models Better for Code Simplification of Pre-trained Large Language Models

TL;DR

LeanCode introduces a novel method for code simplification in large language models by using context-aware attention scores to selectively remove tokens, significantly improving efficiency in code search and summarization tasks.

Contribution

The paper presents a new approach that leverages attention scores for token importance to reduce computational complexity in large language models for code.

Findings

Outperforms DietCode and Slimcode in code search and summarization tasks.

Achieves 60% and 16% improvements in code search.

Achieves 29% and 27% improvements in code summarization.

Abstract

Large Language Models for code often entail significant computational complexity, which grows significantly with the length of the input code sequence. We propose LeanCode for code simplification to reduce training and prediction time, leveraging code contexts in utilizing attention scores to represent the tokens' importance. We advocate for the selective removal of tokens based on the average context-aware attention scores rather than average scores across all inputs. LeanCode uses the attention scores of `CLS' tokens within the encoder for classification tasks, such as code search. It also employs the encoder-decoder attention scores to determine token significance for sequence-to-sequence tasks like code summarization. Our evaluation shows LeanCode's superiority over the SOTAs DietCode and Slimcode, with improvements of 60% and 16% for code search, and 29% and 27% for code summarization, respectively.