一種基于自蒸餾的自適應惡意流量分類算法

潘 嘉，翟江濤

（1.江蘇科技大學電子信息學院，江蘇鎮江 212003；2.南京信息工程大學計算機與軟件學院，江蘇南京 210044）

0 引言

隨著5G 技術的快速發展，網絡流量也井噴式產生。與此同時，黑客技術也不斷發展，各種網絡攻擊產生的惡意流量給網絡管理與服務帶來嚴峻挑戰，嚴重威脅網絡安全。網絡流量分類技術作為網絡安全管理系統的重要組成部分之一，其中一個至關重要的問題即是如何檢測與識別惡意流量［1］。

目前網絡流量分類主要采用3 種方法：基于端口的識別分類方法［2］、基于深度包檢測的分類方法［3］與基于統計的分類方法。基于端口的識別分類方法是早期的網絡流量分類方法，通過數據包端口號區分網絡流量類型。早在1992 年，RFC（Request For Comments）制定了標準的端口空間，端口識別方法自此開始發展并得到了廣泛運用［4］。然而，隨著互聯網技術的快速發展，動態端口降低了端口識別技術的可用性［5］。為彌補端口識別方法的不足，基于深度包檢測的分類方法應運而生，該方法通過分析數據包是否含有網絡應用所具有的獨特“指紋標簽”進行檢測。文獻［6］提出利用網絡流量的深度數據包進行流量檢測分類，基于深度數據包的流量檢測方法識別效率高，同時一定程度上彌補了端口識別的缺陷，但對于加密網絡流量，自有協議無法做到有效識別。近年來網絡流量數據日益龐大，研究者將機器學習方法引入網絡流量分類工作中。文獻［7］通過使用相關向量機（Relevant Vector Machine，RVM）對網絡流量進行分類，實驗結果表明，RVM 的性能指標優于支持向量機（Support Vector Machine，SVM）；文獻［8］采用多種機器學習算法進行網絡流量分類的對比實驗，結果表明，決策樹算法性能優于C4.5 及SVM 算法。然而，機器學習方法可移植性差，針對不同的流量數據，機器學習方法需要通過人工重新提取特征并進行訓練。為實現端到端的網絡流量分類，可采用深度學習技術。深度學習在目標檢測［9］、自然語言處理［10］等各領域都取得了突破性進展，因此將深度學習應用于網絡流量分類是當前的研究熱點。文獻［11］首次提出利用一維和二維卷積神經網絡用于惡意流量與普通流量分類，結果表明一維卷積神經網絡分類效果更佳；文獻［12］將流量數據包處理成三維數據格式，利用三維卷積神經網絡對流量進行分類識別；文獻［4］進一步利用卷積神經網絡（Convolutional Neural Network，CNN）與循環神經網絡（Recurrent Neural Network，RNN）的多種組合進行網絡流量分類。然而，網絡流量樣本分布的不均衡性特點導致流量樣本分類具有難易差別，對于簡單的流量樣本，通常淺層的神經網絡即能對其進行精準分類，而對于難樣本，則需要更深層的網絡資源。因此導致了一個問題：淺層的神經網絡無法滿足難樣本的分類需求，而將深層網絡應用于簡單樣本又會造成網絡資源浪費，降低運算速度。針對此問題，本文提出一種基于自蒸餾的自適應網絡流量分類算法，主要有以下兩方面貢獻：

（1）針對網絡流量的時序特點，設計基于自注意力機制的流量特征提取器SAC（Self-Attention-Convd）網絡，同時引入流量的全局結點信息解決Self-Attention 的低秩瓶頸問題。

（2）設計分支網絡自適應功能，通過利用主干網絡的教師分類器對分支網絡的學生分類器進行自蒸餾訓練，使分支網絡的學生分類器擬合教師分類器的概率分布，并使流量樣本自適應選擇網絡層進行輸出，以避免后續網絡不必要的推理過程。

1 模型方法

本文首先對數據集進行預處理，將預處理后的流量數據作為網絡模型的輸入。訓練階段先訓練主干網絡和教師分類器，之后凍結主干網絡和教師分類器參數。利用教師分類器作為軟標簽對各分支網絡的學生分類器進行自蒸餾指導訓練。測試階段依據分支網絡計算流量樣本的熵值，若流量樣本在該分支網絡的熵值小于設定閾值，則流量樣本在該分支網絡提前返回；否則進入下一層網絡進行判別分類，以此類推。

1.1 數據集預處理

數據集采用公開的USTC-TFC2016 數據集，數據集的流量格式為pcap［11］。一般網絡流量的切分方法分為5 種：TCP 連接、流、會話、服務、主機［13］。本文使用流的切分方法，將原始網絡流量切分成為五元組（源IP、源端口、目的IP、目的端口、傳輸層協議）的數據包。完整預處理流程如下：首先將原始網絡流量切分成多個獨立的網絡流，提取每條網絡流中的前n 個數據包，同時丟棄多余數據包，若網絡流不足n 個數據包，則進行填充0 處理，并進行流量清洗或匿名化處理［14］，即去掉網絡協議層中的IP 地址和MAC 地址；其次截取每個包中前m 個字節，并丟棄包中的多余字節，若不足m 個字節，則進行填充處理；接下來對于每條網絡流共截取n*m 個字節，在此基礎上截取前256 個字節（保證n*m 大于256），并將每一個字節轉化為0～255之間的十進制數值；最后添加流量的類別標簽，對每一條數據進行實例歸一化處理后映射為網絡模型的輸入［15］。

1.2 模型架構

本文的網絡模型架構主要由主干網絡與分支網絡構成。如圖1 所示，主干網絡由6 層相同的SAC 網絡層及教師分類器組成，其中SAC（Self-Attention-Convd）網絡層是本文針對流量數據特點所設計的神經網絡層，主要架構為多頭自注意力層及一維卷積神經網絡。該網絡層的總體設計思路如下：首先輸入流量序列經過多頭自注意力機制處理后，得到具有權值分布的流量序列，然后利用一維卷積神經網絡提取流量分布的顯著特征作為后續網絡輸入。教師分類器由反瓶頸層及多頭自注意力層構成，反瓶頸層首先擴張特征向量，然后由多頭自注意力層進行計算，再將多頭自注意力層的計算結果進行壓縮，最后由帶有Softmax激活函數的全連接層進行輸出分類。分支網絡層結構與教師分類器一致，作為每層SAC 網絡的額外輸出。

Fig.1 Network architecture圖1 網絡架構

1.3 自注意力機制

網絡流量序列本質為時序數據，RNN、LSTM（Long Short-Term Memory）等神經網絡天然契合時序類數據的特點，因此目前大多數文獻選擇此類神經網絡作為基準網絡進行流量分類。但是RNN、LSTM 神經網絡需依時序對流量序列進行計算，對于遠距離的相互依賴特征，需要經過若干時間步驟的信息積累才能將兩者相聯系，距離越遠，有效捕捉的可能性越小。Transformer 的Self-Attention 機制能夠將輸入序列中的任意兩個特征直接進行聯系，更有利于遠距離特征捕獲。然而，考慮到本文任務是對網絡流量序列進行分類，分類任務更關注序列整體信息，因此本文在Self-Attention 中引入全局節點S 構建序列的全局信息。原始的Self-Attention 計算步驟如下：

該公式描述了Q、K兩個向量之間的聯系，其中d為K向量維度。假設序列長度為n，則該分布共有n2個值。由于Transformer 中引入多頭注意力機制，使得建模存在瓶頸，h為多頭數量。當Q、K分別投影至低維時，總共參數量只有2n*（d/h），公式（1）相當于用只有2n*（d/n）的參數量逼近本身有n2值的量，變相減弱了模型表達能力。如果要增加模型表達能力，直觀的想法就是增加2n*（d/h）參數量。由于多頭機制本身確實能夠使模型在不同位置關注不同子空間的信息，如果減少h的數量則會得不償失，因此本文選擇引入全局信息S融入Q、K向量，進而增大d以解決低秩瓶頸問題。首先對輸入的query、key、value進行線性映射，其次取映射向量Q的高斯分布作為全局信息S，并分別與Q、K向量進行concat 融合操作，最后利用公式（1）進行自注意力計算。

在多頭自注意機制中，多頭數目h會影響多頭自注意力層的性能。為針對流量分類任務選取合理的多頭數目，本文進行了h的尋優實驗：對單層的SAC 網絡進行單次訓練，期間僅改變多頭自注意力層中的多頭數目，分別觀察網絡的分類準確率及損失變化，實驗結果如圖2 所示。隨著多頭數目的增加，準確率不斷上升，損失隨之下降。當多頭數目達到4 時，準確率趨于穩定，繼續增加多頭數目對于網絡性能基本無影響，因此本文在SAC 網絡層中選擇多頭數目h=4。此外，在實驗中發現去除scale 操作更有利于顯著特征的提取，進而可提升模型的整體準確率。

1.4 一維卷積神經網絡

SAC 網絡層的另一個核心架構為一維卷積神經網絡。流量數據的實質為一維時序數據，基于以上考慮，不同于Transformer 的前饋神經網絡結構，本文使用一維卷積神經網絡及最大池化操作替代前饋神經網絡，以更好地捕獲流量數據特征。不同超參數會影響網絡性能，因此本文設計了參數尋優實驗。一維卷積神經網絡的主要參數為卷積核尺寸及特征通道數。為選取合理的超參數，設計了6 種候選模塊，如表1 所示。針對參數選擇有如下說明：①流量序列經多頭自注意層力處理后維度為256，一般核尺度為3 或5，同時為避免邊緣信息丟失，引入Padding 以保持邊界信息；②對于流量數據所包含的信息，特征面一般不超過64，否則容易陷入局部最優狀況，因此特征面一般設為8、16、32。

Fig.2 Accuracy of different heads圖2 不同head精度

Table 1 Candidate modules表1 候選模塊

將以上設計的6 種候選模塊分別嵌入單層SAC 網絡，采用最大池化及平均池化進行多次對比實驗（單次訓練epoch=1），分別對網絡的分類精度及時間開銷進行檢測，實驗結果如表2 所示。對比檢測結果可發現，使用最大池化的模塊整體效果略優于平均池化，特征通道變化對于精度影響不大，但花費時間提升較為明顯，同時小尺度核精度更高。綜合以上實驗分析，最終選擇1 號候選模塊的參數作為一維卷積神經網絡的超參數。

Table 2 Results of experimental tests表2 檢測結果

1.5 架構優化

Transformer 中對于兩兩子層之間使用Add 和Layer-Norm 操作，且使用Gelu 作為前饋神經網絡層的激活函數［16］。本文所設計網絡層之間的操作不同于Transformer子層之間的操作，具體說明如下：對于殘差連接，可較大程度上減緩梯度消失問題，進而可訓練更深層次的神經網絡。本文殘差操作表示如下：

其中，α表示殘差連接權值，該參數初始化為0，后續可由網絡層自行學習最優參數；Xi、Xi-1分別表示網絡層的兩兩子層；f表示子層的函數映射。這種殘差連接方式可訓練更深的網絡層，且收斂速度更快［17］。

本文在模型延遲實驗中發現：在設計的網絡層中，LayerNorm 操作以及Gelu 激活函數占據了大量推理時間。為加速網絡的推理計算速度，在保證網絡準確率的前提下，采用HadamardProduct 操作替代LayerNorm，同時使用Swish 激活函數替代Gelu 激活函數。不同架構推理時間如表3 所示，分析得到LayerNorm 的推理時間幾乎是Hadamard Product 的3 倍，Swish 激活函數的推理時間遠小于Gelu 激活函數。基于以上結果，最終選擇的架構操作為HadamardProduct，激活函數為Swish。

Table 3 Reasoning time of different architectures表3 不同架構推理時間

1.6 自蒸餾

分支網絡（學生分類器）作為每層SAC 網絡的額外輸出，其網絡架構與教師分類器結構保持一致。由于網絡流量樣本分類有難易之分，對于簡單的流量樣本，通常單一的淺層網絡即能進行區分，而對于難樣本，往往需要更深層次的網絡，如果僅使用單層次的淺層或深層網絡則會導致誤分類、網絡資源浪費等后果。因此，本文設計了分支網絡，并使用自蒸餾的訓練方法解決該問題。分支網絡總體設計思路如下：首先訓練主干網絡及教師分類器，然后凍結主干網絡及教師分類器參數，利用訓練后的教師分類器作為軟標簽指導學生分類器進行自蒸餾訓練。本文使用KL 散度表示教師分布與學生分布的損失：

其中，pt表示教師輸出概率分布，ps表示學生輸出概率分布，N表示流量類別數目。由于主干網絡由6 層SAC 網絡堆疊而成，所以自蒸餾的總損失可表示為：

其中，L代表網絡層數，psi代表第i個分支網絡的輸出概率分布。流量樣本的自適應網絡層采用分支網絡進行處理：通過對流量樣本的熵值進行判別，通常樣本熵值越低，代表樣本的置信度越高，分類準確率越高。若樣本的熵值低于設定閾值，則流量樣本在該分支網絡提前返回，否則進入下一層網絡進行判別分類。網絡流量樣本的熵值定義如公式（5）所示，其中P（s）為分支網絡的輸出概率分布，N為流量類別數目。

2 實驗與結果分析

2.1 實驗環境及數據集簡介

實驗使用基于Python3.7 環境的Pytorch1.2.0，操作系統為Ubuntu18.04，并采用英偉達1050TiGPU 進行加速訓練。其中，Mini_batch 大小為128，Epochs 為40 輪，初始學習率為0.001，訓練過程中采用余弦退火方式更新學習率，模型評估指標為準確率、召回率及F1指數。

實驗使用公開的USTC-TFC2016 數據集，該數據集由中國科技大學的Wang等［11］創建，主要由兩部分構成：一部分是由捷克科技大學科研人員在真實環境中采集的10 種惡意流量，另一部分源于專業網絡流量仿真工具IXIA BPS仿真的10 種正常流量。該數據集格式與Wireshark 抓取原始流量的pcap 格式相同，數據集總量為3.71G。表4 顯示了該數據集中部分正常流量及惡意流量的屬性信息。

Table 4 Selected traffic attributes of the dataset表4 數據集部分流量屬性

2.2 實驗結果

本文實驗主要由兩部分組成：第一部分為正常流量及惡意流量檢測對比實驗，以驗證本文模型的可行性；第二部分為分支網絡自適應實驗，利用測試樣本驗證分支網絡是否具備自適應網絡層的功能。其他部分超參數實驗在本文1.3、1.4小節，涉及模型參數尋優及架構優化。

2.2.1 流量檢測實驗

流量檢測識別不僅在于檢測惡意流量類別，對于正常流量的檢測識別也十分必要。因此，流量檢測實驗分為兩部分：第一部分實驗檢測混合流量（含有正常流量及惡意流量），由此得到模型對正常流量及惡意流量的檢測率；第二部分實驗檢測惡意流量類別。本文選取USTC-TFC2016數據集中5 類正常流量與5 類惡意流量構成的數據集進行第一部分實驗。為避免模型過擬合，本文采取5 折交叉驗證方法進行對比實驗，兩部分檢測實驗對比對象分別為文獻［12］使用的三維卷積神經網絡模型與文獻［11］使用的一維卷積神經網絡模型。值得一提的是，文獻［11］首次將卷積神經網絡應用于惡意流量分類，后續很多研究也將卷積神經網絡用于流量檢測分類，所以一維卷積神經網絡是深度學習在該領域的典型算法。表5 詳細對比了文獻［12］及本文模型在混合流量中的檢測指標。

Table 5 Comparison of testing indicators表5 檢測指標對比（%）

實驗結果如下：無論是正常流量或惡意流量檢測，本文模型的表現均優于文獻［12］的模型。其中，對于正常流量的檢測表現，兩者相差甚微，對于惡意流量的檢測表現卻相差甚大。文獻［12］利用三維卷積神經網絡進行流量特征提取，將流量數據歸一化至0～16 字段，并對每個字節進行one-hot 編碼，之后轉化為三維數據。但流量本質為一維時序數據，三維卷積過程并不契合流量數據的特點，且one-hot 編碼稀疏性較大，無法捕獲流量內部的聯系。由表5 可進一步看出文獻［12］的模型在檢測Neris 及Virut兩類樣本時，檢測效果遠不如本文模型。

為進一步檢測惡意流量的具體類別，本文進行了第二部分惡意流量類別檢測。對比模型為文獻［11］的一維卷積神經網絡，表6 詳細對比了文獻［11］及本文模型針對10類惡意流量的檢測指標。在深度學習分類中，難易樣本的判定是由訓練損失（loss）決定的，難樣本的訓練損失較簡單樣本大的多，進而會最終影響精度、召回率等檢測指標。反映到具體指標上，一般可由精度、召回率進行判斷。由表6 可以看出，文獻［11］使用單一結構的淺層一維卷積神經網絡檢測Gridex 及Nsis-ay 兩類簡單樣本時效果很好，然而對于Neris、Virut 兩類難樣本的檢測表現不佳，精度、召回率僅為93%左右。由之前實驗得到文獻［12］使用的三維卷積神經網絡在檢測Neris 及Virut 兩類難樣本時，檢測效果也不佳。本文主干網絡針對Neris、Virut 兩類難樣本的檢測效果明顯優于文獻［12］及文獻［11］的模型，各項檢測指標提升了近5 個百分點，總體樣本檢測率高達99.9%左右。由于本文模型中自注意力機制的存在，使網絡能夠學習流量序列的內部特征聯系，同時主干網絡中的架構優化大大緩解了梯度衰減問題，從而使得主干網絡的分類效果更佳。

Table 6 Comparison of testing indicators表6 檢測指標對比（%）

2.2.2 分支網絡自適應實驗

本文所設計的分支網絡層自適應功能是目前網絡流量分類中所缺乏的，通過給予流量樣本自適應功能，使流量樣本自適應選擇網絡層進行輸出，可避免后續不必要的計算推理。為驗證分支網絡是否具備自適應功能，本文從測試集的流量樣本中隨機均勻抽樣500 個樣本（各類別樣本數量相差不大），觀察各分支網絡的流量樣本輸出，如圖3 展示的測試樣本在各分支網絡的輸出分布。從實驗結果可以看出，幾乎90%的流量樣本在第一層分支網絡輸出，從而避免了后續不必要的網絡層計算，加快了網絡推理速度。與此同時，實驗中發現第一層分支網絡輸出的流量樣本中Gridex 及Nsis-ay 等簡單樣本占88%，驗證了主干網絡的對比實驗結果。其余10%的流量樣本的輸出分布在第二層及第三層分支網絡，10%的流量樣本為Neris、Virut兩類難樣本，可見難樣本需要更深層次的網絡來提取特征。

Fig.3 Layer distribution of the sample圖3 樣本的層分布

3 結語

本文提出一種基于自蒸餾的自適應惡意流量分類算法，可避免人為提取流量特征，實現了端到端的自我學習，同時具備流量樣本自適應網絡的功能。在網絡架構設計中，主干網絡首先采用自注意力機制構建流量的權值分布，其次利用一維卷積神經網絡提取權值流量分布中的顯著特征，最后經過教師分類器輸出流量類別。分支網絡架構與主干網絡的教師分類器保持一致，作為每層SAC 網絡層的額外輸出，分支網絡根據每個流量樣本的熵值自適應選擇網絡。同時在網絡架構的優化中，設計帶權值的殘差連接減緩網絡梯度衰減問題，采用Hadamard Product 及Swish 激活函數降低網絡的推理延遲，進而提升網絡的檢測率。通過對比實驗，驗證了本文模型較現有方法的優勢。由于實際流量數據具有多源化、分布不均衡等特點，都會給分類算法帶來挑戰，未來將針對該問題以及更高效的深度學習模型進行研究。