今天分享清華大學的一篇文章：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.