[論文レビュー] Using Machine Learning for Vulnerability Detection and Classification

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

Software Engineering Research被引用数 452

ひとこと要約

この論文では、コードスニペットをその抽象構文木（AST）内の合成的経路として符号化し、LSTMを用いて経路を埋め込み、デコード中にそれらの経路に注目する注目メカニズムを用いることで、自然言語のシーケンスを生成する新しいニューラルモデル、CODE2SEQを紹介する。このモデルは、2つの言語と4つのデータセットにおいて、コード要約作成およびキャプション生成タスクで、先行のコード特化型および最先端のニューラル機械翻訳モデルを大きく上回る性能を発揮する。

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経路をサンプリングすることで、データレベルの正則化を実現し、一般化性能を向上させる。
• 教師強制を用いたシーケンス生成のための交差エントロピー損失を用いて、モデルをエンドツーエンドで学習する。
• コードおよび自然言語の入出力の両方に対してサブワードトークン化（バイトペアエンコーディング）を採用する。

実験結果

リサーチクエスチョン

• RQ1 コードをAST経路としてモデル化することで、トークンレベルのシーケンスモデル化と比較して、コードからテキストへの生成が向上するか？
• RQ2 構造的経路に対する注目が、コード生成において生のトークンに対する注目に比べて優れた性能をもたらすか？
• RQ3 学習中に動的経路サンプリングを適用することで、モデルの一般化性能および性能にどのような影響を与えるか？
• RQ4 構造的符号化は、機能的に同等のコードにおける表面的な構文的変化に対して、どれほど感受性が低下するか？
• RQ5 AST経路で学習したモデルは、トレーニング中に見られなかった未確認の構文的パターンに対しても一般化できるか？

主な発見

• CODE2SEQは、コード要約作成およびキャプション生成タスクで最先端の性能を達成し、特にコードに特化したモデルおよび最先端のNMTモデルを上回る。
• C#のコメント生成データセットにおいて、線形化されたASTを用いた先行の最先端モデルと比較して、BLEUスコアを62%向上させる。
• アブレーションスタディの結果、構文的経路を除去（トークンのみを用いる）すると性能が著しく低下し、構造的符号化の重要性が裏付けられる。
• 各学習イテレーションでの動的経路サンプリングは、固定経路サンプリングよりも一般化性能を向上させ、性能を向上させる。
• CODE2SEQにおける注目メカニズムは不可欠であり、注目なしのバージョンでは性能が著しく低下し、標準的なNMTにおけるその役割と同様である。
• モデルは未学習のコードパターンに対しても良好に一般化できる。LSTMを介して、トレーニング時に見られなかった任意の構文的経路を表現・注目できるからである。