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

TL;DR

LeanCode addresses the token-length bottleneck in pre-trained code models by introducing context-aware, category-local attention-guided token removal. It leverages CLS attention for classification and encoder-decoder attention for sequence-to-sequence tasks to rank token importance, formulating a combinatorial optimization to remove a fixed number of tokens per snippet. Empirical results on code search and code summarization show LeanCode outperforms DietCode and SlimCode up to substantial margins while delivering sizable inference-time savings, with modest training overhead and evidence of cross-model transfer via GPT-4o. The work demonstrates that context-aware token pruning can meaningfully accelerate code-understanding models without large performance losses, and it establishes a path toward broader language support and model generalization.

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.

Figures (7)

• Figure 1: The variance of the top and bottom 10 tokens of CodeBERT for code search and summarization.
• Figure 2: The example of LeanCode for code simplification.
• Figure 3: CLS and Encoder-Decoder attention scores based on the example in Fig. \ref{['fig:motiv-example']}.
• Figure 4: An Example of Bubble-Sort Code
• Figure 5: Example of accumulated self attention in $\textsc{DietCode}\xspace$ on Fig. \ref{['fig:motiv-example']}. Dotted lines with different colors represent the self-attentions of different tokens. $s_{len}$ is the accumulated self-attention score of the token 'len'.