一种基于自蒸馏的自适应恶意流量分类算法

潘 嘉，翟江涛

（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 激活函数降低网络的推理延迟，进而提升网络的检测率。通过对比实验，验证了本文模型较现有方法的优势。由于实际流量数据具有多源化、分布不均衡等特点，都会给分类算法带来挑战，未来将针对该问题以及更高效的深度学习模型进行研究。