DeepMind 和谷歌的這項新研究聚焦檢索和匹配圖結構物件這一極具挑戰性的問題,做出了兩個重要貢獻。
首先,研究者展示瞭如何訓練圖神經網路(GNN),使之生成可在向量空間中執行高效相似性推理的圖嵌入。其次,研究者提出了新型圖匹配網路模型(GMN),該模型以一對圖作為輸入,通過基於跨圖注意力的新型匹配機制進行聯合推理,從而計算它們之間的相似性分數。
研究者證明 GMN 模型在不同領域中的有效性,包括極具挑戰性的基於控制流圖的函式相似性搜尋問題,這個問題在檢索軟體系統的漏洞中起著非常重要的作用。實驗分析表明 GMN 模型不止能在相似性學習的環境下利用結構,還能超越針對這些問題手動精心設計的特定領域基線系統。
研究主題:圖相似性學習問題
圖是編碼關係結構的自然表徵,常常出現在多個領域中。根據圖結構資料定義的計算可以用在各種領域中,從計算生物學和化學的分子分析到自然語言理解中知識圖或圖結構解析的分析都可以。
近幾年來,圖神經網路(Graph Neural Network,GNN)已經成為可以有效學習結構資料表徵、解決各種基於圖的監督預測問題的模型了。這樣的模型在迭代聚合區域性結構資訊的傳播過程中設計並計算圖節點表徵,從而對圖元素的排列(permutation)具有不變性。然後直接將這些節點表徵用於節點分類,或者將它們合併到用於圖分類的圖向量中。而 GNN 在監督分類或迴歸以外的問題的相關研究相對較少。
DeepMind 的這篇論文研究的是圖結構物件的相似性學習問題,這個問題在現實生活中有很多重要的應用,尤其是在圖資料庫中基於相似性的搜尋。還有一個應用是涉及電腦保安的二元函式相似性搜尋,給定的二元函式可能包含有已知漏洞的程式碼,我們要檢查這個二元函式中是否有和已知易受攻擊的函式相似的控制流圖(control-flow-graph)。這有助於識別閉源軟體中易受攻擊的靜態連結函式庫,這是一個很常見的問題 (CVE, 2010; 2018),現在還沒有很好的解決方法。圖 1 展示了一個例子,在這個例子中用匯編語言註釋的控制流圖來表示二元函式。這種相似性學習問題極具挑戰性,因為就算是圖之間細微的差別也會造成語義上極大的不同,但結構不同的圖語義上可能非常相似。因此,對這個問題而言,一個成功的模型應該(1)利用圖結構;(2)能從圖的結構和學習到的語義中推匯出圖的相似性。
圖 1:二元函式相似性學習問題。檢查兩個圖是否相似需要推理圖的結構和語義。左邊兩個控制流圖對應使用不同編譯器編譯的相同函式(因此二者比較相似),但右側圖對應的是不同函式。
解決方案
為了解決圖相似性學習問題,該論文研究了 GNN 在這種情況中的使用,探討了如何用 GNN 將圖嵌入到向量空間,並學習這種嵌入模型,從而使向量空間中相似的圖靠近、不相似的圖分開。這個模型的一個重要特性是它可以將每一個圖獨立地對映到一個嵌入向量,然後在向量空間中執行相似性計算。因此,可以預先計算並索引大型資料庫中的圖嵌入,這樣就能用快速的最近鄰搜尋資料結構(如 k-d 樹) 或區域性敏感雜湊演算法 (Gionis et al., 1999) 執行高效的檢索。
研究者進一步擴充套件 GNN,提出新型圖匹配網路(Graph Matching Networks,GMN)來執行相似性學習。GMN 沒有單獨計算每個圖的圖表徵,它通過跨圖注意力機制計算相似性分數,來關聯圖之間的節點並識別差異。該模型依賴成對圖計算圖表徵,因此它比嵌入模型更強大,並在準確率和計算之間做出了很好的權衡。
研究者在三個任務上評估了 GMN 和基線模型:僅捕獲結構相似性的合成圖編輯距離學習任務(synthetic graph edit-distance learning tas),以及兩個現實世界任務——二元函式相似性搜尋和網格檢索,這兩項任務都需要推理結構相似性和語義相似性。在所有任務中,GMN 都比基線和結構不可知(structure agnostic)模型的效能更好。在更詳細的模型簡化測試中,研究者發現 GMN 始終優於圖嵌入模型和 Siamese 網路。
該研究的貢獻如下:
展示瞭如何用 GNN 產生用於相似性學習的圖嵌入;
提出了新型圖匹配網路(GMN),該網路基於跨圖注意力匹配來計算相似性;
實驗證明,該研究提出的圖相似性學習模型 GMN 在多個應用中都有良好的表現,比結構不可知模型和現有的手動建立的基線模型都要好。
深度圖相似性學習
給定兩個圖 G1 = (V1, E1) 和 G2 = (V2, E2),我們需要一個可以計算兩圖之間相似性分數 s(G1, G2) 的模型。每個圖 G = (V, E) 被表示為節點 V 和邊 E 的集合,每個節點 i∈V 都可以和特徵向量 x_i 相關聯,每條邊 (i, j) ∈ E 都可以和特徵向量 x_ij 關聯起來。這些特徵可以表示節點型別、邊的方向等。如果一個節點或者一條邊不能關聯任何特徵,那麼我們可以將對應向量設定成值為 1 的常量。研究者提出了兩個圖相似性學習模型:一個是基於標準 GNN 的學習圖嵌入的模型;另一個是更為嶄新也更加強大的 GMN。圖 2 展示了這兩個模型:
圖 2:圖嵌入模型(左)和圖匹配模型(右)圖示。
圖嵌入模型
圖嵌入模型可以將每一個圖都嵌入到向量中,然後用向量空間中的相似性矩陣衡量圖之間的相似性。GNN 嵌入模型包括三個部分:編碼器、傳播層和聚合器。
圖匹配網路
圖匹配網路以一對圖作為輸入,計算它們之間的相似性分數。和嵌入模型相比,圖匹配模型聯合成對圖計算相似性分數,而不是先將每個圖獨立地對映到向量上。因此,圖匹配模型可能比嵌入模型更加強大,但它需要額外的計算效率。
圖匹配網路改變了每個傳播層中的節點更新模組,這樣不僅可以考慮到每個圖的邊上的聚合資訊,還可以考慮到衡量一個圖中的一個節點和其他圖中的一或多個節點匹配程度的跨圖匹配向量:
和圖嵌入模型相比,圖匹配模型能根據對比圖改變圖的表徵。圖匹配模型可以調整圖的表徵,在它們不匹配時放大它們之間的差異。
實驗
研究者在三個任務上評估了圖相似性學習(Graph Similarity Learning,GSL)框架、圖嵌入模型(GNN)以及圖匹配網路(GMN)的效能,並將這些模型與其他方法進行了對比。總體上,實驗結果表明在圖相似性學習任務上,GMN 表現優異,而且始終優於其他方法。
學習圖編輯距離(GED)
圖 G1 和 G2 之間的圖編輯距離即將 G1 變換為 G2 所需的最小編輯操作。通常這些編輯操作包括新增/移除/替換節點和邊。圖編輯距離是衡量圖之間相似性的自然指標,在圖相似性搜尋中有很多應用。
從下表 1 中可以看出,通過學習特定分佈的圖,GSL 模型的效能優於一般的基線模型,而 GMN 的效能持續優於圖嵌入模型(GNN)。
表 1:在來自不同分佈的圖上訓練的圖嵌入(GNN)模型和圖匹配(GMN)模型和基線的比較,衡量指標為pair AUC / triplet accuracy (×100)。
圖 3:5 個傳播層後 GMN 跨圖注意力的視覺化。在每一對圖中,左圖展示了從左到右的注意力,右圖展示的是從右到左的注意力。
基於控制流圖的二元函式相似性搜尋
二元函式相似性搜尋是電腦保安領域中的重要問題。當我們無法獲取原始碼時,可以通過二元函式執行分析和搜尋,例如在處理商業或嵌入式軟體或可疑的可執行程式時。
下圖 4 展示了具備不同傳播步和不同資料設定的不同模型在二元函式相似性搜尋任務上的效能。從圖中,我們可以看到:
圖嵌入模型和圖匹配模型的效能隨著傳播步的增加而持續提升;
在傳播步足夠的情況下,圖嵌入模型持續優於基線模型;
圖匹配模型在所有設定和傳播步的情況下都優於圖嵌入模型。
圖 4:不同模型在二元函式相似性搜尋任務上的效能(*100)。
更多基線和模型簡化測試
研究者檢測了GMN 模型中不同元件的效果,並將 GMN 模型與圖卷積網路(GCN)、圖神經網路(GNN)和 GNN/GCN 嵌入模型的 Siamese 版本進行對比。
下表 2 展示了實驗結果,表明:
GNN 嵌入模型是具備競爭力的模型(比 GCN 模型強大);
使用 Siamese 網路架構基於圖表徵學習相似性要比使用預先指定的相似性指標(Euclidean、Hamming 等)好;
GMN 優於Siamese 模型,這表明在計算過程早期進行跨圖資訊交流是非常重要的。
表 2:函式相似性搜尋任務上的更多實驗結果,以及模型在 COIL-DEL 資料集上的結果。