今天分享清華大學(xué)的一篇文章：NetMamba: Efficient Network Traffic Classification via Pre-training Unidirectional Mamba。網(wǎng)絡(luò)流量分類是一個重要的研究領(lǐng)域，旨在提高服務(wù)質(zhì)量、簡化網(wǎng)絡(luò)管理并保護網(wǎng)絡(luò)安全。現(xiàn)有流量分類模型面臨兩方面挑戰(zhàn)：（1）模型效率不佳：廣泛使用的Transformer架構(gòu)具有平方復(fù)雜度，在高速流量識別場景不可用；（2）流量表征不可靠：現(xiàn)有表征方案往往忽略重要字節(jié)信息或保留不必要的偏置。為了解決上述問題，本文提出了全新的預(yù)訓(xùn)練流量模型NetMamba。NetMamba使用專為網(wǎng)絡(luò)領(lǐng)域選擇和優(yōu)化的Mamba架構(gòu)替代Transformer，提升了模型效率；同時，NetMamba設(shè)計了一套更全面的流量表征方案，在保留海量流量有效信息的同時去除偏置信息。本文在3類主要分類任務(wù)的6個公開數(shù)據(jù)集上對NetMamba進行評測，結(jié)果顯示NetMamba準(zhǔn)確率接近99%（部分超過99%）。同時，NetMamba比現(xiàn)有最佳模型的推理速度快2.24倍，且顯存占用較低。此外，NetMamba具備更加出色的少樣本學(xué)習(xí)能力。NetMamba是首個為網(wǎng)絡(luò)領(lǐng)域量身定制的Mamba架構(gòu)模型。

• ?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

?提高模型效率：當(dāng)前基于深度學(xué)習(xí)的流量分類模型主要采用Transformer架構(gòu)，由于其核心的Self-Attention機制存在關(guān)于序列長度的平方時空復(fù)雜度，導(dǎo)致模型在處理長序列時推理速度過慢且顯存開銷巨大，難以滿足在線、高速流量檢測的需求。

?優(yōu)化流量表征：準(zhǔn)確的流量表征方案對流量模型的效果至關(guān)重要。已有基于原始字節(jié)的流量表征方案存在忽略包頭信息、保留非匿名偏置、忽視字節(jié)分配平衡、切分方法不合理等缺陷。

2 Method

為了提高模型效率，不同于已有的基于Transformer架構(gòu)的方案，本文選擇Mamba框架作為基礎(chǔ)架構(gòu)，并通過大量測試，選用了更適合流量序列特性的單向Mamba塊。為了優(yōu)化流量表征，本文設(shè)計了一套全新的網(wǎng)絡(luò)流量表征方案，該方案保留了包頭的有效信息，并通過報文匿名、字節(jié)平衡和stride切分去除偏置信息。

方法：NetMamba包括三個階段：1）從原始流量中提取表征信息，存儲為stride序列；2）基于MAE在無標(biāo)簽的流量數(shù)據(jù)上對NetMamba進行預(yù)訓(xùn)練，優(yōu)化目標(biāo)為最小化masked stride的重構(gòu)誤差；3）在帶標(biāo)簽流量數(shù)據(jù)上對NetMamba進行微調(diào)，優(yōu)化目標(biāo)為最小化分類的交叉熵?fù)p失函數(shù)。

特點：1）網(wǎng)絡(luò)領(lǐng)域首個基于Mamba架構(gòu)的流量分類模型；2）更優(yōu)的網(wǎng)絡(luò)流量表征方案

3 Conclusion

?在多個公開網(wǎng)絡(luò)流量數(shù)據(jù)集上有更好的分類表現(xiàn)。CrossPlatform(Android) (F1 +4.65%)， CrossPlatform(iOS) (F1 +5.86%)，ISCXVPN2016 (F1 +0.51%)，USTC-TFC2016 (F1 +0.40%)

?實現(xiàn)了2.24倍推理速度提升，并且顯存占用很低。

?具備良好的few-shot性能，使用少量帶標(biāo)簽數(shù)據(jù)訓(xùn)練出良好的分類能力。

4 Future Work

NetMamba除了能夠解決流量分類問題外，其模型的泛化性也具備在QoE推斷等任務(wù)上的可行性。目前，Mamba的性能優(yōu)化依賴于特定GPU，未來也將進一步探索在多數(shù)算力受限的網(wǎng)絡(luò)設(shè)備上部署的可能性。

二、詳細(xì)內(nèi)容

1 總體評估：NetMamba模型輕量且分類準(zhǔn)確

結(jié)論1：相比于其他深度學(xué)習(xí)模型，**NetMamba僅擁有兩百二十萬預(yù)訓(xùn)練參數(shù)和一百九十萬微調(diào)參數(shù)，是更輕量級的模型。**

結(jié)論2：在5個數(shù)據(jù)集上，NetMamba準(zhǔn)確率領(lǐng)先全部baseline，除了在CICIoT2022數(shù)據(jù)集上落后TFE-GNN 0.15%，但該模型在其他數(shù)據(jù)集上表現(xiàn)遠(yuǎn)不如NetMamba。

2 性能對比：NetMamba推理速度最快，顯存占用具備優(yōu)勢

多個batch大小比較：

batch大小為64（fine-tune環(huán)節(jié)的參數(shù)選擇）時：

結(jié)論1：NetMamba推理速度最快。

結(jié)論2：NetMamba顯存占用具備優(yōu)勢。僅略微弱于YaTC（優(yōu)化顯存占用，去除后的YaTC(OF)顯存開銷極大）和FS-Net（分類表現(xiàn)很差）。

3 消融實驗：模型設(shè)計和數(shù)據(jù)處理都很重要

結(jié)論1：單向Mamba更適應(yīng)流量的序列特性，比雙向Mamba準(zhǔn)確率高（且少一次前傳，速度更快）

結(jié)論2：positional embedding引入的位置信息能提升分類能力

結(jié)論3：模型預(yù)訓(xùn)練學(xué)習(xí)流量通用表征，有助于提升分類能力

結(jié)論4：流量的包頭和載荷信息都很重要，缺失導(dǎo)致分類能力嚴(yán)重降低

結(jié)論5：去除IP字段能減少偏置信息，提升分類能力 

結(jié)論6：一維stride切分去除了二維patch切分引入的干擾信息，提升分類能力

4 Few-Shot評估：能通過少量樣本學(xué)習(xí)到有效流量表征

結(jié)論：NetMamba魯棒性強，分類能力受訓(xùn)練數(shù)據(jù)量減少的影響較小，能通過少量樣本學(xué)習(xí)到有效的網(wǎng)絡(luò)流量知識

三、總結(jié)

1. Transformer架構(gòu)存在時空復(fù)雜度高的問題，本文實現(xiàn)了基于單向Mamba的預(yù)訓(xùn)練流量分類模型，具備分類準(zhǔn)、速度快、顯存省的優(yōu)點。

2.通過保留包頭有效信息和去除偏置信息，本文設(shè)計了一套更優(yōu)的流量表征方案，有效提升模型的分類能力。

3. 經(jīng)過預(yù)訓(xùn)練，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.