[Paper Review] Using Machine Learning for Vulnerability Detection and Classification

Uri Alon, Shaked Brody+1 more|arXiv (Cornell University)|Aug 4, 2018

Software Engineering ResearchComputer Science452 citations

TL;DR

This paper introduces CODE2SEQ, a novel neural model that generates natural language sequences from source code by encoding code snippets as compositional paths in their abstract syntax trees (ASTs), using LSTMs to embed paths and attention mechanisms to attend over them during decoding. The model significantly outperforms prior code-specific and state-of-the-art neural machine translation models on code summarization and captioning tasks across two languages and four datasets.

ABSTRACT

The work described in this paper aims at developing a machine learning based tool for automatic identification of vulnerabilities on programs (source, high level code), that uses an abstract syntax tree representation. It is based on FastScan, using code2seq approach. Fastscan is a recently developed system aimed capable of detecting vulnerabilities in source code using machine learning techniques. Nevertheless, FastScan is not able of identifying the vulnerability type. In the presented work the main goal is to go further and develop a method to identify specific types of vulnerabilities. As will be shown, the goal will be achieved by optimizing the model’s hyperparameters, changing the method of preprocessing the input data and developing an architecture that brings together multiple models to predict different specific vulnerabilities. The preliminary results obtained from the training stage, are very promising. The best f1 metric obtained is 93% resulting in a precision of 90% and accuracy of 85%, according to the performed tests and regarding a trained model to predict vulnerabilities of the injection type.

Motivation & Objective

• To improve code-to-text generation by leveraging syntactic structure rather than just token sequences.
• To address the limitations of sequence-based models that ignore code's hierarchical structure.
• To develop a model that generalizes better to unseen code patterns by using compositional AST paths.
• To demonstrate that structural encoding via AST paths leads to superior performance in code summarization and captioning.

Proposed method

• Representing a code snippet as a set of k randomly sampled, fixed-length paths between terminal nodes in its AST.
• Encoding each AST path as a vector using a unidirectional LSTM to capture compositional syntactic structure.
• Using an attention mechanism over the encoded path vectors during decoding to generate each output token.
• Sampling k new AST paths at every training iteration to provide data-level regularization and improve generalization.
• Training the model end-to-end using cross-entropy loss with teacher forcing for sequence generation.
• Employing subword tokenization (Byte-Pair Encoding) for both code and natural language inputs and outputs.

Experimental results

Research questions

• RQ1 Can modeling code as AST paths improve code-to-text generation compared to token-level sequence modeling?
• RQ2 Does attention over structural paths lead to better performance than attention over raw tokens in code generation?
• RQ3 How does the use of dynamic path sampling during training affect model generalization and performance?
• RQ4 To what extent does structural encoding reduce sensitivity to surface-level syntactic variations in functionally equivalent code?
• RQ5 Can a model trained on AST paths generalize to unseen syntactic patterns not seen during training?

Key findings

• CODE2SEQ achieves state-of-the-art performance on code summarization and captioning tasks, outperforming models specifically designed for code and SOTA NMT models.
• The model improves BLEU scores by 62% on a C# comment generation dataset compared to a prior SOTA model using linearized ASTs.
• Ablation studies show that removing syntactic paths (using only tokens) reduces performance significantly, proving the importance of structural encoding.
• Dynamic path sampling at each training iteration improves generalization and boosts performance over fixed path sampling.
• Attention mechanisms in CODE2SEQ are crucial, with the no-attention variant showing a substantial drop in performance, similar to their role in standard NMT.
• The model generalizes well to unseen code patterns, as it can represent and attend over any syntactic path via LSTMs, even if not seen during training.