面向特征融合與知識蒸餾的惡意軟件分類

莊賢, 陳志豪, 蔡鐵城, 陳開志, 廖祥文

(福建省網絡計算與智能信息處理重點實驗室, 數字福建金融大數據研究所, 福州大學計算機與大數據學院, 福建 福州 350108)

0 引言

惡意軟件分類[1]旨在提取惡意軟件的關鍵特征, 并用其來構建分類器, 從而對該惡意軟件的類別進行判斷, 本質上是一個多分類任務. 對惡意軟件進行分類, 可以有效分析不同惡意軟件的演變過程, 更好地保障軟件應用的安全性. 當前, 隨著代碼混淆等技術的不斷發展, 惡意軟件分類面臨著更大的挑戰. 傳統的惡意軟件分類方法難以有效應對新型惡意軟件[2], 基于深度學習的惡意軟件分類方法則因其良好的擴展性和穩定性受到廣泛關注[3]. 目前, 基于深度學習的惡意軟件分類方法, 依據惡意軟件特征提取方法的不同, 主要可以分為如下兩類.

1) 基于惡意軟件特征的分類方法. 通過靜態分析技術和動態分析技術分別提取惡意軟件中的靜態特征和動態特征, 用于訓練不同的模型. Krcál等[4]通過靜態分析技術來提取字節序列, 用于訓練所提出的深度神經網絡模型. Xie等[5]通過靜態分析工具提取出匯編指令, 利用LSTM獲得每個塊的代表性特征, 并結合注意力機制來減少垃圾指令. 鄭銳等[6]利用動態分析技術提取API調用序列, 通過雙向LSTM對前后API的調用概率關系進行建模. 這類方法需要人工制定特征工程來提取惡意軟件特征, 分類器的構建具有較高的成本. 此外, 在新型惡意軟件數量激增的情況下, 特征工程難以保持良好的特征抽取能力, 導致模型性能難以提升.

2) 基于惡意軟件圖像的分類方法. 通過惡意軟件可視化算法, 將惡意軟件可視化為圖像, 并利用卷積神經網絡, 從圖像中自動抽取惡意軟件特征, 并進行學習和分類. Nataraj等[7]最早提出將惡意軟件可視化為圖像的方法, 根據惡意軟件文件的大小, 將惡意軟件轉換為二維數組, 并生成相應的灰度圖. Yuan等[8]則通過計算字節傳輸概率矩陣, 將惡意軟件轉換為馬爾可夫圖, 將其用于訓練所提出的DCNN模型(改良的VGG16模型). Kalash等[9]通過VGG16模型, 從灰度圖中抽取特征, 并進行學習與分類. Vasan等[10]將灰度圖映射成彩色圖像, 并對圖像進行數據增強, 進而微調預訓練的VGG16模型. Aslan等[11]集成兩個預訓練卷積神經網絡, 構建一種混合模型, 從灰度圖中抽取特征, 并進行學習和分類. Parihar等[12]利用3個不同的卷積神經網絡, 分別從灰度圖中抽取特征, 將其用于模型的訓練和分類. 這類方法僅通過單一圖像語義表示來抽取惡意軟件特征, 導致模型無法捕獲惡意軟件序列特征和分布特征之間的依賴關系, 且模型的規模較大, 對計算資源較為依賴.

針對上述問題, 本研究提出面向特征融合與知識蒸餾的惡意軟件分類方法, 基于注意力機制, 模型對不同特征之間的關聯性和依賴關系進行挖掘與捕獲. 同時, 利用知識蒸餾來降低模型的規模, 以期減少模型對計算資源的消耗.

1 模型方法

本研究模型的框架如圖1所示, 共由6個部分組成. 1) 編碼層. 通過不同的可視化算法, 將惡意軟件二進制序列分別可視化為灰度圖和馬爾可夫圖. 2) 特征抽取層. 使用殘差網絡, 分別從編碼層生成的灰度圖和馬爾可夫圖中, 自動抽取惡意軟件的序列特征和分布特征. 3) 特征融合層. 由自注意力模塊和雙重注意力(attention on attention, AoA)模塊[13]構成, 利用注意力機制, 挖掘兩種不同特征之間的關聯性. 4) 特征修正層. 利用高速神經網絡[14], 對特征抽取層和特征融合層的輸出向量進行修正. 5)分類層. 通過3層感知機輸出預測概率分布. 6) 模型壓縮. 首先, 教師網絡向多個學生網絡進行知識遷移; 然后, 學生網絡互相之間進行協作學習; 最后, 選擇性能最優的學生網絡作為最終的分類模型.

圖1 模型框架圖Fig.1 Overall structure of model

1.1 編碼層

圖2為將惡意軟件可視化為灰度圖和馬爾可夫圖的過程.

圖2 惡意軟件可視化過程Fig.2 Malware visualization process

首先, 計算每8位二進制的字節大小(A), 將惡意軟件二進制序列轉換為字節向量; 然后, 分別根據Nataraj等[7]和Yuan等[8]的方法, 將字節向量重塑為二維數組和計算字節傳輸概率矩陣; 最后, 通過二維數組和概率矩陣生成對應的灰度圖和馬爾可夫圖. 字節大小(A)的計算公式為

(1)

式中:ai為8位二進制位中的第i位二進制位.

1.2 特征融合層

自注意力模塊對惡意軟件的序列特征和分布特征進行自注意力計算, 捕捉不同特征之間的依賴關系. 該模塊將序列特征向量和分布特征向量進行拼接, 得到輸入向量(fatt_in), 分別使用WQ、WK、WV這3個矩陣作為Query、 Key、 Value, 將其映射到相應的特征空間; 然后再進行注意力計算, 得到融合惡意軟件的序列特征和分布特征的融合向量(fatt_out). 融合向量的生成過程可用公式表示為

fatt_out=Att(WQfatt_in,WKfatt_in,WVfatt_in)

(2)

式中: Att(·)為自注意力函數.

fatt_out可能存在結果向量與查詢向量不匹配的情況, 導致誤導信息的產生, 這會影響模型的性能. 將fatt_in和fatt_out作為輸入, 利用AoA進行二次注意力計算, 可以過濾掉無關的注意力結果. 最終特征融合層的輸出向量, 即AoA的輸出向量(fAoA_out), 可定義為

fAoA_out=Dropout((fatt_out⊕fatt_in)W1+b1)?δ((fatt_out⊕fatt_in)×W2+b2)

(3)

式中:W1和W2為可訓練的矩陣;b1和b2為偏差向量;δ(·)為Sigmoid激活函數; ?和⊕分別代表矩陣點乘和拼接. 為防止參數過多帶來的過擬合, 本研究以50%的概率刪除輸出向量.

1.3 特征修正層

為充分利用惡意軟件的序列特征和分布特征, 同時增強融合模塊的泛化性, 利用高速神經網絡的門機制, 可以對不同的特征向量進行修正, 其輸出向量定義為

(4)

式中:x為特征抽取層得到的惡意軟件的序列特征向量、 分布特征向量或特征融合層得到的融合特征向量; 函數H(·)表示的是一個線性層和非線性層的組合; 矩陣WH、WT為線性層的可學習參數; 非線性層使用Sigmoid激活函數. 函數T(x)的表達式為

T(x)=δ((WT)Tx+bT)

(5)

式中:bT為-1或者-3. 由定義可知, 高速神經網絡將會直接保留輸入特征向量中的有效信息, 并對其中的無效信息進行映射, 進一步增強特征向量的有效性.

1.4 模型壓縮

模型壓縮的過程可分為4步. 1) 預訓練特征抽取層為ResNet-34的大模型, 對其進行非結構化剪枝[15]后, 將其作為教師網絡; 2) 固定教師網絡參數, 對多個特征抽取層為ResNet-18的學生網絡進行知識遷移[16]; 3) 多個不同的學生網絡互相之間進行協作學習; 4) 選擇性能最優的學生網絡作為最終的分類模型, 用于輸出分類結果.

知識蒸餾和協作學習的核心, 都是用于更新學生網絡參數的損失計算. 當對第i個學生網絡進行參數更新時, 對應的損失(LTL)將由散度值(L1)和混合損失值(L2)進行加權計算得到, 其表達式為

(6)

L2=Mix Loss(pi,y)

(7)

LTL=αL1+βL2

(8)

式中:DKL(·)代表KL散度函數;pi為第i個學生網絡的輸出向量; 在知識蒸餾階段和協作學習階段,pk分別為教師網絡和性能最優學生網絡的輸出向量;y為惡意軟件類別對應的數值;α+β=1, 且α>>β.混合損失函數的定義為

Mix Loss(pi,y)=LCE(pi,y)+LFL(pi,y)+LOHEM(pi,y)

(9)

式中:LCE(·)、LFL(·)和LOHEM(·)分別代表交叉熵損失函數Cross Entropy、 聚焦損失函數Focal Loss[17]和難樣本損失函數OHEM[18].

優化器選用隨機梯度下降(stochastic gradient descent, SGD), 通過SGD來更新模型的參數, 最小化訓練數據的損失.

2 實驗設計

2.1 數據集描述

本研究采用微軟數據集[19]和CCF數據集[20]來驗證模型的性能. 微軟公開數據集來源于微軟于2015年在Kaggle舉辦的大型惡意軟件分類競賽, 該數據集共有21 741個惡意軟件樣本, 其中的10 868個樣本為訓練集, 剩下的10 873個樣本為測試集. 該數據集中的每個樣本都分別屬于9個不同的惡意軟件類別或家族, 訓練集中的樣本數(Ntrain)分布情況如圖3(a)所示. CCF公開數據集來源于CCF于2021年舉辦的基于人工智能的惡意軟件家族分類競賽, 訓練集中共有5 841個惡意軟件樣本, 分別屬于10個不同的惡意軟件類別或家族, 其樣本數分布如圖3(b)所示. 兩個數據集中的樣本均包含去除PE頭的PE文件和.asm文件, 分別保存惡意軟件的十六進制序列和反匯編代碼.

圖3 數據集訓練樣本分布Fig.3 Dataset training sample distribution

2.2 評價指標

對于微軟數據集, 可以將測試集的測試結果上傳到Kaggle, 官方將以下式計算多分類對數損失. 即

(10)

式中:Ntest為測試集的樣本數;M為惡意軟件類別數; 當樣本i屬于y時,yij為1, 否則為0;pij為分類器預測的樣本i屬于類別j的概率.

對于CCF公開數據集, 只能在訓練集上進行兩折交叉驗證, 采用精確度(PACC)來衡量模型的性能. 其計算式為

(11)

式中:n表示預測樣本的總數;ntrue表示做出正確預測的樣本總數.

2.3 基準模型

選取5個模型作為基準模型. 1) M-CNN模型[9]. 將灰度圖作為輸入, 對 VGG16模型進行訓練. 2) IMCFN模型[10]. 將灰度圖映射為彩圖, 并對圖像進行數據增強, 用于微調VGG16模型. 3) DCNN模型[8]. 將馬爾可夫圖作為輸入, 用于訓練改良的VGG16模型. 4) NMCD模型[11]. 將灰度圖作為輸入, 訓練所提出的一種基于深度學習的新的混合模型. 5) S-DCNN模型[12]. 將灰度圖作為輸入數據, 訓練所提出的由3個不同卷積神經網絡組成的模型.

2.4 參數設定

本研究模型需要50個迭代輪次才能完成訓練, 且每訓練20個迭代輪次, 都需要將學習率降低10倍. 批大小、 學習率、 權重衰退、 動量、 非結構化剪枝率、 學生網絡數據的最優參數分別為8、 0.005、 0.000 5、 0.9、 40%、 4.

3 實驗結果與分析

3.1 有效性實驗

為驗證本研究模型的有效性, 分別測試該模型和基準模型在微軟數據集上的Loss和在CCF數據集上的PACC, 如表1所示. 其中, 基準模型的實驗結果均是根據對應模型原論文進行復現得出的結果. 在基準模型中, DCNN模型的參數量最小, 且與本研究模型的參數量最為相近, 但該模型的性能明顯不如本研究模型. 性能最好且計算量最低的基準模型是NMCD模型, 但該模型的參數量明顯高于本研究模型. 由此可知, 相較于基準模型, 本研究模型可以在降低模型規模和計算量的同時, 顯著提升模型性能.

表1 各模型在微軟和CCF數據集上的實驗結果對比Tab.1 Comparison of experimental results of various models in Microsoft and CCF datasets

3.2 單輸入實驗

表2為通過不同卷積神經網絡, 從灰度圖或馬爾可夫圖中抽取特征并進行分類的實驗結果. 在以灰度圖作為輸入數據訓練的模型中, 與VGG16模型相比, VGG19模型的性能顯著降低, 而ResNet-34模型的性能顯著提升. 在以馬爾可夫圖作為輸入數據構建的模型中, 與VGG16模型相比, VGG19模型的性能顯著降低, DCNN模型的性能略有提升, ResNet-34模型的性能顯著提升. 這主要是因為, 與VGG16模型相比, VGG19模型雖然有更深層的網絡結構, 但其全連接層的神經元梯度消失或者爆炸, 會導致模型性能的大幅降低. DCNN模型性能會有略微的提升, 是因為該模型的全連接層具有較少的神經元. 在ResNet模型的網絡結構中, 殘差塊的恒等映射可較好地減少模型退化問題所帶來的影響, 該模型可以從圖像中抽取到更為豐富的圖像語義特征, 因此其性能得到有效提升.

表2 單輸入實驗結果Tab.2 Experimental results of single input

3.3 消融實驗

為驗證本研究模型中特征融合層和特征修正層的有效性, 進行相應的消融實驗, 分別測試不同結構下模型的LTL和PACC, 結果如表3所示. 可以看出, 特征融合層和特征修正層都可以進一步提升模型性能, 同時包含特征融合層和特征修正層的模型, 其性能最優. 這主要是因為, 特征抽取層通過殘差網絡從灰度圖和馬爾可夫圖中抽取到的特征向量, 分別具備惡意軟件的序列特征和分布特征. 特征融合層利用注意力機制對這兩種不同特征進行注意力計算, 可以挖掘不同特征之間的關聯性, 捕獲不同特征之間的依賴關系. 此外, 特征修正層利用高速神經網絡, 對分類結果產生誤導的向量進行映射, 進而對輸出向量進行修正.

3.4 模型壓縮實驗

表4為教師網絡、 直接訓練的學生網絡和通過本研究模型壓縮算法訓練的學生網絡的參數量、 浮點運算數、 單個樣本預測時間(t)、LTL和PACC. 可以看出, 與教師網絡相比, 學生網絡具有更小的規模、 計算量和更短的預測時間. 與教師網絡相比, 直接訓練的學生網絡, 其性能將會有所下降, 通過本研究模型壓縮方法訓練的學生網絡, 其性能則會有所提升. 這主要是因為, 通過本研究模型壓縮算法對學生網絡進行訓練時, 教師網絡會向其傳遞有效的知識點, 經此訓練的學生網絡可以比從頭訓練的學生網絡更快地找到真正有用知識點. 不同學生網絡互相之間進行協作學習, 可及時對性能較差的學生網絡進行調整. 此外, 非結構化剪枝操作對蒸餾過程起到正則化的作用, 混合損失函數中所包含聚焦損失函數和難樣本損失函數, 可以對困難惡意軟件樣本進行挖掘.

3.5 可解釋性實驗

為驗證通過卷積神經網絡從惡意軟件圖像中抽取特征進行分類的科學性, 實驗生成“熊貓燒香”惡意軟件的熱力圖, 并根據熱力圖中高熱值區域像素點的偏移, 定位原始文件中對最終分類結果產生較大影響的字節. 圖4為“熊貓燒香”對應的灰度圖和熱力圖. 圖5為熱力圖中熱值較高區域所對應的字節通過IDA Pro反匯編的結果. 該區域的字節所對應的字符串為“熊貓燒香”病毒專有的字符串, 具有唯一性. 由此可見, 基于惡意軟件圖像和卷積神經網絡的分類模型可以較為準確地捕獲到惡意軟件中的關鍵字節并用于分類.

圖4 “熊貓燒香”的灰度圖和熱力圖Fig.4 Grayscale and thermograph of “Worm.WhBoy”

圖5 “熊貓燒香”的專有字符串Fig.5 Proprietary string of “Worm.WhBoy”

4 結語

提出一種面向特征融合與知識蒸餾的惡意軟件分類方法, 基于注意力機制融合惡意軟件的序列特征和分布特征, 對不同特征之間的關聯性進行挖掘. 通過知識蒸餾和協作學習技術, 降低分類模型的參數量和計算量. 在微軟和CCF兩個公開數據集上的實驗結果表明, 與基準模型相比, 本研究模型具有更加良好的性能和更小的規模. 此外, 實驗分析證明, 卷積神經網絡可以從惡意軟件圖像中抽取關鍵特征, 并將其用于模型學習和分類. 后續研究將在對惡意軟件進行可視化之前, 對惡意軟件中無用的垃圾字節進行消除, 以降低作為模型輸入的惡意軟件圖像噪聲.