比Transformer更快更省，Mamba终于卷到网络领域了

今天分享清华大学的一篇文章：NetMamba: Efficient Network Traffic Classification via Pre-training Unidirectional Mamba。网络流量分类是一个重要的研究领域，旨在提高服务质量、简化网络管理并保护网络安全。现有流量分类模型面临两方面挑战：（1）模型效率不佳：广泛使用的Transformer架构具有平方复杂度，在高速流量识别场景不可用；（2）流量表征不可靠：现有表征方案往往忽略重要字节信息或保留不必要的偏置。为了解决上述问题，本文提出了全新的预训练流量模型NetMamba。NetMamba使用专为网络领域选择和优化的Mamba架构替代Transformer，提升了模型效率；同时，NetMamba设计了一套更全面的流量表征方案，在保留海量流量有效信息的同时去除偏置信息。本文在3类主要分类任务的6个公开数据集上对NetMamba进行评测，结果显示NetMamba准确率接近99%（部分超过99%）。同时，NetMamba比现有最佳模型的推理速度快2.24倍，且显存占用较低。此外，NetMamba具备更加出色的少样本学习能力。NetMamba是首个为网络领域量身定制的Mamba架构模型。

• •Title:NetMamba: Efficient Network Traffic Classification via Pre-training Unidirectional Mamba
• •Authors: Tongze Wang, Xiaohui Xie, Wenduo Wang, Chuyi Wang, Youjian Zhao, Yong Cui
• •URL:​​ https://arxiv.org/abs/2405.11449​​
• •Github: ​​https://github.com/wangtz19/NetMamba​​

一、概述

1 Motivation

•提高模型效率：当前基于深度学习的流量分类模型主要采用Transformer架构，由于其核心的Self-Attention机制存在关于序列长度的平方时空复杂度，导致模型在处理长序列时推理速度过慢且显存开销巨大，难以满足在线、高速流量检测的需求。

•优化流量表征：准确的流量表征方案对流量模型的效果至关重要。已有基于原始字节的流量表征方案存在忽略包头信息、保留非匿名偏置、忽视字节分配平衡、切分方法不合理等缺陷。

2 Method

为了提高模型效率，不同于已有的基于Transformer架构的方案，本文选择Mamba框架作为基础架构，并通过大量测试，选用了更适合流量序列特性的单向Mamba块。为了优化流量表征，本文设计了一套全新的网络流量表征方案，该方案保留了包头的有效信息，并通过报文匿名、字节平衡和stride切分去除偏置信息。

方法：NetMamba包括三个阶段：1）从原始流量中提取表征信息，存储为stride序列；2）基于MAE在无标签的流量数据上对NetMamba进行预训练，优化目标为最小化masked stride的重构误差；3）在带标签流量数据上对NetMamba进行微调，优化目标为最小化分类的交叉熵损失函数。

特点：1）网络领域首个基于Mamba架构的流量分类模型；2）更优的网络流量表征方案

3 Conclusion

•在多个公开网络流量数据集上有更好的分类表现。CrossPlatform(Android) (F1 +4.65%)， CrossPlatform(iOS) (F1 +5.86%)，ISCXVPN2016 (F1 +0.51%)，USTC-TFC2016 (F1 +0.40%)

•实现了2.24倍推理速度提升，并且显存占用很低。

•具备良好的few-shot性能，使用少量带标签数据训练出良好的分类能力。

4 Future Work

NetMamba除了能够解决流量分类问题外，其模型的泛化性也具备在QoE推断等任务上的可行性。目前，Mamba的性能优化依赖于特定GPU，未来也将进一步探索在多数算力受限的网络设备上部署的可能性。

二、详细内容

1 总体评估：NetMamba模型轻量且分类准确

结论1：相比于其他深度学习模型，**NetMamba仅拥有两百二十万预训练参数和一百九十万微调参数，是更轻量级的模型。**

结论2：在5个数据集上，NetMamba准确率领先全部baseline，除了在CICIoT2022数据集上落后TFE-GNN 0.15%，但该模型在其他数据集上表现远不如NetMamba。

2 性能对比：NetMamba推理速度最快，显存占用具备优势

多个batch大小比较：

batch大小为64（fine-tune环节的参数选择）时：

结论1：NetMamba推理速度最快。

结论2：NetMamba显存占用具备优势。仅略微弱于YaTC（优化显存占用，去除后的YaTC(OF)显存开销极大）和FS-Net（分类表现很差）。

3 消融实验：模型设计和数据处理都很重要

结论1：单向Mamba更适应流量的序列特性，比双向Mamba准确率高（且少一次前传，速度更快）

结论2：positional embedding引入的位置信息能提升分类能力

结论3：模型预训练学习流量通用表征，有助于提升分类能力

结论4：流量的包头和载荷信息都很重要，缺失导致分类能力严重降低

结论5：去除IP字段能减少偏置信息，提升分类能力 

结论6：一维stride切分去除了二维patch切分引入的干扰信息，提升分类能力

4 Few-Shot评估：能通过少量样本学习到有效流量表征

结论：NetMamba鲁棒性强，分类能力受训练数据量减少的影响较小，能通过少量样本学习到有效的网络流量知识

三、总结

1. Transformer架构存在时空复杂度高的问题，本文实现了基于单向Mamba的预训练流量分类模型，具备分类准、速度快、显存省的优点。

2.通过保留包头有效信息和去除偏置信息，本文设计了一套更优的流量表征方案，有效提升模型的分类能力。

3. 经过预训练，NetMamba具有很强的few-shot能力。

四、参考

T. Wang, X. Xie, W. Wang, C. Wang, Y. Zhao, and Y. Cui, “Netmamba: Efficient network traffic classification via pre-training unidirectional mamba,” 2024.