《 A Novel Self-Supervised Framework Based on Masked Autoencoder for Traffic Classification 》

基于掩码自编码器的新型自监督流量分类框架

论文于2024年6月发表于IEEE/ACM Transactions on Networking(TON)原文链接

文章开源了代码和数据集,项目地址:https://github.com/NSSL-SJTU/YaTC

PS:这篇文章的工作最先发表在AAAI 2023上,文章题目为《Yet Another Traffic Classifier: A Masked Autoencoder Based Traffic Transformer with Multi-Level Flow Representation》,补充完善后发表的TON 2024,模型没有变化,代码是通用的。

note

两三个月没有正经写过博客,一个是到春季学期感觉事情莫名地多了起来,每周光上课就占了三分之一的时间,还要找时间读论文做实验来应付每周组会上的工作汇报,再加上一些零散的诸如开会的事情,总是感觉一天时间莫名其妙就结束了。另一个是打开IDE也不知道写些什么,写教程?写论文笔记?(还是写点小作文?hhhhh),写教程的话最近做实验确实没碰到什么困难的问题,google搜一搜也就解决了,写一些什么配环境装软件的教程又没这个必要,csdn上大把的教程;想了一下那还是写点笔记,正好上上周上课汇报了这篇强专业相关的论文,也是最近感兴趣的小方向,本着不浪费的原则,顺便就写成博客记录下来。

——————————————–分割线———————————————

引言

论文核心工作是提出了一个新型的流量分类器,利用高性能的CV模型来尝试解决问题,最近几年流量分类的工作其实不是很多,大部分的研究工作都是做指纹(fingerprinting),网站指纹(WF)、应用指纹(AF)或者是一些小众的网页指纹(WPF),其实本质上都差不多,只不是考虑的细粒度不一样,指纹可以看成是更细粒度的流量分类(of course这是不严谨的说法,两者的目标其实是不太一样的,只做到粗粒度的类别划分更像是运营商ISP会去关注的问题,而指纹精确到具体的网站、应用甚至是网页的话对隐私的破坏非常强,说白了就是攻击行为,以前这方面的研究在国外是个挺敏感的话题,相关的文章经常被拒稿,最近五六年风向转变,相关的论文发表量激增,每年四大上都会有这个方向的论文)。

论文背景

流量分类的应用和挑战

流量分类技术广泛运用于网络安全管理、入侵检测系统和服务质量管理等领域,但是随着加密流量的占比快速增加以及Tor等匿名网络技术的发展,以深度包检测(DPI)为代表的传统技术已经基本上无力应对当下复杂的网络流量。

根据谷歌透明度报告(如无法访问请科学上网)显示,当前Chrome中绝大部分流量都是HTTPS加密流量,根据StatCounter2023年的统计,Chrome有68%的浏览器市场份额,可以推断全球整体网络环境中HTTPS流量占比和这个图是类似的。

根据Tor的统计报告(如无法访问请科学上网),整个2024年Tor网络的活跃用户量最低也有大约200万,峰值约780万(相当于新西兰全国人口规模),拥有庞大的用户群。

现有方法对比

方法类型 简述 优势 局限性
规则匹配 基于预定义规则(如端口号、协议类型)匹配流量类型 简单高效,实时性强 无法识别加密流量(如HTTPS)或动态端口应用(如P2P)
机器学习 提取统计特征(如包长、传输间隔)训练分类模型 自动化特征提取,适应部分复杂场景 依赖人工特征工程,特征冗余影响性能
深度学习 处理数据包原始字节(如CNN/RNN自动学习特征) 无需人工设计特征,表征能力强 长数据包淹没关键信息(如首包过长)需大量标注数据

文中对于深度学习方法的描述有一点偏差,并不是深度学习方法都是处理数据包的原始字节,实际上相当大的一部分工作是处理传统的统计特征,直接处理原始数据包字节数据的深度学习绝大部分都是预训练的方法,此处的深度学习方法进一步指明是预训练方法可能更合理一些?

主要贡献

现有的方法核心有三个不足:

表征缺陷:现有方法直接使用原始字节,导致信息丢失(如长包淹没短包)。

模型结构不适配:通用深度学习模型忽略流量层次性(包/流依赖)。

数据依赖:标注成本高,难以适应动态场景(如新应用/攻击类型)。

在受到CV与NLP的成功案例启发:MAE(掩码自编码器)和BERT通过自监督学习减少对标注数据的依赖,文章工作通过设计层次化多级流量表示(MFR)矩阵、包级与流级双层级注意力机制的专用Transformer结构,并结合基于MAE的自监督预训练框架,有效解决流量分类中的信息丢失、计算复杂度和标注数据依赖问题,显著提升分类精度与泛化能力。

方法设计

方法总览

文章提出的模型YaTC整体结构分为预训练和微调两阶段,过程上和其他的预训练-微调模型没有太大区别。

多级流量表示矩阵(MFR)

MFR矩阵构建过程主要有下面四个步骤:

  1. 根据五元组(源/目的IP、源/目的端口、协议)将原始流量划分为流(双向流)。
  2. 每条流取前5个数据包。(这个数量非常少,以前的传统方法大多数都需要使用一条流的所有数据包的信息)
  3. 删除以太网头,置零端口号,替换IP为随机地址(PS:以前有些使用字节数据的工作不对包头部进行处理,将原始的IP和端口号也放进去训练,按照我们的理解,这本身就是很强的特征,会导致模型在这一批数据集上表现良好,但是在陌生的IP和端口的数据集上表现很差,不符合实际的场景)。其中每个包级矩阵包含2行包头,6行负载,共8行,每行40字节,不足时进行补零对齐。
  4. 对每条流的5个包级矩阵按顺序纵向堆叠(文中将这一过程称为流级堆叠),形成MFR矩阵(共计40行,每行40字节)。

Traffic Transformer

在上一部分中,我们已经将每一条流特征表示为一个MFR矩阵,对于每个MFR矩阵,进一步将其转化为一个40×40的灰度图(每一字节对应一个灰度值),那么现在我们的问题就转变为对这些灰度图像进行分类,变成了一个CV问题。文章在Vision Transform(ViT)的基础上,设计实现了图示的Traffic Transform。主要变化体现在下面三点:

  • 嵌入模块:分割为2x2非重叠块,每个块通过线性层映射为192维向量并添加位置编码,生成块嵌入。
  • 包级注意力:使用ViT结构,仅在同一个数据包的块间执行多头自注意力(MSA)。
  • 流级注意力(主要变化):对包级特征进行行池化(Row Pooling, RP),生成行块并在行块间应用MSA,捕获流级依赖关系。最后通过列池化(Column Pooling, CP)得到全局特征。

下面是谷歌提出的ViT模型结构,可以看出Traffic Transformer主要变动为增加了一层流级注意力,其他部分未作太多修改:

MAE预训练

既然问题转变为对MFR矩阵灰度图的分类,那么很自然地就会想到CV里面常用的MAE预训练方法,大致过程就是随机掩码MFR矩阵中90%的块,仅输入未掩码块至编码器提取特征后,通过轻量解码器重构掩码区域。MAE预训练阶段使用全局注意力(和微调阶段的Traffic Transform的双层注意力不同)和高掩码率,迫使模型从极少信息中学习特征,增强鲁棒性。

实验分析

实验设置

  • 数据集
数据集名称 简述
ISCXVPN2016 VPN流量(6类)
ISCXTor2016 Tor匿名流量(10类)
USTC-TFC2016 常规加密恶意流量(10类)
CICIoT2022 物联网攻击流量(7类)
Cross-Platform 跨平台移动应用流量(18类)
说明: 前四个数据集合并为大规模未标注预训练数据用于预训练,微调时各数据集独立训练。Cross-Platform未参与预训练,用于迁移学习实验。
  • baseline
基线方法 方法类别
FlowPrint(统计特征) ML
AppScanner(随机森林) ML
DF(CNN) DL
Deeppacket(CNN) DL
2D、3D-CNN DL
FS-Net(时序模型) DL
CL-ETC 预训练
PEAN 预训练
PERT 预训练
ET-BERT 预训练

对比实验

相对于基线模型,YaTC在所有数据集上都取得了显著优势,在加密流量(ISCXTor)任务中表现最优。整体来看,预训练方法(YaTC/PERT)显著优于非预训练模型。

消融实验

说明:GA:全局注意力;PA:包级注意力;FA:流级注意力;FS:流级堆叠;PS:参数共享;MFR:多级流表示;PT:预训练

主要结论:

  • 包级注意力机制:移除后性能显著下降,验证其关键性(如ISCXTor2016准确率从99.72%降至78.02%),证明其对捕捉包内细粒度特征至关重要。
  • 流级注意力机制:移除后性能小幅下降(如ISCXVPN2016准确率从98.07%降至95.62%),表明其辅助性作用,用于建模全局流量关系。
  • 参数共享策略:移除后模型参数量增加但性能仅轻微降低(如CICIoT2022的F1从96.58%降至95.39%),证实共享机制在轻量化和性能间的有效平衡。
  • 预训练阶段:移除后准确率大幅下降(如ISCXTor2016从99.72%降至92.03%),凸显自监督学习对提取通用流量特征的不可替代性。
  • 多级流量表示(MFR):直接使用原始字节或移除流级堆叠均导致性能下降,证明分层结构(字节-包-流)能全面保留关键流量信息。
  • 头部与载荷作用差异:头部移除比载荷移除影响更大(如ISCXTor2016准确率降至42.86%),但载荷在新VPN场景贡献显著,体现两者互补性。

单从模型的角度来看,Traffic Transform的性能没有比单纯使用全局注意力Transform提升很多,例如在USTC-TFC2016数据集上F1分数甚至是一样的,Traffic Transform的优势主要体现在更小的参数量带来的更低开销。

迁移学习实验

使用全新的VPN协议(即WiredGuard,IkeV2和PPTP)评估了方法对新的现实流量数据的表征能力,能够反映模型的迁移能力。

PS:此数据集来源:《AppSniffer: Towards Robust Mobile App Fingerprinting Against VPN》

结论:YATC的分类性能远远领先于其他基线,即使面对使用新协议的流量数据,YATC也保持高可行性。

在Cross-Platform数据集上进行实验,其微调和测试数据与预训练无关,YATC将F1得分从69.93%显着提高到82.35%,比没有预训练的情况好12.42%,同时显著优于其他几种预训练的方法。

总结

  1. 核心创新
  • 多级流量表征(MFR)
    通过分层矩阵(字节/包/流级)整合流量信息,解决长报文覆盖关键内容的问题,支持跨协议流量分析。
  • 流量变换器(Traffic Transformer)
    结合包级/流级注意力机制,复杂度降低至 $O(N^2/M)$,参数共享策略减少50%参数量。
  • 两阶段训练范式
    自监督预训练(90%掩蔽率) + 微调,标注数据需求降低80%,模型迁移能力显著提升。
  1. 实验结果
  • 性能优势
    在5个真实数据集(VPN/Tor/IoT等)上平均F1达 97.3%,较SOTA方法(PERT/ET-BERT)提升12.4-34.7%。
  • 少样本鲁棒性
    仅10%标注数据时,分类准确率仍保持89.2%(较监督基线高21.5%)。
  • 推理效率
    GPU推理速度达 2929样本/秒,较同类Transformer方法快3-6倍。
  1. 未来方向
  • 轻量化部署:探索模型剪枝与量化技术
  • 实时性优化:动态时间窗口策略改进
  • 多模态扩展:结合元数据与统计特征