PaperlayerView exact revision on arXiv

Source-linked AI summary

Pre-training by Predicting Program Dependencies for Vulnerability Analysis Tasks

Zhongxin Liu, Zhijie Tang, Junwei Zhang, Xin Xia, Xiaohu Yang

arXiv:2402.00657v1cs.SE

TL;DR

Existing code pre-training objectives often ignore semantic dependencies, although such structure matters for vulnerability analysis. This paper trains PDBERT with objectives that predict dependencies from source code, and reports stronger results on three vulnerability tasks alongside high dependency-prediction scores.

  • Problem

    Existing pre-training objectives ignore code structure or focus on syntax, and the feasibility and benefits of learning to analyze code dependencies for vulnerability analysis remain uninvestigated.

    Source: §1 · p. 2

  • Method

    PDBERT is pre-trained with objectives that predict statement-level control and token-level data dependencies from source code, alongside masked language modeling.

    Source: §1 · p. 2

  • Results

    PDBERT outperforms the best-performing baselines by 5.9%-9.0% across vulnerability detection, classification, and assessment.

    Source: §1 · p. 2

  • Takeaways & Limitations

    PDBERT can directly analyze control and data dependencies in partial and complete functions, as well as support vulnerability-analysis tasks.

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

    PDBERT cannot properly analyze snippets over 512 tokens for dependence analysis, and truncates long snippets for downstream tasks.

    Source: §6.1 · p. 9

Abstract

from arXiv · show

Vulnerability analysis is crucial for software security. This work focuses on using pre-training techniques to enhance the understanding of vulnerable code and boost vulnerability analysis. The code understanding ability of a pre-trained model is highly related to its pre-training objectives. The semantic structure, e.g., control and data dependencies, of code is important for vulnerability analysis. However, existing pre-training objectives either ignore such structure or focus on learning to use it. The feasibility and benefits of learning the knowledge of analyzing semantic structure have not been investigated. To this end, this work proposes two novel pre-training objectives, namely Control Dependency Prediction (CDP) and Data Dependency Prediction (DDP), which aim to predict the statement-level control dependencies and token-level data dependencies, respectively, in a code snippet only based on its source code. During pre-training, CDP and DDP can guide the model to learn the knowledge required for analyzing fine-grained dependencies in code. After pre-training, the pre-trained model can boost the understanding of vulnerable code during fine-tuning and can directly be used to perform dependence analysis for both partial and complete functions. To demonstrate the benefits of our pre-training objectives, we pre-train a Transformer model named PDBERT with CDP and DDP, fine-tune it on three vulnerability analysis tasks, i.e., vulnerability detection, vulnerability classification, and vulnerability assessment, and also evaluate it on program dependence analysis. Experimental results show that PDBERT benefits from CDP and DDP, leading to state-of-the-art performance on the three downstream tasks. Also, PDBERT achieves F1-scores of over 99% and 94% for predicting control and data dependencies, respectively, in partial and complete functions.

1 INTRODUCTION

The paper targets vulnerability analysis with pre-training objectives that teach models to analyze code dependencies from source code. PDBERT improves vulnerability-task results and predicts dependencies in partial and complete functions.

1 INTRODUCTION

  • Deep-learning methods automate vulnerability detection, classification, patch identification, and assessment, motivated by the growing number of vulnerabilities.The paper focuses on improving vulnerable-code understanding through pre-training.Source: §1 · pp. 1–2
  • Prior pre-training objectives either ignore code structure or target syntax, while dependency-aware approaches rely on extracted structure and may require correctly parsed code.The paper notes that partial code is useful in settings such as vulnerability detection in snippets shared on Stack Overflow.Source: §1 · p. 2
  • Control Dependency Prediction and Data Dependency Prediction train models to predict statement-level control and token-level data dependencies from source code alone.A masking strategy filters impossible token pairs to address DDP's data imbalance.Source: §1 · p. 2
  • PDBERT outperforms the best-performing baselines by 5.9%-9.0% across vulnerability detection, classification, and assessment.The authors report these gains for the three downstream tasks.Source: §1 · p. 2
  • Dependency-prediction F1-scores exceed 99% for control dependencies and 95% for partial-function data dependencies, and 99% and 94%, respectively, for complete functions.PDBERT's throughput is 23 times higher than Joern's, with the authors noting it suits settings tolerant of some low-level imprecision where throughput matters.Source: §1 · p. 2
  • PDBERT combines CDP, DDP, and masked language modeling, learning dependence-analysis knowledge that supports partial-code analysis and diverse downstream tasks.The paper presents this as more general than learning to use extracted dependencies and as complementary to static analysis tools on partial code.Source: §1 · p. 2

2 PRELIMINARY

This section introduces program dependence graphs and explains their relevance to partial-code analysis. A PDG represents control and data dependencies, while partial-code support is useful for assessing vulnerable snippets before they are reused.

2.1 Program Dependence Graph

  • A program dependence graph represents control and data dependencies as distinct edge types between statement or predicate nodes.Control edges encode predicate-to-statement control relationships; data edges encode variable-labeled def-use relationships.Source: §2.1 · p. 3
  • Control-dependency edges indicate the conditions on which a statement depends.In Figure 1, these are shown as dotted green lines.Source: §2.1 · p. 3, Figure 1 · p. 3
  • Data-dependency edges represent def-use relationships and label the variable defined at the source and used at the target.In Figure 1, these are shown as solid red lines.Source: §2.1 · p. 3, Figure 1 · p. 3

2.2 Usage Scenario of Handling Partial Code

  • Analyzing partial code can help identify vulnerable Stack Overflow snippets before copied code is distributed into production software.The paper cites evidence that insecure Stack Overflow snippets have been copied into popular Android applications and that vulnerable snippets affected GitHub projects.Source: §2.2 · p. 3

3 APPROACH

The approach section describes PDBERT's input representation and the stages that follow it, including dependency ground-truth construction, pre-training tasks, and model usage. Its input pipeline tokenizes source code and obtains contextual token embeddings with a Transformer.

3.1 Input Representation

  • PDBERT takes source code as input and uses a subword tokenizer, such as BPE, to produce tokens preceded by a [CLS] token.The resulting sequence is denoted T = [t_cls, t_1, ..., t_n].Source: §3.1 · p. 3
  • A multi-layer Transformer processes the token sequence to produce a contextual embedding for each token.The paper denotes the contextual embeddings as H_t.Source: §3.1 · p. 3
  • The representation records each token's start and end character indices in the source code.Source: §3.1 · p. 3

3.2 Program Dependency Extraction

The paper constructs dependency targets from program graphs, representing control dependencies at statement level and data dependencies at a finer variable-aware level. It uses Joern to extract the graph structure from source code.

3.2 Program Dependency Extraction

  • Joern constructs a program dependence graph and abstract syntax tree from source code, which the method combines into a joint graph for extracting dependencies.Source: §3.2 · p. 3

3.2.1 Control Dependency Extraction.

  • Control dependencies are encoded in a binary matrix whose entries indicate whether one PDG statement depends on a predicate.The matrix is sized using a maximum of 50 PDG nodes by default.Source: §3.2.1 · p. 3
  • Data dependencies are represented at variable-aware granularity by pairing AST nodes for the variable's definition and use.Each dependency is represented as D_i = [s_i, e_i], the start and end AST nodes identified from its PDG edge.Source: §3.2.2 · p. 3

3.3 Token-Level Data Dependency Construction

The authors convert AST-level data-dependency edges into token-level ground truth for DDP. Each dependency becomes one edge between the first tokens mapped to its source and destination nodes.

3.3 Token-Level Data Dependency Construction

  • Each AST node is mapped to overlapping subword tokens by matching character spans, so code elements split into multiple tokens retain their token sequence.For example, each temp_flag maps to [Ġtemp, _, flag].Source: §3.3 · p. 4
  • Rather than connect every token pair in the source and destination sequences, each dependency is represented by an edge between their first tokens.This choice avoids producing many unnecessary edges.Source: §3.3 · p. 4
  • The authors set the corresponding source-to-destination entry in G_d to 1 in an initially zero-filled matrix, which serves as DDP ground truth.For the illustrated edge, G_d[1,13] is set to 1.Source: §3.3 · p. 4

3.4 Pre-training Tasks

PDBERT combines masked-token reconstruction with control- and data-dependency prediction to learn code syntax and semantic structure. Its DDP objective addresses sparse token-pair labels by masking non-identifier pairs.

3.4 Pre-training Tasks

  • MLM captures code naturalness and some syntax, while CDP and DDP target semantic structure and are described as complementary tasks.The authors say the dependency objectives strengthen attention between computationally related code elements.Source: §3.4 · p. 4

3.4.1 Masked Language Model (MLM).

  • MLM reconstructs randomly masked input tokens from their contextual embeddings, and the authors mask 15% of tokens with additional replacements and unmasking.Of sampled tokens, 10% are replaced with random tokens and another 10% are left unmasked.Source: §3.4.1 · p. 4

3.4.2 Statement-Level Control Dependency Prediction (CDP).

  • CDP predicts directed control dependencies between PDG nodes using node representations derived from averaged token embeddings and a bilinear predictor.The asymmetric prediction reflects that control dependencies are directional.Source: §3.4.2 · p. 4

3.4.3 Token-Level Data Dependency Prediction (DDP).

  • DDP predicts token-level data dependencies, but its pairwise label matrix is large and sparse, so direct prediction can encourage constant no-dependency predictions.With 512 input tokens, the matrix has over 262,144 elements.Source: §3.4.3 · p. 5
  • To address sparsity, DDP masks matrix entries involving non-identifiers, using lexer token types mapped to the model tokens by character-span overlap.The paper notes that stronger priors, such as predicting only between identical identifiers, are possible but uses token types.Source: §3.4.3 · p. 5
  • The authors say masking can be disabled for end-to-end dependence inference, unlike explicitly supplying identifier pairs, which requires additional tools.They also state that masking implicitly guides attention toward identifier pairs and can help the model learn code syntax.Source: §3.4.3 · p. 5

3.5 The Usages of PDBERT

PDBERT uses source code alone to produce contextual embeddings, which can support direct dependence analysis or task-specific fine-tuning. The authors hypothesize that learned dependence knowledge can benefit downstream tasks sensitive to program dependencies.

3.5 The Usages of PDBERT

  • PDBERT accepts a partial or complete code function and produces contextual embeddings without requiring parsing or PDG construction during downstream fine-tuning.Task-specific headers convert embeddings into outputs such as classifications or token sequences.Source: §3.5 · p. 5
  • After pre-training, CDP and DDP headers can be attached to infer statement-level control and token-level data dependencies directly from a code snippet.For downstream tasks, a new task-specific header is fine-tuned together with PDBERT.Source: §3.5 · p. 5
  • The authors hypothesize that PDBERT absorbs dependence knowledge during pre-training and that this knowledge can boost downstream tasks sensitive to program dependencies.They argue the pre-training-task headers are not used downstream and therefore have no side effects.Source: §3.5 · p. 5

4 PRE-TRAINING SETUP

The authors pre-train PDBERT on C/C++ functions and configure it as a 12-layer Transformer initialized from CodeBERT. The setup limits dependency analysis to the truncated input sequence.

4.1 Pre-training Dataset

  • The pre-training dataset contains over 2.28M C/C++ functions from GitHub projects and the Draper dataset, split into about 1.9M training, 155.0K validation, and 60.4K test functions.Functions that Joern could not analyze within a 30-minute timeout were filtered; validation and test data were deduplicated against training data.Source: §4.1 · p. 6

4.2 Pre-training Model Configurations

  • PDBERT uses a 12-layer Transformer with 768-dimensional hidden states and a maximum sequence length of 512, initialized with CodeBERT parameters.Long snippets are truncated, and CDP and DDP cover only dependencies within the input tokens.Source: §4.2 · p. 6
  • Pre-training jointly minimizes the three task losses and takes approximately 72 hours for 10 epochs on four Nvidia RTX 3090 GPUs.The loss weights are set to 5, 20, and 1 based on validation-set grid search.Source: §4.2 · p. 6

5 EVALUATION

The evaluation tests PDBERT on direct dependence analysis of partial and complete functions and on three function-level vulnerability tasks. It compares pre-training objectives and models against pre-trained and non-pre-trained baselines.

5 EVALUATION

  • Extrinsic evaluation fine-tunes models for vulnerability detection, classification, and assessment on complete functions because some baselines cannot handle partial functions.For intrinsic evaluation, PDBERT is tested directly on dependencies in partial and complete functions without further training or fine-tuning.Source: §5 · p. 6, §5.1.1 · p. 6

5.1.3 Results.

  • On partial code, PDBERT achieves about 99% F1 for CDP and over 95% for DDP across tested K values, with overall F1 above 96%.On complete functions, control, data, and overall F1-scores exceed 99%, 94%, and 96%, respectively, across LOC ranges.Source: §5.1.3 · pp. 6–7
  • PDBERT derives a complete function’s program-dependence graph in 19 ms on a GPU, versus over 460 ms for Joern on two CPUs; the comparison is not entirely fair.The authors frame this as a throughput advantage for use cases tolerant of some low-level imprecision, not as evidence that PDBERT replaces Joern in all use cases.Source: §5.1.3 · p. 7

5.4.2 Experimental Setup.

  • The detection comparison includes ReVeal, Big-Vul, and Devign, while classification uses 7.6K vulnerable functions across 14 CWE categories.The classification dataset excludes CWE categories with fewer than 100 instances.Source: §5.3.2 · p. 7, §5.4.2 · p. 8

5.3.3 Results.

  • PDBERT improves the strongest vulnerability-detection baseline by 4.3% on ReVeal, 9.0% on Big-Vul, and 5.6% on Devign in F1-score or Accuracy, and consistently outperforms all baselines.MLM+CDP+DDP outperforms MLM on all three datasets; the combined objectives outperform either CDP or DDP alone.Source: §5.3.3 · p. 8

5.4.3 Results.

  • PDBERT improves CodeBERT by 11.7% in Macro F1, 9.1% in Weighted F1, and 12.2% in MCC for vulnerability classification.The pre-trained models outperform the non-pre-trained Bi-LSTM and Transformer baselines by large margins.Source: §5.4.3 · p. 8

5.5 Vulnerability Assessment

  • Assessment predicts Availability, Confidentiality, Integrity, and Access Complexity from vulnerable functions; its dataset contains 9.9K functions with valid CVSS scores.The authors note that CVSS assessment can require project-specific context that is difficult to find in a vulnerable function, limiting gains from improved code understanding.Source: §5.5.1 · p. 9, §5.5.2 · p. 9, §5.5.3 · p. 9

5.5.3 Results.

  • For vulnerability assessment, MLM+CDP+DDP improves GraphCodeBERT by 2.8% in Macro F1 and 3.3% in MCC on average.MLM+CDP performs better than MLM+CDP+DDP on Access Complexity; the paper concludes CDP is more important than DDP for this task.Source: §5.5.3 · p. 9

6 DISCUSSION

The discussion considers limits on input length and statement-level control-flow coverage, alongside threats involving language generalization, Joern-derived dependence labels, and possible data contamination.

6.1 Limitations

  • PDBERT’s 512-token limit prevents proper dependence analysis of longer snippets and leads it to truncate long inputs for downstream tasks.Over 83% of the test-set C/C++ functions fit the limit; the authors report downstream gains despite possible performance effects from truncation. The limit comes from the Transformer architecture, and longer context requires more computation resources and advanced hardware.Source: §6.1 · pp. 9–10
  • CDP predicts statement-level control dependencies but does not handle control flow within a statement, such as ternary operators.The authors identify this as an area for future improvement.Source: §6.1 · p. 10

6.2 Threats to Validity

  • PDBERT is pre-trained on C/C++ and may not suit other languages, although CDP and DDP are language-agnostic.The authors state that practitioners can pre-train models with these objectives for other or multiple languages.Source: §6.2 · p. 10
  • Joern supplies the PDGs used as CDP and DDP ground truth, so PDBERT is expected to predict Joern’s dependencies and may share its limitations.The authors argue that Joern’s limitations do not affect their conclusions, citing its use and their intrinsic and extrinsic evaluations.Source: §6.2 · p. 10
  • Pre-training and fine-tuning data may overlap because both come from open-source repositories, but task labels are unknown during pre-training and the authors consider this threat limited.The discussion notes prior work [10] reporting that data contamination may have little impact on pre-trained-model performance.Source: §6.2 · p. 10

7 RELATED WORK

Prior code models use objectives for natural-language tasks, code syntax, or functional similarity, while PDBERT targets code’s semantic structure. Related vulnerability-analysis approaches often rely on extracted dependencies or task-specific training, unlike PDBERT’s source-code-based dependence pre-training.

7.1 Pre-Trained Code Models

  • PDBERT targets code’s semantic structure, whereas prior encoder-only models use natural-language objectives, AST prediction, or contrastive learning for functional similarity.The paper distinguishes these goals from its pre-training objectives for capturing semantic structure.Source: §7.1 · p. 10
  • Unlike Code-MVP [62] and GraphCodeBERT [26], PDBERT learns control and data dependencies together for end-to-end analysis from source code, including unparsable code.The paper says this design supports direct dependence analysis and yields knowledge it considers more general for downstream tasks.Source: §7.1 · pp. 10–11

7.2 Deep Learning for Vulnerability Analysis

  • Many vulnerability-detection approaches extract PDGs before neural modeling [41] [44] [68], whereas PDBERT needs no dependency input and can handle partial code.The paper contrasts task-specific learning to use dependencies with PDBERT’s pre-training to learn dependence-analysis knowledge for multiple vulnerability-analysis tasks.Source: §7.2 · p. 11
  • Prior vulnerability-classification studies mostly use curated descriptions [51], while assessment work often uses descriptions or requires identified vulnerable statements or known fixes [36].The paper also notes μVulDeePecker as a source-code CWE classifier and DeepCVA [38] as a commit-level assessment approach.Source: §7.2 · p. 11

8 CONCLUSION AND FUTURE WORK

The paper proposes CDP and DDP to train models to analyze program dependencies from source code, and reports strong dependence-analysis and vulnerability-task results. It identifies encoder-decoder pre-training, more tasks, multilingual data, and combining learned objectives with extracted dependencies as future directions.

8 CONCLUSION AND FUTURE WORK

  • CDP and DDP predict statement-level control and token-level data dependencies from source code, and PDBERT achieves F1-scores over 99% and 94%, respectively, in partial and complete functions.Source: §8 · p. 11
  • PDBERT is more effective than state-of-the-art baselines on vulnerability detection, classification, and assessment.The reported results indicate that it benefits from CDP and DDP.Source: §8 · p. 11
  • Future work includes pre-training encoder-decoder models, evaluating more downstream tasks, and using a multilingual corpus.The authors also propose investigating whether adding explicitly extracted dependencies improves performance when static analyzers can extract them successfully.Source: §8 · p. 11
Loading 2402.00657v1…