PaperlayerView exact revision on arXiv

Source-linked AI summary

Diet Code Is Healthy: Simplifying Programs for Pre-trained Models of Code

Zhaowei Zhang, Hongyu Zhang, Beijun Shen, Xiaodong Gu

arXiv:2206.14390v5cs.SE

TL;DR

Pre-trained code models are computationally expensive, especially as attention complexity grows with input length. DietCode analyzes CodeBERT’s attention and simplifies code inputs through pruning strategies that retain important tokens and statements. Across two downstream tasks, it achieves comparable performance with substantially lower computation.

  • Problem

    Pre-trained code models deliver strong software-engineering performance but impose high computational costs, including attention complexity that grows quadratically with sequence length.

    Source: §1 · p. 1, §4 · p. 5

  • Method

    DietCode simplifies code inputs using word dropout, frequency filtering, and attention-based pruning informed by CodeBERT’s learned attention.

    Source: §1 · p. 1

  • Results

    DietCode achieves comparable results to CodeBERT across two tasks while reducing fine-tuning and testing computational cost by nearly 40%.

    Source: §1 · p. 2, §8 · p. 11

  • Takeaways & Limitations

    Selecting important tokens and statements enables lightweight use of pre-trained code models without hampering reported model performance.

    Source: §1 · p. 2

    The evaluation covers Java and Python and only code search and code summarization, so applicability to other languages and software-engineering tasks remains unverified.

    Source: §6.5 · p. 10

Abstract

from arXiv · show

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 the input sequence. Our empirical analysis of 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 DietCode, which aims at lightweight leverage of large pre-trained models for source code. DietCode 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 reduction in the computational cost without hampering the model performance. Experimental results on two downstream tasks show that DietCodeBERT provides comparable results to CodeBERT with 40% less computational cost in fine-tuning and testing.

1 INTRODUCTION

Pre-trained code models achieve strong software-engineering performance but impose substantial computational costs. DietCode analyzes CodeBERT’s attention to simplify inputs, retaining comparable performance while reducing computation.

Motivation

  • CodeBERT has 125 million parameters and requires 28 hours to pre-train on 8.5M code, motivating computational simplification.Fine-tuning also remains costly because pre-trained models use large parameter counts and training datasets.Source: §1 · p. 1

Empirical Analysis

  • Keywords and data types receive the most attention, while method signatures and return statements are especially important statement types.These statement types reflect overall method functionality in the attention analysis.Source: §1 · p. 1

Method

  • DietCode prunes unimportant tokens and statements using word dropout, frequency filtering, and attention-based pruning.Its attention-based strategy selects high-attention tokens and statements, formulating statement selection as a 0-1 knapsack problem under a target length.Source: §1 · p. 1

Evaluation

  • The paper evaluates DietCode on code search and code summarization to assess computational efficiency and retained performance.These tasks are compared against the original pre-trained code model and SIVAND [34].Source: §1 · pp. 1–2

Experimental Results

  • Nearly 40% less computational cost in fine-tuning and testing accompanies DietCode’s comparable results to RoBERTa (code).The evaluation covers code search and code summarization, using relative length, FLOPs, and time cost alongside baseline comparisons.Source: §1 · pp. 1–2

Contributions

  • The contributions combine an analysis of CodeBERT’s critical information with a program simplification approach evaluated on two downstream tasks.The stated goal is to reduce computation cost while retaining comparable performance.Source: §1 · p. 2

2 BACKGROUND

Pre-trained language models use Transformer-based attention to represent sequences and support downstream adaptation. CodeBERT applies this architecture to programming and natural-language inputs through masked language modeling and replaced token detection.

2.1 Pre-Trained Language Models

  • Transformer-based pre-trained models use self-attention to relate tokens across positions and learn dependencies from distant tokens in parallel.Transformer layers combine multi-head self-attention with feed-forward networks in an encoder-decoder architecture.Source: §2.1 · p. 2
  • The Transformer’s attention mechanism produces token representations using learned query, key, and value parameters across multiple attention heads.The input sequence is represented through attention-based transformations whose head outputs are concatenated.Source: §2.1 · p. 2
  • Large pre-trained models consume substantial computational resources, with BERT-base and BERT-large containing 110M and 340M parameters, respectively.This resource demand motivates lightweight leverage of pre-trained models.Source: §2.1 · p. 2

2.2 CodeBERT

  • CodeBERT is a Transformer encoder model for natural and programming languages whose representations support code search and code summarization.It was proposed for software engineering tasks and captures semantic representations from programming languages.Source: §2.2 · p. 2
  • CodeBERT is pre-trained with masked language modeling and replaced token detection before downstream fine-tuning.MLM predicts masked tokens, while RTD uses a discriminator to identify whether tokens are original.Source: §2.2 · p. 2

3 EMPIRICAL ANALYSIS

The empirical analysis measures CodeBERT’s token and statement importance using averaged self-attention weights, then examines which code elements receive the most attention across Java and Python. CodeBERT emphasizes functional and structurally informative elements, while giving less attention to many control-flow and auxiliary statements.

3.1 Study Design

  • Averaged self-attention weights across CodeBERT’s layers and heads define token importance, using CodeSearchNet code snippets as input.Statement importance is then computed from token attentions rather than treated as a directly reusable dictionary entry.Source: §3.1 · p. 2, §3.1.1 · pp. 2–3

3.1.2 Computing Attentions of Statements.

  • Statements are grouped into 21 categories, including method signatures, variable declarations, and if conditions, to compare category-level attention.Function invocation and method signature are the most common categories in CodeSearchNet, while while-loop statements are relatively rare.Source: §3.1.2 · p. 3, Table 1 · p. 3

3.2.1 What does CodeBERT learn about code tokens? (RQ1).

  • Among Java tokens, boolean has the highest attention weight at 2.94𝑒-2, while coverage has the lowest at 5.43𝑒-5.The analysis summarizes tokens from more than 100,000 Java functions and removes tokens appearing fewer than 50 times for visualization.Source: §3.2.1 · p. 3, Figure 2 · p. 4
  • CodeBERT emphasizes Java keywords and data-oriented identifiers such as public, boolean, Map, List, and String because they convey code structure or key data structures.The keyword private receives the highest Java-keyword attention, 1.15𝑒-2, whereas interface and transient receive 7.14𝑒-4 and 9.62𝑒-4.Source: §3.2.1 · pp. 3–4
  • The same token-level pattern appears in Python, where return and def receive high attention and while ranks above if and for among branching and loop keywords.In a Java heatmap, functional tokens such as File and read receive high attention, while auxiliary tokens such as null and Exception receive lower attention.Source: §3.2.1 · pp. 3–4

3.2.2 What does CodeBERT learn about code statements? (RQ2).

  • Method signature receives the highest Java statement attention at 4.11𝑒-3, while case receives the lowest at 2.32𝑒-3.CodeBERT focuses more on statements indicating a method’s overall functionality, including method signatures and return statements.Source: §3.2.2 · p. 4, Figure 4 · p. 5
  • Arithmetic expressions and function invocations also receive relatively high attention because their operations, names, and parameters convey functionality or identify other method signatures.These statement types provide information that can be interpreted as sequences of operations or references to callable behavior.Source: §3.2.2 · p. 4
  • CodeBERT assigns relatively little attention to control-flow statements such as while, for, and case, indicating stronger learning of plain-text-like content than code structure.For Python, method signatures receive 6.86𝑒-3 while break and continue receive the lowest weights; variable initialization and an if statement are also low-attention in a Java example.Source: §3.2.2 · pp. 4–5

4 DIETCODE: PROGRAM SIMPLIFICATION FOR CODEBERT

DietCode simplifies CodeBERT inputs by pruning tokens and statements to reduce attention-related computation while preserving information useful for downstream tasks. It offers random, frequency-based, and attention-based pruning, with the latter selecting statements and tokens using learned attention.

4 DIETCODE: PROGRAM SIMPLIFICATION FOR CODEBERT

  • DietCode reduces CodeBERT’s input complexity by pruning unimportant tokens and statements while targeting a maximum sequence length.The problem is motivated by attention complexity that scales quadratically with sequence length.Source: §4 · p. 5

4.1 Word Dropout

  • Word dropout randomly removes tokens until the input meets a specified sequence length, with each token independently marked for retention or pruning.Its selection probability is p = L/|C|.Source: §4.1 · p. 5

4.2 Frequency Filtering

  • Frequency filtering keeps the most frequently occurring tokens, addressing random removal but potentially biasing selection toward common tokens.Token frequency is measured over the whole code corpus.Source: §4.2 · p. 5, §4.3 · p. 5

4.3 Attention-based Code Pruning

  • Attention-based pruning selects high-attention statements with a 0-1 knapsack procedure, then removes low-attention tokens until the target length is met.Min-max normalization multiplied by statement token count reduces bias toward short statements; the procedure uses two phases: statement selection and token pruning.Source: §4.3 · pp. 5–6
  • In the example, DietCode reduces a function from 118 to 60 tokens while retaining key semantic elements such as the method name, variable, and invoked function.The simplified code is not runnable but retains information intended for downstream tasks such as code search and summarization.Source: §4.3 · p. 6

5 EVALUATION

DietCode is evaluated on code search and summarization using CodeSearchNet, standard complexity and task metrics, and comparisons with pre-trained, traditional, and simplification baselines. Across experiments, it reduces computation while retaining comparable performance, with attention-based pruning generally strongest and roughly 60% relative length offering a useful trade-off.

5.1 Downstream Tasks

  • The evaluation covers code search and code summarization on CodeSearchNet, using Relative Length, FLOPs, time cost, MRR, and BLEU-4.CodeSearchNet contains comment-code pairs, and the experiments compare DietCode with CodeBERT, CodeT5, RoBERTa, traditional models, and SIVAND.Source: §5.1 · p. 6, §5.2 · p. 6, §5.3 · p. 7

5.2 Datasets

  • The study preprocesses code by parsing token sequences into statements and delexicalizing most string and numerical constants before evaluation.Java statements are split using brackets and semicolons, while Python uses symbols such as “def”, “=”, “:”, and “)”.Source: §5.2 · p. 7

5.3 Evaluation Metrics

  • Relative Length is the percentage of simplified-code tokens retained from the original, while FLOPs and execution time quantify computational cost.MRR evaluates code search and BLEU-4 evaluates code summarization.Source: §5.3 · p. 7

5.5 Experimental Results (RQ3)

  • 40% less computational cost accompanies comparable accuracy when DietCode reduces CodeBERT inputs to 60% relative length on code search and summarization.For Java code search, MRR is 0.71 versus CodeBERT's 0.74, while FLOPs fall from 16.99G to 10.19G; Java summarization reaches BLEU-4 17.29 while FLOPs fall from 24.328G to 15.325G.Source: §5.5 · pp. 8–9, Table 3 · p. 8, Table 4 · p. 8
  • Attention-based pruning outperforms dropout and frequency filtering on Java, reaching MRR 0.71 versus 0.68 and 0.66, respectively.For CodeT5 summarization, attention-based pruning reaches BLEU 19.25 versus 17.65 for dropout and 18.74 for frequency filtering; the advantage is less significant in Python.Source: §5.5 · p. 9
  • SIVAND cannot be run over the full datasets because its estimated processing time exceeds 30,000 hours, so its timing is extrapolated from 50 sampled functions without accuracy validation.The authors report that the sampled data are insufficient to fine-tune SIVAND and validate its accuracy.Source: §5.5 · p. 9

5.6 Effect of Relative Length (RQ4)

  • Performance drops sharply as relative length decreases, while attention pruning remains strongest; around 60% relative length provides the best reported trade-off.Above 70%, performance is comparable to original CodeBERT but savings are modest; below 50%, performance becomes deficient.Source: §5.6 · p. 9, Figure 6 · p. 9

5.7 Effect of Token and Statement Pruning (RQ5)

  • Statement pruning preserves performance better than token pruning as relative length decreases, suggesting statements carry more critical representational information.The reported explanation is that statement attention captures relationships among tokens, whereas token attention tends to measure local information.Source: §5.7 · p. 9, Figure 7 · p. 9

6 DISCUSSION

The discussion explains why simplified code can remain useful to CodeBERT, weighs DietCode’s efficiency trade-off, and identifies formulation choices, structural limitations, and validity boundaries.

6.1 How Can Pre-Trained Models Understand Simplified Code?

  • CodeBERT can use simplified code because it emphasizes the natural-language channel, including keywords, rather than requiring complete executable behavior.This aligns with programmers’ tendency to prioritize function names, method invocations, and variable names while skipping some structural information.Source: §6.1 · p. 10

6.2 What users may adopt the techniques?

  • DietCode offers a practical option for resource-limited users: slight performance drops are offset by roughly 40% savings in fine-tuning time.It still significantly outperforms traditional non-pretraining methods such as BiRNN and SelfAttn.Source: §6.2 · p. 10

6.3 The role of the 0-1 knapsack formulation

  • At 60% relative length, token-and-statement pruning reaches MRR=0.71 versus around 0.68 for statement-only pruning on Java code search.The same advantage appears in Java summarization, where BLEU is 17.3 versus 16.8, supporting the 0-1 knapsack formulation.Source: §6.3 · p. 10

6.4 Implications for Future Research

  • CodeBERT rarely recognizes grammar and structural information, motivating future methods that verbalize structural symbols before encoding.The discussion specifically identifies if, for, and while conditions as structures that are rarely considered.Source: §6.4 · p. 10
  • DietCode is easier to deploy than model compression or distillation because it prunes input programs without modifying the model.The authors present these model-level techniques as future directions rather than components of DietCode.Source: §6.4 · p. 10

6.5 Threats to Validity

  • The evidence is limited by excluded rare statement categories, parser errors, experiments in only Java and Python, and evaluation on code search and summarization.These issues may introduce dataset noise, produce different attention patterns in other languages, and leave applicability to other tasks unverified.Source: §6.5 · p. 10

7 RELATED WORK

Related work studies what pre-trained code models attend to and how programs can be simplified. DietCode differs by targeting critical tokens and statements and by avoiding repeated auxiliary-model evaluations.

7.1 Understanding Pre-trained Models of Code

  • DietCode focuses on critical tokens and statements learned by pre-trained models, then uses those findings to simplify programs.This differs from work that probes the whole model or captures broader structural knowledge.Source: §7.1 · pp. 10–11

7.2 Program Simplification

  • Traditional delta-debugging simplifiers repeatedly run an auxiliary deep learning model while splitting code fragments to preserve a target score.SIVAND [34] and P2IM [42] are described as examples of this approach.Source: §7.2 · p. 11
  • DietCode processes 10,000 functions within two minutes, whereas SIVAND requires hundreds of hours because it evaluates an auxiliary model at each iteration.The comparison highlights DietCode’s computational efficiency relative to iterative simplification methods.Source: §7.2 · p. 11

8 CONCLUSION

The paper analyzes CodeBERT’s critical code information and proposes DietCode to simplify inputs using attention-guided selection. Across two tasks, DietCode achieves comparable results with improved computational efficiency.

8 CONCLUSION

  • CodeBERT emphasizes keywords and data-relevant statements such as method signatures and return statements in Java and Python.The paper’s empirical analysis identifies these as critical information for its program representations.Source: §8 · p. 11
  • DietCode selects important statements and tokens to meet a target input length using attention weights and a 0-1 knapsack algorithm.This provides lightweight leverage of pre-trained models by simplifying the input code rather than changing the model.Source: §8 · p. 11
  • Across two downstream tasks, DietCode provides results comparable to CodeBERT while improving computational efficiency.The authors plan to study additional syntax and semantic information and model compression or distillation in future work.Source: §8 · p. 11

Related papers

Loading 2206.14390v5…