[Code Representation Learning 研究筆記（四）] Path-Sensitive Code Embedding via Contrastive Learning for Software Vulnerability Detection
作者：Xiao Cheng, Guanqin Zhang, Haoyu Wang, Yulei Sui
論文：https://arxiv.org/pdf/2009.08366.pdf
開源程式碼：unknown
1. Introduction
- 現有應用於靜態漏洞檢測的NN-based方法，通常的學習方法是以 Graph Neural Network 的方式學習與獲取程式碼的抽象結構訊息（e.g., DFG, CFG, AST, etc.），且這些方法的學習目標與思路通常皆聚焦在分類的問題上。
- 此研究提出了一個名為 ContraFlow 方法，主要針對基於靜態漏洞檢測的應用場景所提出。ContraFlow 通過 self-supervised contrastive learning 的方式來賽選出可行的 value-flow（program dependence）paths 而不是 graph structure，從而針對 path-based vulnerability detection 的問題顯著減少其所需的高昂訓練與標記數據的成本。
2. Proposed Method
(a) Contrastive Value-Flow Embedding
- 此階段主要目的在於通過 contrastive learning 的方式訓練一個 value-flow embedding model 來學習與獲得 value-flow paths 的特征表示, 並名為 Value-flow Path Encoder (VPE)。
- 給定一組 value-flow paths (unlabeled source code 藉由 static analyzer SVF 萃取而來）。透過一種 data augmentation 方法（SimCSE）來產生 contrastive value-flow representations。
- 並且利用 standard Noise Contrastive Estimate (NCE) loss function 來 maximize 相似 value-flow path 之間的相似度。透過此方式來更新 VPE的參數，以提示它保留 value-flow paths 的深層語義（deep semantics）。 與供接下來的兩個階段所用。
- 如範例中，給定一批從某 code fragment 中萃取出來的 value-flow paths，裡頭包含了 𝜋1, 𝜋2, 𝜋3, 𝜋4 四個 paths。它們皆被餵給 VPE 兩次（分別採用不同的 dropout masks）。第一次 dropout 獲得的向量表示定義為 v𝜋1, v𝜋2, v𝜋3, v𝜋4，而第二次 dropout 獲得的則表示為 contrastive representation 與分別定義為 v𝜋1+, v𝜋2+, v𝜋3+, v𝜋4+。由相同 𝜋 生成的 embedding representations 被視為正樣本，而同一批不相同 𝜋 生成的則被視為負樣本。
- VPE 以此作為預訓練學習，並在反向傳播中使用 NCE loss function，使得相對應的向量表示如 v𝜋1 與 v𝜋1+ 彼此靠近；非對應的向量表示如 v𝜋1 與 v𝜋3+ 彼此疏遠。
- 進一步詳細定義解釋可參閱 2.1。
(b) Value-Flow Path Selection
- 此階段則是用於精準賽選出可行且具有代表性的 value-flow paths 來表示code fragment，並排除不可行的 value-flow paths，以達到讓下個階段的 Detection Model Training 加速與有效學習的效果。
- 一開始先從 (a) 獲得輸入路徑的特征向量/表示，並使用 self-supervised active learning 的方式對路徑進行排名並採樣一部分定量的樣本。這些 paths 被進一步餵到 Path-Feasibility Checker 以刪除不可行的 paths。
- 如範例中，v𝜋2（path 3 →9 →13）就是不可行的，因為對應的 statement 3 →9 是成立在 !FLG，而 statement 9 →13 則是成立在 FLG 的條件下。v𝜋3（path 2 →6 →15）亦是如此。v𝜋1 與 v𝜋4 則是可行 path 並且會被提供至下個階段 (c) detection model 所使用。
- 進一步詳細定義解釋可參閱 2.2。
(c) Detection Model Training
- 此階段將訓練一個 Vulnerability Detection Model，使用的輸入則是階段 (b) 中被選中的 paths 以及它對應的 label（vulnerable or safe）。
- 首先先為每個被選中的 value-flow path 從 VPE 萃取出其向量表示。
- 接著利用 Transformer 來生成 contextual vector 以捕獲 path 之間的交互關係。
- 然後這些 vector 會被餵進 soft attention layer 對它們進行評分以及聚合它們成一個 vector，用以最終的 detection model 訓練。
- 此外此模型還可以根據它們對模型輸出的貢獻程度（attention score）來表示 value-flow paths 和 statements 的重要程度。
- 模型完成訓練後，在預測過程中可用於檢測未曾見過的程式其潛在的 vulnerabilities。檢測的輸出結果則是根據 attention score 將排名靠前的 value-flow paths 視為 buggy paths，備註這些排名中的 value-flow paths 僅包含由階段 (b) 選出來的可行 paths。
2.1 Contrastive Value-Flow Embedding
2.1.1 Guarded value-flow path（𝜋）
是由一串程式 statement 語句所組成，表示一個變量 variable 從宣告定義到被使用的過程（def-use chain）。
2.1.2 Noise Contrastive Estimate (NCE)
NCE loss function 為學習目標以最大化正樣本。對同批中任兩 value-flow path 向量 v𝜋i 與 v𝜋j 計算其 cosine similarity。
2.1.3 Local Encoding
負責使用 statement encoder 來獲得 value-flow path（𝜋）的 AST-subtree vector representations。定義 si 為 𝜋 中的第 i 個statement，而 local encoding 則從 𝜋 的一系列 statements [s0, s1, …, sN] 產生出其 vector representations [vs0, vs1, …, vsN]。
2.1.3.1 Statement Encoder
旨在對 statement s 建立 in-depth feature，其中包含 s 的 node token，types 以及其 syntactic structure (AST edges)。
- 設 ni 為某個 statement s 的某個節點，透過其 node type 與 node tokens初始化其向量為 vni。
- 設 𝐶𝑖 表示為子節點 ni（children node）的 AST-subtree。
- 隨後計算出其權重矩陣 W 與 ni 和 nj 的 attention 權重 aij，以此對 vni 進行更新從而獲得 vni’。
- as 則是 global attention vector 與 Wa 是 weighted matrix。
- || 表示 concatenation 而 𝜎(·) 表示 activation function。
- 計算所有 [vn0’, vn1', …, vnN’] 的 element-wise mean 並且與 N 個 node vectors 的 max pooling 進行 concatenation，由此生成其 summary vector vsm。
- 由於模型可能設有 L 個 layer，因此使用了 JK-net architecture 來將 L 個 layer 所產生的不同尺度 vsm 進行聚合加總，最終產生出 final statement vector vjk。
- 最後配置一個 fully connected layer 與一個 dropout layer 將 final statement vector 由 vjk 更新為 v𝑠 = dropout(Wv𝑗𝑘 )。
2.1.4 Global Encoding
負責將 vector representations [vs0, vs1, …, vsN] 利用 Bidirectional Gated Recurrent Unit (BGRU) 生成對應的 hidden state vectors [hs0, hs1, …, hsN]。隨後為每個 hs 透過 attention mechanism 賦予相應的 weights 並求加總，從而獲得一個 fixed length vector v𝜋 來充當 𝜋 的語義表示。
2.2 Value-Flow Path Selection
2.2.1 Value-Flow Active Learning
此階段則是用於精準賽選出可行且具有代表性的 value-flow paths 來表示code fragment，並排除不可行的 value-flow paths，以達到讓下個階段的 Detection Model Training 得以：
- 有效學習：在保留了可行路徑，相當於為模型過濾了 noise input。
- 加速訓練：在保留了可行路徑，相當於減少了訓練 input。
具體實踐過程如下
- 給定一批 value-flow paths H = [v𝜋1, …, v𝜋n] ∈ Rnxd，旨在賽選出可行並且具有代表性的子集 H’ ∈ Rkxd ⊆ V。
- 使用的模型架構為一個 encoder-decoder，其中包含了兩個需優化的參數矩陣 Q 和 P ∈ Rnxn。
- 使用的 loss 為 active learning 論文裡提出的 joint reconstruction loss 來 minimize H 與 QH 的距離。
- 使用 K-means 獲得 H 的 cluster centroid matrix C。
- 最終從 decoder 中輸出 G。
- 後續對兩個參數矩陣 Q 和 P 的每一列進行 l2-norm 處理以及 normalize 到 [0, 1]，由此產生出兩個 ranking vectors qhat, phat ∈ Rn。
- 最後 qhat 和 phat 彼此進行 merging 以及 descending order sort。
- 其中 top-k value-flow paths 被視為最具代表性的 paths，將其帶入 feasibility analyses。
2.2.2 Value-Flow Path Feasibility Analysis
針對由 value-flow active learning 獲得的被選中的 value-flow paths，將其帶入 feasibility checking 進一步提煉出更精確的 code embedding。
Feasibility checking 透過 guarded value-flow graph 方法來對 edges 用 guarded 進行標註，來描述 control-flow graph 的 control-flow transfer conditions。（其實這裡本人沒有很懂，因為沒看過 guarded value-flow graph 這個論文，之後可能會補上）。
- 給定一個 value-flow path 𝜋 =[s0, s1, …, sN]，定義一個 boolean function guardv(𝜋)，透過 𝜋 的 control-flow paths 從 s0 至 sN 依照編譯過程中每對可能觸發的路徑，如果 guardv(𝜋) 返回 True 則 𝜋 是可行的，否則則意味著不可行。
- CP(si, si+1) 定義為 statement si 和 si+1 所有可能的 control flow paths。
- CE(p) 則是表示 control flow path p 當中的 control-flow edges。
- guarde(e) 則是表示 edge e ∈ CE(p) 的 control-flow transfer condition。
- si 與 si+1 中一個 guard 是基於 CFG 上的 path p 計算其 disjunctions （類似 or 判斷）。若當下分支的 control-flow transfer 屬於 conditional（if-else branch），則將 condition c 編碼為 True c 反之則為 False !c。若當下分支的 control-flow trasfer 並非屬於 conditional 則無條件分配其 control-flow transfer 為 True。
- path p （2 → 9）為可行的（feasible），因為 9 的 conditional c 是可行發生的。
- path p （2 →11）為不可行（infeasible），因為 11 的 conditional !c 是實際在編譯過程中是不會成立的，因此將它從 code embedding 中剔除。
2.3 Detection Model Training
應該就是傳統的 Transformer 模型，所以就省略介紹吧。 XD
