[论文解读] Using Machine Learning for Vulnerability Detection and Classification

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

Software Engineering Research被引用 452

一句话总结

该论文提出 CODE2SEQ，一种新颖的神经模型，通过将代码片段编码为其抽象语法树（AST）中的组合路径，利用LSTM嵌入路径并使用注意力机制在解码过程中关注这些路径，从而从源代码生成自然语言序列。该模型在两种语言和四个数据集上的代码摘要和字幕生成任务中，显著优于先前的专用代码模型和最先进的神经机器翻译模型。

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.

研究动机与目标

• 通过利用语法结构而非仅标记序列来改进代码到文本的生成。
• 解决仅基于序列的模型忽略代码层次结构的局限性。
• 通过使用组合式 AST 路径，开发一种能更好地泛化到未见代码模式的模型。
• 证明通过 AST 路径进行结构编码可显著提升代码摘要和字幕生成任务的性能。

提出的方法

• 将代码片段表示为从其 AST 中随机采样、长度固定的终端节点之间的 k 条路径集合。
• 使用单向 LSTM 将每条 AST 路径编码为向量，以捕捉组合式语法结构。
• 在解码过程中，对编码后的路径向量使用注意力机制，以生成每个输出标记。
• 在每次训练迭代中采样 k 条新的 AST 路径，以提供数据级别的正则化并提升泛化能力。
• 使用交叉熵损失和教学 forcing 方法端到端训练模型，以进行序列生成。
• 对代码和自然语言的输入与输出均采用子词分词（字节对编码）。

实验结果

研究问题

• RQ1 与基于标记的序列建模相比，将代码建模为 AST 路径是否能提升代码到文本的生成性能？
• RQ2 在代码生成中，对结构化路径的注意力机制是否优于对原始标记的注意力机制？
• RQ3 在训练过程中使用动态路径采样如何影响模型的泛化能力和性能？
• RQ4 结构化编码在多大程度上减少了对功能等价但表面语法不同的代码的敏感性？
• RQ5 在训练期间未见过的语法模式下，基于 AST 路径训练的模型是否能泛化到未见的语法模式？

主要发现

• CODE2SEQ 在代码摘要和字幕生成任务中达到最先进性能，优于专门针对代码设计的模型和最先进 NMT 模型。
• 与先前最先进模型（使用线性化 AST）相比，该模型在 C# 注释生成数据集上的 BLEU 分数提升了 62%。
• 消融研究显示，若移除语法路径（仅使用标记），性能显著下降，证明结构化编码的重要性。
• 在每次训练迭代中进行动态路径采样，可提升泛化能力并优于固定路径采样。
• 注意力机制在 CODE2SEQ 中至关重要，无注意力变体性能显著下降，其作用与标准 NMT 中类似。
• 该模型能很好地泛化到未见的代码模式，因为它能通过 LSTM 表示并关注任何语法路径，即使这些路径在训练中未出现过。