论文阅读——NetMamba：通过预训练单向 Mamba 进行高效网络流量分类

文章基本信息

标题：NetMamba: Efficient Network Traffic Classification via Pre-training Unidirectional Mamba
作者：

• 王同泽（清华大学网络科学与网络空间研究院）
• 谢晓辉（清华大学计算机科学与技术系，通讯作者，电子邮件：xiexiaohui@tsinghua.edu.cn）
• 王文铎（清华大学智利书院）
• 王楚怡（清华大学计算机科学与技术系）
• 赵有建（清华大学计算机科学与技术系）
• 崔勇（清华大学计算机科学与技术系）
  机构：
• 清华大学网络科学与网络空间研究院
• 清华大学计算机科学与技术系
• 清华大学智利书院
  期刊：为发表期刊，目前为预印版本
  发表时间：提交于2024年5月19日
  解决问题：使用mamba网络进行流量分类，提高准确率和效率
  贡献：

1. 提出了NetMamba，这是第一个专门为网络流量分类设计的状态空间模型。与现有的基于Transformer的方法相比，NetMamba具有更好的性能和推理效率。
2. 提出了一种网络流量数据的综合表示方案，该方案保留了有价值的流量特征，同时消除了不必要的偏差。
3. 在一系列流量分类任务上进行了广泛的实验。提供了总体比较以及详细的评估–包括消融研究、效率分析和不太可能的学习调查。
   代码：https://github.com/wangtz19/NetMamba
   文章：https://arxiv.org/abs/2405.11449

摘要

摘要部分文章主要说明了网络流量分类的重要性，以及现有方法的两个挑战：

• 广泛使用的Transformer架构存在一个二次复杂度问题，对资源要求较高。
• Transformer丢弃了重要的字节信息而保留了不必要的偏置

为了解决这些问题，文章提出了NetMamba，一个高效的线性时间状态空间模型（an efficient linear-time state space model），并且还设计了一个配套的流量表示方案，用于从数据中提取有效信息并去除有偏差的信息。最后文章在6个数据集进行了评估，将推理速度提高了60倍。

结论

NetMamba在六个公共数据集上的评估实验表明，NetMamba具有优越的有效性、效率和健壮性。作者认为除了流量分类任务，该模型还可以应用在服务质量预测和网络性能预测中。然而，目前NetMamba的实现依赖于专用的GPU硬件，这限制了它在现实世界的网络设备上的部署。

引言

文章引言部分还是从近年来加密技术TLS的普及和匿名技术VPN和ToR的普及出发，说明加密流量分类的困难。紧接着说明使用传统机器学习方法、深度学习方法、预训练和Transformer方法进行流量分类的优缺点，总结了两个现有模型的挑战即摘要部分的两个问题。
接下来就提出了NetMamba模型，并简单介绍了模型：使用Mamba架构代替Transformer作为模型的backbone。然后文章简单介绍了Mamba模型，引出了单向Mamba因为其没有全向扫描和冗余块，适合学习顺序的网络流量中的潜在模式。
并且针对于文章提出的NetMamba模型，还设计了一个新的对流量数据预处理的方法。最后简单介绍了实验的结果。

相关工作

这一部分主要从下面几个方面介绍相关工作，不详细展开：

1. 基于Transformer的流量分类：介绍了几个使用Transformer进行流量分类的工作
2. 基于Mamba的表示学习：介绍了Mamba网络在CV和NLP领域的研究与应用
3. 流量表示方法：介绍了近来的流量分类工作在流量特征选取的局限性，比如只使用统计信息而不使用原始流量的字节数据。

基本概念

文章的这一节简单介绍了Mamba块的定义，也是理解NetMamba的基础。其中包括什么是状态空间模型SSM、如何将SSM进行离散化、如何使用卷积/循环表示SSM（实际上这两项都是S4中创新的工作）以及Mamba模型中选拔机制和硬件感知算法。（PS：文章中也没有详细介绍这部分内容，博主写的上一个学习笔记介绍了Mamba相关的内容，感兴趣的读者可以看一下上一篇文章）

NetMamba框架

NetMamba的整体框架如下：

大致的流程为NetMamba从原始二进制流量中提取分层信息，并将其转换为基于步长的表示（在Mamba中是参数$\Delta$），文章受到CV中的掩码自动编码器MAE预训练模型的启发，在NetMamba中也是用了双阶段训练方法：使用自监督预训练来获取流量的表示，再使用监督微调的方法来为流量感知任务定制模型。

原文表述：NetMamba extracts hierarchical information from raw binary traffic and converts it into stride-based representation. Inspired by the Masked AutoEncoders (MAE) pre-training model in computer vision, NetMamba employs a dual-stage training approach. Specifically, self-supervised pre-training is utilized to acquire traffic representation, while supervised fine-tuning is employed to tailor the model for downstream traffic understanding tasks.

（文章在框架章节简单地介绍了这三个模块，并在后续详细地分节介绍了每一部分，这里就不再按照原文的顺序而是直接在框架中详细介绍每一个模块。）

NetMamba架构

由于预训练过程中使用到了NetMamba中的结构，这部分的内容在原文中的位置是在流量表示和预训练之间。NetMamba和最经典的Mamba相比主要创新在以下两点：

1. 步长嵌入：对于给定的步长数组，文章首先将每个步长$s_i$线性投影到大小为$D_{enc}$，并结合位置嵌入${\mathbf{E}}_{\text{enc}}^{\text{pos}}\in {\mathbb{R}}^{N_s\times D_{\text{enc}}}$，最后的结果如下：

${\mathbf{X}}_0=\left[{\mathbf{s}}_1\mathbf{W};{\mathbf{s}}_2\mathbf{W};\cdots ;{\mathbf{s}}_{N_s}\mathbf{W};{\mathbf{x}}_{\text{cls}}\right]+{\mathbf{E}}_{\text{enc}}^{\text{pos}}$
其中$W$表示可学习的投影矩阵，相当于增加了参数，提高了模型的性能。并且作者收到VIT和BERT的启发，也加入了一个class token$x_{cls}$来表示整个步幅序列。由于单向Mamba是从前到后处理序列信息，所以将class token附加到序列的末尾，以增强信息聚合。

2. NetMamba块：作者在研究后发现，最初的单向Mamba结构很适合做流量分类任务，所以没有使用新的Mamba变体如双向Mamba或Graph-Mamba等。但是为了提高模型效率，作者去掉了全向扫描和冗余块，并且在后面的消融实验部分进行了论证。

最终，NetMamba的架构就是使用单向Mamba块实现NetMamba编码器和解码器。如下图所示：

可以对比一下S4和经典Mamba的算法：

作者用算法1概述了NetMamba块的前向传递操作过程：

NetMamba块的过程主要是：

• 输入：对于来自第$t-1$个NetMamba块的给定输入token序列$X_{t-1}$，具有批大小$B$和序列长度$L$
• 第1行：首先进行归一化
• 第2、3行：将其线性投影到维度大小为$E$的$x$和$z$中
• 第4行：对$x$进行1D卷积和SiLU函数激活，得到$x^{\prime }$
• 第5、6、7行：基于$x^{\prime }$，计算与输入有关的步长$\Delta$、以及维度为N的参数B和C
• 第8、9行：随后进行离散化
• 第10行：使用硬件感知的SSM计算y
• 第11行： $y$和经过激活函数的$z$进行点积得到$y^{\prime }$，即用$z$的SiLU激活结果对$y$进行门控
• 第12行：将$y^{\prime }$进行线性变换后，加上原始输入$X_{t-1}$进行残差连接，最终得到输出$X_t$

相比于S4和Mamba，NetMamba在SSM操作之间加入了更多的与处理步骤如归一化、线性变换和激活等，并且和Mamba一样允许参数B和C根据输入x动态变化。

流量表示

为了增强网络中的流量知识（获取到的特征），NetMamba采用了一种基于跨度的方法来全面表示网络流量，具体过程分为4步：

1. 分流（Flow Splitting）：这一步是处理网络流量的常规操作，将网络流量分割成不同的流，每个流属于特定协议的数据包组成，并按照它的五元组（源IP、目的IP、源端口、目的端口和协议）进行分类。
2. 数据包解析（Packet Parsing）：对于每一个流，所有数据包都经过几次顺序操作来处理，以保存有价值的信息并消除不必要的干扰。在缩小与特定应用或服务相关的流量数据范围时，作者排除了非IP协议（如ARP协议和DHCP协议）的数据包。除此之外，作者保存了包中包含的关键信息（如总长度字段）和有效载荷。此外，为了减轻可识别信息的偏差，作者删除了潜在的以太网头部和掩码IP来匿名化所有数据包。
3. 数据包裁剪、填充和连接：作者考虑到同一个流中数据包大小会发生变化，以及各个数据包中的报头长度和有效载荷长度都是会发生变化的，所以标准化数据包的大小，方法就是从单个流中选择前M个包，将报头长度设置为$N_h$字节，有效载荷长度设置为$N_p$字节，任何超过此长度的数据包都被裁剪，而对长度不足的包进行填充。最终，初始M个包的所有字节被连接为统一的数组$\left[b_1,b_2,\dots ,b_{L_b}\right]$，其中$L_b=M\times (N_h+N_p)$，$b_i$表示第$i$字节。
4. 步长分割：考虑到当给定的$L_b$过大（超过1000），会带来巨大的计算和内存消耗，文章受到时间序列预测中的patching方法的启发，在原始数据上采用了一维步长分割方法，将字节数组分成大小为$1\times L_s$的不重叠步长，从而得到步长总数$N_s=L_b/L_s$，每一个步长${\mathbf{s}}_i\in {\mathbb{R}}^{1\times L_s}$，被定义为$\left[b_{L_s\times i},b_{L_s\times i+1},\dots ,b_{L_s\times (i+1)-1}\right]$，其中$0\le i<N_s$.这个策略的目的是为了减少偏差，同时保留数据中的基本顺序信息。

预训练

我们从下往上看这个预训练部分的结构图，首先最下面一层是嵌入层，NetMamba首先使用了步长嵌入Stride Embedding+位置嵌入Position Embedding，然后是做了一个叫做随机掩码Random Masking的操作。

1. 随机掩码Random Mask：对于给定的嵌入步长token${\mathbf{X}}_0\in {\mathbb{R}}^{\mathrm{L}\times {\mathrm{D}}_{\mathrm{enc}}}$，随机抽样一部分步长token保留，其余的移除。使用参数$r\in (0,1)$来控制删除的比例，$r$越大说明移除的越多（PS：在预训练过程中，mask是一种关键的技术，它通过将部分输入语句遮蔽掉，强制模型去预测被遮蔽部分的含义，从而让模型学会关注上下文信息。）。处理后的可视tokens的长度可以表示为${\mathrm{L}}_{\mathrm{vis}}=\lceil (1-r)\mathrm{L}\rceil$，然后按照下面的公式对这些可视tokens进行采样：

${\mathbf{X}}_0^{\mathrm{vis}}=\mathbf{Shuffle}({\mathbf{X}}_0)[1:{\mathbf{L}}_{\mathrm{vis}},:]\in {\mathbb{R}}^{{\mathrm{L}}_{\mathrm{vis}}\times {\mathrm{D}}_{\mathrm{cnc}}}$
其中$Shuffle()$函数的作用是随机排列token序列，但是要保证最后的class token在整个过程中不被mask，因为它会在聚合整体序列信息中用到。（Random Mask的目的是）

2. 掩码预训练Masked Pre-training：NetMamba编码器的任务是使用可视token来捕获潜在的步长之间的联系。而解码器的目标是利用编码器输出的tokens和mask tokens来重建掩码步长。每一个masked token都是一个可学习的向量，并且还添加了新的位置嵌入向masked token提供位置信息。

随后作者用三个公式表示了接下来预训练的操作：
$\begin{array}{rl}& {\mathbf{X}}_{\mathrm{enc}}^{\mathrm{out}}=\mathrm{MLP}(\mathrm{Encoder}({\mathbf{X}}_0^{\mathrm{vis}}))\in {\mathbb{R}}^{{\mathrm{L}}_{\mathrm{vis}}\times {\mathrm{D}}_{\mathrm{dec}}}\\ & {\mathbf{x}}_{\mathrm{dec}}^{\mathrm{in}}=\mathrm{Unshuffle}(\mathrm{Concat}({\mathbf{X}}_{\mathrm{enc}}^{\mathrm{out}},{\mathbf{X}}_{\mathrm{mask}}))+{\mathbf{E}}_{\mathrm{dec}}^{\mathrm{pos}}\\ & {\mathbf{X}}_{\mathrm{dec}}^{\mathrm{out}}=\mathrm{Decoder}({\mathbf{X}}_{\mathrm{dec}}^{\mathrm{in}})\end{array}$
其中Unshuffle操作恢复原始序列的顺序，并且${\mathbf{E}}_{\mathrm{dec}}^{\mathrm{pos}}$表示解码器特定的位置嵌入。用于计算自监督重建的均方误差MSE损失如下：
$\begin{array}{rl}& {\mathbf{y}}_{\mathrm{real}}=\mathbf{Shuffle}({\mathbf{X}}_0)[{\mathbf{L}}_{\mathrm{vis}}+1:\mathbf{L},:]\\ & {\mathbf{y}}_{\mathrm{rec}}=\mathbf{Shuffle}({\mathbf{X}}_{\mathrm{dec}}^{\mathrm{out}})[{\mathbf{L}}_{\mathrm{vis}}+1:\mathbf{L},:]\\ & {\mathcal{L}}_{\mathrm{rec}}=\mathbf{MSE}({\mathbf{y}}_{\mathrm{real}},{\mathbf{y}}_{\mathrm{rec}})\end{array}$
其中，$y_{real}$表示真实掩码标记，$y_{rec}$表示预测的掩码标记。

微调

对于下游任务，所有编码器参数，包括嵌入模块和Mamba块，都是从预训练中加载的。为了对已标记的流量数据进行分类，用MLP头替换解码器。鉴于所有步长token都是可视的（没有被mask），NetMamba的微调使用监督方式进行：
$\begin{array}{rl}& \mathbf{X}=\mathbf{Encoder}({\mathbf{X}}_0)\in {\mathbb{R}}^{\mathrm{L}\times {\mathrm{D}}_{\mathrm{enc}}}\\ & \mathbf{f}=\mathbf{X}[\mathbf{L},:]\in {\mathbb{R}}^{{\mathrm{D}}_{\mathrm{enc}}}\\ & \hat{\mathbf{y}}=\mathbf{MLP}(\mathbf{Norm}(\mathbf{f}))\end{array}$
其中$f$表示尾部的class token，$\hat{\mathbf{y}}\in {\mathbb{R}}^{\mathsf{C}}$表示预测分布，其中$C$是流量的类别数量。然后，通过最小化预测值$\hat{y}$和真实值$y$之间的交叉熵损失来优化分类过程：
${\mathcal{L}}_{\mathrm{cls}}=\mathrm{CrossEntropy}(\hat{\mathbf{y}},\mathbf{y})$

实验

实验设置

数据集
作者将6个公开数据集分为了3类：

1. 加密应用程序分类任务：对各种加密协议下的应用程序流量进行分类。使用了CrossPlatform(Android)数据集、CrossPlatform(iOS)数据集、ISCXTor2016数据集中8个类别的Tor流量数据和ISCXVPN2016数据集中的7个类别的VPN流量数据。
2. 攻击流量分类任务：识别攻击流量。使用了CICIoT2022中的6个数据类别。
3. 恶意流量分类任务：区分恶意软件生成的流量和正常流量。使用了USTCUFC2016数据集的所有20个数据类别。

作者对这些数据集的类别做了预处理，首先对于数据量过多的类进行随机采样，对于数据量过少的类直接舍弃。
模型
作者对比的模型：

1. 传统的机器学习方法：APPScaner、FlowPrint
2. 深度学习方法：FS-Net、TFE-GNN
3. 基于Transformer的方法：ET-BERT、YaTC、YaTC(OF)

细节
在预训练中：

• batch size：128
• steps：15万
• optimizer：AdamW
• r：0.9（控制random mask比例的参数）

在微调中：

• batch size：64
• LR：2.0×10-3
• 数据集划分：8:1:1
• epochs：120

NetMamba结构：编码器部分使用4个mamba块，解码器使用2个mamba块
参数表：

评估指标：准确率、精确度、召回率、F1分数

实验结果

表4和表5是在8中方法在6个数据集上的准确率等性能对比：


图3是效率的对比，包括时间和空间（速度和显存）：

图4是推理效率的对比：

如果对详细的实验结果分析感兴趣的读者可以去看一下原文中作者的总结，这里不多赘述。

消融实验

作者对模型和流量表示两个方面进行了消融实验，通过对模型中不同的结构进行调整对比得出了文章中确定的NetMamba模型拥有最高的性能，以及文章的流量表示是最优的方法。

小样本实验

为了验证NetMamba的健壮性和泛化能力，作者在四个数据集上进行了少量评估，标签数据大小分别为完整训练集的10%、40%、70%和100%(占总数据的80%)。结果如下：