Zenn
📄

【論文紹介】ChatGPTは脆弱性検知にどの程度有用か

2024/01/14に公開
機械学習
Security
ChatGPT
LLM
tech

はじめに

近年、ChatGPTをはじめとする大規模言語モデル(Large language model; LLM)は目覚ましい進歩を遂げています。ソフトウェア開発の分野においても、コード生成やコードレビューにLLMが活用されているようです。今回は、そんなLLMがソースコード中の脆弱性の検出(vulnerability detection)にどの程度有用かを調査した論文をまとめました。

論文情報

概要

この研究は、190,000以上のC/C++の関数を対象に、ファインチューニング無しのChatGPT(gpt-3.5-turboおよびgpt-4)の脆弱性検出性能を検証し、ファインチューニングを行った他のstate-of-the-artなベースライン(CodeBERT, GraphCodeBERT, およびAIBugHunter)と比較したものです。なお、ファインチューニングとは、公開されている既存の学習済みモデルに対して、特定のタスク(今回の場合は脆弱性診断に関するタスク)での性能向上を目的に独自の追加データでモデルを再学習させることを指します。この研究では以下の４つの問いを検証しました。詳細は後述しますが、先に結論を述べると次の通りです。

1. ChatGPTは、関数ごとおよび行ごとの脆弱性検知において、どの程度の性能を発揮するか

ChatGPT(gpt-4)は、関数ごとの脆弱性検知において、F1-scoreで10%, top-10 accuracyで25%の性能を示した。また、行ごとの脆弱性診断において、F1-scoreで29%, top-10 accuracyで65%の性能を示した。これらの値は比較対象のファインチューニング済みベースラインの中で最も低い。

2. ChatGPTは、脆弱性の種類(CWE-ID)の分類において、どの程度の性能を発揮するか

ChatGPT(gpt-4)は、multiclass accuracyで20%の性能を示した。この値は比較対象のファインチューニング済みベースラインの中で最も低い。

3. ChatGPTは、脆弱性の深刻度(CVSS)の予測において、どの程度の性能を発揮するか

ChatGPT(gpt-4)は、平均二乗誤差(mean squared error; MSE)で5.85、平均絶対誤差(mean absolute error; MAE)で1.86の性能を示した。この値は比較対象のファインチューニング済みベースラインの中で最も悪い。

4. ChatGPTは、脆弱性のある関数の修正タスクにおいて、どの程度の性能を発揮するか

ChatGPT(gpt-4)は、テストデータ中の脆弱性のある関数全てにおいて、適切な修正パッチを生成することはできなかった。一方、他のファインチューニング済みベースラインは、脆弱性のある関数のうち7%から30%を修正できた。

結果として、ファインチューニング無しのChatGPTは上記４つの全てのタスクにおいてファインチューニング済みのベースラインよりも低い性能を示しました。この研究は、数千億パラメータをもつとされるChatGPTであっても、ファインチューニングを行わない場合、ファインチューニングを行った1-2億パラメータのモデルに脆弱性診断性能で劣ることを示しています。

実験対象のモデルの概要

本研究では、脆弱性の検知、分類、深刻度の評価、および脆弱性の自動修正の4つのタスクにおけるChatGPT(gpt-3.5-turboおよびgpt-4)の性能を、次の3つのベースラインと比較しています。

  1. CodeBERT：ソースコードに関するタスクに特化した言語モデル。コードから自動でドキュメントを生成する作業(code documentation generation)や、開発者の意図する挙動を示すコードをデータベースから自然言語で検索する作業(code search)などのタスクで高い性能を示している。

  2. GraphCodeBERT：CodeBERT同様、ソースコードに関するタスクに特化して学習された言語モデル。CodeBERTとの違いは、ソースコードおよびテキストのみならず、コード中の変数間の依存関係を示す有向グラフも学習に用いる点。前述のcode searchなどのタスクにおいてCodeBERTよりも高い性能を示している。

  3. AIBugHunter：ソフトウェアの脆弱性の検知および分類を行ってくれるツール。脆弱性の自動修正機能についても、公式サイト上でcoming soonとなっている(2024/01/14時点)。

CodeBERT, GraphCodeBERT, およびAIBugHunterについては、脆弱性に関するタスクのためのファインチューニングを行います。一方、ChatGPT(gpt-3.5-turboおよびgpt-4)についてはファインチューニングを行いません。なお、上記3つのベースラインのパラメータ数が1億~2億程度であるのに対し、ChatGPTでは、gpt-3.5-turboで1,750億、gpt-4のパラメータ数は公開されていないようですが、数千億、あるいは1兆以上とされています。この研究の注目すべき点は、1-2億程度のパラメータ数のファインチューニングモデルと、数千億パラメータを持つファインチューニング無しモデルのどちらが脆弱性診断において良い性能を示すかを検討した点である思います。

実験内容

実験1: ChatGPTは、関数ごとおよび行ごとの脆弱性検知において、どの程度の性能を発揮するか

実験設定と評価指標

実験1では、1) 与えられた関数が脆弱であるかどうかの二値分類(binary classification)性能および、2) どの行が脆弱かを判断する性能の２点を評価します。1)の二値分類性能は、recall, precision, およびF1-scoreによって測ります。また、2)の行ごとの脆弱性検知性能については、top-10 accuracyによって評価します。
なお、recallとは、テストデータ中の陽性サンプルのうち、モデルが実際に陽性であると言い当てられたサンプルの割合です。例えば、テストデータ中に1000個の脆弱な関数があり、そのうち800個についてモデルが正しく脆弱であると判断できた場合、recallは0.8となります。また、precisionは、モデルが陽性であると判断したサンプルのうち、実際に陽性であるサンプルの割合です。例えば、モデルが脆弱であると判断した関数が1000個あり、そのうち800個が実際に脆弱な関数である場合、precisionは0.8となります。つまり、recallが高いほどモデルが脆弱性のある関数を漏れなく検出できていることを示し、precisionが高いほど誤報率が低いことを示します。F1-scoreはrecallとprecisionの調和平均であり、(2×recall×precision)/(recall+precision)で計算されます。
recall, precision, F1-scoreは、どれも値が大きいほどモデルの性能が良いことを示します。
また、top-k accuracyとは、テスト対象の全関数の数をN、モデルが、それぞれの関数中の各行を脆弱である可能性の高さでランク付けし、その上位k行のうちに実際に脆弱な行が入っていた関数の数をnとして\frac{n}{N}で表されます。top-k accuracyも、値が大きいほどモデルの性能が高いことを示します。

実験に用いるデータセット

実験1で用いるデータはBig-Vulデータセットです。Big-Vulデータセットは、348のGitHubプロジェクトから収集した188,000個のC/C++関数からなり、合計91種類、3,754個の脆弱な関数を含んでいます。それぞれの脆弱な関数には、脆弱性の種類を示すCWE識別子(CWE-ID)および脆弱性の深刻度を示す共通脆弱性評価システム(Common Vulnerability Scoring System; CVSS)があらかじめ付与(ラベリング)されています。全データの80%を訓練用、10%を検証用、10%をテスト用に使います。

プロンプト設計

脆弱性検知のためのプロンプトは以下の図のとおりです。１つ目のプロンプトは関数ごとの脆弱性検知のためのプロンプトであり、入力関数に脆弱性があるかどうかを1, 0で(Yes/Noで)出力するように要求します。２つ目のプロンプトは、行ごとの脆弱性検知のためのプロンプトであり、入力関数が脆弱であることを宣言した上で、関数中の各行について脆弱である可能性をスコア化し、その上位10行を出力するよう要求します。


関数ごとおよび行ごとの脆弱性検知のためのプロンプトの例(元論文から引用)

実験2: ChatGPTは、脆弱性の種類(CWE-ID)の分類において、どの程度の性能を発揮するか

実験設定と評価指標

実験2では、脆弱性のある関数をモデルに入力し、その関数の持つ脆弱性の種類(CWE-ID)の判断性能を測ります。すなわち、この問題はCWE-IDを予測する多クラス分類(multiclass classification)です。この実験では、評価指標としてmulticlass accuracyを報告しています。multiclass accuracyは、値が大きいほどモデルの性能が良いことを意味します。

実験に用いるデータセット

実験1と同様、Big-Vulデータセットを使用します。実験1と同様、全データの80%を訓練用、10%を検証用、10%をテスト用に使います。上述のように、それぞれの脆弱な関数にはCWE-IDがラベル付けされおり、このラベルをモデルがどの程度正確に言い当てられるかを検証するというのが実験2の目的です。

プロンプト設計

以下の図のように、入力関数が脆弱であることを宣言した上で、CWE-IDを予測するよう要求します。

脆弱性の分類(CWE-IDの予測)タスクのためのプロンプトの例(元論文から引用)。

実験3: ChatGPTは、脆弱性の深刻度(CVSS)の予測において、どの程度の性能を発揮するか

実験設定と評価指標

実験3は、入力関数が持つ脆弱性の深刻度(CVSS)を連続値で評価する回帰問題です。CVSSは情報システムの脆弱性を評価する指標であり、0.0から10.0までの連続値を取ります。評価指標は、平均二乗誤差(mean squared error; MSE)および平均絶対誤差(mean absolute error; MAE)です。MSEとMAEのどちらも、予測値が真の値からどれだけずれているかを評価する指標であり、値が小さいほどモデルの性能が良いことを意味します。

実験に用いるデータセット

実験1, 2と同様、Big-Vulデータセットを使用します。実験1と同様、全データの80%を訓練用、10%を検証用、10%をテスト用に使います。上述のように、それぞれの脆弱な関数にはCVSSがラベル付けされおり、このCVSSの値をモデルが言い当てられるかを検証するというのが実験3の目的です。

プロンプト設計

以下の図のように、入力関数が脆弱であることを宣言した上で、そのCVSSを予測するよう要求しています。

脆弱性の深刻度(CVSSS)の予測タスクのためのプロンプトの例(元論文から引用)。

実験4: ChatGPTは、脆弱性のある関数の修正タスクにおいて、どの程度の性能を発揮するか

実験設定と評価指標

実験4では、ChatGPTが脆弱性を持つ入力関数に関数に対して適切な修正パッチを生成できるかどうかを検証します。評価指標はpercentage of perfect prediction (%PP)です。%PPは次式で計算されます。

{\text{\%PP} = \frac{\text{total correct  predictions}}{\text{total testing samples}}}
すなわち、テスト(予測)に用いた全サンプルのうち、正しく修正パッチを生成できたサンプルの割合が%PPです。ただし元論文では、モデルが生成した修正パッチが真の修正パッチ(あらかじめ正解として用意している修正パッチ)と完全に一致する場合のみ「正しく修正パッチを生成できた (correct prediction)」としてカウントしています。

実験に用いるデータセット

実験4では、先述のBig-Vulデータセットに加え、CVEFixesデータセットを用います。CVEFixesデータセットはBig-Vulデータセットと同様、ソースコードとその脆弱性の情報を含む大規模なデータセットです。「実験設定と評価指標」で述べたように、実験4では、モデルが正しい修正パッチを生成できるかどうかを検証することが目的です。そしてその「正しさ」とは、あらかじめ用意しておいた正解の修正パッチと一致しているかどうかで判断します。つまり、正解の修正パッチを用意する必要があります。そこで、機械学習を用いてソースコードの脆弱性を自動で修正する手法(提案者の名前を取ってChenらの方法と呼びます)を用いて、Big-VulデータセットおよびCVEFixesデータセット中の脆弱なコードに対する正解の修正パッチを準備します。すなわち、「実験設定と評価指標」で述べた、「あらかじめ正解として用意している修正パッチ」は、機械学習手法により自動的に生成されたものであり、完全な修正パッチではないという点に注意が必要です。全データの70%を訓練用、10%を検証用、20%をテスト用に使います。

プロンプト設計

実験4ではまず、脆弱な関数とそれに対する正解の修正パッチ(Chenらの方法によって生成した修正パッチ)の組を３組提示します。その上で、脆弱な関数を与え、それに対する修正パッチの生成を要求します。

脆弱性の自動修正のためのプロンプトの例(元論文から引用)。

実験結果

実験1: ChatGPTは、関数ごとおよび行ごとの脆弱性検知において、どの程度の性能を発揮するか

引用した図が示すように、ChatGPTの脆弱性検知性能は、他のファインチューニング済みモデルと比較すると低いことがわかります。例えば、与えられた関数が脆弱であるかどうかを予測する二値分類タスクにおいては、gpt-4でF1-scoreは29%です。これは、テストデータ中の全ての脆弱な関数のうち、29%の関数について脆弱性を検知したことを意味しています。また、行ごとの脆弱性診断性能において、gpt-4のtop-10 accuracyは65%です。これは、テストデータ中の全ての脆弱な行のうち65%の行を、脆弱性をもつ可能性が高いとgpt-4が判断したことを意味しています。このタスクで最高の性能を達成したのはAIBugHunterです。gpt-4のパラメータ数が数千億であるとされる一方、AIBugHunterのパラメータ数は1.2億であり、gpt-4のパラメータ数の数千分の1にすぎません。この結果は、LLMによる脆弱性検知におけるファインチューニングの重要性を示唆しています。

実験1の結果。どの指標も、値が高いほどモデルの性能が良いことを示す(元論文から引用)。

実験2: ChatGPTは、脆弱性の種類(CWE-ID)の分類において、どの程度の性能を発揮するか

引用した図が示すように、ChatGPTのCWE-ID判断性能は、他のファインチューニング済みモデルと比較して低いことがわかります。multiclass accuracyは、gpt-4で20%である一方、他のファインチューニング済みモデルは62%から65%の性能を示しています。この結果も、実験1と同様、LLMによる脆弱性検知におけるファインチューニングの重要性を示唆しています。

実験2の結果。値が大きいほどモデルの性能が良いことを示す(元論文から引用)。

実験3: ChatGPTは、脆弱性の深刻度(CVSS)の予測において、どの程度の性能を発揮するか

実験3においても、ChatGPTは他のファインチューニング済みモデルと比較して性能が低いことがわかります。やはり、LLMによる脆弱性診断ではファインチューニングが重要であるようです。

実験3の結果。MSEもMAEも値が小さいほどモデルの性能が良いことを示す(元論文から引用)。

実験4: ChatGPTは、脆弱性のある関数の修正タスクにおいて、どの程度の性能を発揮するか

この実験でも、３つの指標全てにおいて、chatGPTはファインチューニング済みのモデルと比較して限られた性能しか発揮しませんでした。特に、%PPを見ると、ChatGPTは脆弱性のある全ての関数において適切な修正パッチを生成できていないことがわかります。ただし、他の3つのファインチューニング済みモデルにおいても%PPは7%から30%にとどまっており、LLMを用いた修正パッチの自動生成は困難なタスクであることがわかります。


実験4の結果。どの指標も値が小さいほどモデルの性能が良いことを示す(元論文から引用)。

まとめ

本研究の要点は、脆弱性診断のタスクでファインチューニングを行った1-2億パラメータのモデルの方が、ファインチューニングを行っていない数千億パラメータのモデルよりも脆弱性診断において良い性能を示すことを示した点だと思います。脆弱性診断に言語モデルを用いる場合は、ファインチューニングを行い、ドメイン固有の知識を学習させる必要であるようです。

Discussion