学术报告：基于Masked Autoencoder的恶意流量分类研究

本研究题为《Flow-MAE: Leveraging Masked Autoencoder for Accurate, Efficient and Robust Malicious Traffic Classification》，由Zijun Hang、Yuliang Lu、Yongjie Wang和Yi Xie完成，研究所在中国的National University of Defense Technology。这项研究发表于2023年10月16日至18日召开的“26th International Symposium on Research in Attacks, Intrusions and Defenses (RAID ’23)”会议上，是一项重要的原创性研究工作。

研究背景、必要性和目标

网络安全重要性与挑战

恶意流量的分类是入侵检测系统(IDS)的核心功能之一，其目的是通过区分不同类型的恶意网络流量来检测网络安全漏洞。然而，目前常用的传统机器学习方法在特征提取上需要依赖领域专家知识，同时需要大量标注数据，这既费时又费力。此外，绝大多数实际的网络流量是良性的，恶意流量占比极小，这进一步增加了分类的难度。同时，现有的方法往往专注于特定任务，面对新任务时需要重新设计或重新训练模型，导致效率较低、鲁棒性不足。

深度学习与Transformer的引入

近年来，Transformer架构的深度学习模型逐渐突破传统机器学习方法的局限性。这种模型最初在自然语言处理(NLP)和计算机视觉(CV)领域中获得了显著成果。例如，ET-BERT将BERT模型引入网络流量分类，实现了当时的最优性能。然而，ET-BERT存在两个主要局限：(1) 输入长度固定为128字节，不能充分代表范围在千字节到兆字节的网络流量，且这种短输入长度带来较高的计算和内存开销；(2) 模型采用Byte Pair Encoding (BPE)固定字典，对流量模式变化的鲁棒性有限。

流量分类领域的创新性转变

研究团队提出了一种基于计算机视觉领域Masked Autoencoders (MAE)的预训练模型Flow-MAE。该模型通过“自监督学习”的方式对网络流量的通用表示进行学习，从而在恶意流量分类任务中实现高精度、高效性和强鲁棒性。研究目标包括： 1. 解决输入长度限制并更高效地处理大规模网络流量。 2. 引入新的自监督预训练任务“Masked Patch Model (MPM)”，从不同长度和不同模式的流量突发(Burst)中捕捉无偏的表示。 3. 通过实验验证该方法在精度、效率和鲁棒性上的提升。

工作流程

数据预处理与表示

Flow-MAE的核心基于“Burst-Level”网络流量表示，这是针对会话(Session)的子级划分。研究采用以下预处理步骤： 1. 会话分割：依据五元组（源IP地址、目标IP地址、协议类型、源端口和目标端口）将网络流量划分为不同会话。 2. 突发分割：将一个会话（Session）划分为多个短时间窗口内的双向流量突发（Bursts）。 3. 裁剪和填充：对突发流中的字节序列进行填充或裁剪，确保输入长度固定为1024字节。

Patch Embedding与位置嵌入

每个突发按照固定窗口（即Patch）进行分割后，数据被编码为768维特征向量。这一过程利用一维卷积层完成。此外，为了保留网络流量中数据包的序列关系，模型引入了位置嵌入。

预训练任务：Masked Patch Model (MPM)

研究采用自监督学习方式，预训练任务包括： 1. 随机遮盖15%输入Patch。 2. 编码器（Encoder）对可见Patch生成潜在表示。 3. 解码器（Decoder）通过编码器输出重建被遮盖的Patch。

这部分的优化目标是最小化重构的均方误差(MSE)。

微调（Fine-Tuning）

预训练后的Flow-MAE编码器通过微调阶段适配到恶意流量分类这一下游任务。微调利用小规模、特定任务的标注数据，替换解码器为一个全连接线性分类器。

数据集与实验设置

研究使用6个公开数据集进行训练和测试，包括CIC-IDS2018、USTC-TFC-2016、ISCX-VPN-2016、ISCX-TOR-2016和Cross-Platform(Android)。模型先在背景流量上预训练，再在目标任务数据上微调。所有实验使用的硬件平台包括i7-12700k CPU和两块NVIDIA Geforce 3090Ti GPU。

研究结果

模型性能

Flow-MAE在所有数据集上的分类精度均超过99%，相比现有最优模型ET-BERT，其精度提升了0.41%-1.93%，速度提升了7.8-10.3倍。此外，Flow-MAE仅占用ET-BERT的0.2% FLOPs和44%的内存。

1. CIC-IDS2018：Flow-MAE在此最大数据集的准确率为0.9958，比ET-BERT进一步提升0.03%-0.26%。这种性能提升得益于模型更长的输入序列以及结合数据包头与有效负载的特点。
2. ISCX-TOR-2016：在此数据集中，Flow-MAE展现了卓越的鲁棒性，对多层加密流量和对抗性混淆具有良好的区分力。
3. Cross-Platform(Android)：实验表明，在多达71类的应用流量分类任务中，Flow-MAE表现出优异的迁移性能，分类准确率高达99.21%。

微调效率

Flow-MAE所需微调时间显著低于ET-BERT，仅需原模型的7%-15%。处理速度也大幅提升，每秒处理能力从ET-BERT的121提升至982。

少样本学习能力

在仅有30%样本量的情况下，Flow-MAE在所有数据集上的分类精度及F1评分均超过95%。即使在面对高难度的ISCX-TOR-2016和Cross-Platform(Android)数据集时，Flow-MAE依然保持了良好的鲁棒性和泛化能力。

研究意义与创新

1. 方法学创新：首次将计算机视觉领域的MAE引入网络流量分类领域，并通过增设Patch Embedding和Masked Patch Model扩展数据表示能力。
2. 性能提升：在六个数据集中实现了新一代性能表现，大幅提升了检测精度、鲁棒性和效率。
3. 实践价值：Flow-MAE的快速微调能力和适应性使其非常适合动态网络威胁检测场景。
4. 少样本与迁移学习能力：可在极少标注数据下或迁移至多类别任务时表现优异。

结论

Flow-MAE是网络安全领域一次重要的技术突破，它不仅在恶意流量分类性能上超越了现有方法，还通过创新的算法和优化的训练流程显著提升了模型效率和适应性。这项研究为基于Transformer模型的网络流量分类提供了新思路，并对未来在网络攻击检测、实时分类和加密流量分析等方向的应用具有重要意义。