[Paper Review] Lens: A Knowledge-Guided Foundation Model for Network Traffic

Xiaochang Li, Qian Chen+8 more|arXiv (Cornell University)|Feb 6, 2024

Network Traffic and Congestion Control4 citations

TL;DR

Lens is a foundation model for network traffic that leverages the T5 encoder-decoder architecture to learn representations from large-scale unlabeled traffic data using a novel multi-task pre-training objective combining Masked Span Prediction, Packet Order Prediction, and Homologous Traffic Prediction. It achieves state-of-the-art performance in both traffic understanding and generation tasks, reducing fine-tuning data requirements by 50%–95% compared to prior methods.

ABSTRACT

Network traffic refers to the amount of data being sent and received over the Internet or any system that connects computers. Analyzing network traffic is vital for security and management, yet remains challenging due to the heterogeneity of plain-text packet headers and encrypted payloads. To capture the latent semantics of traffic, recent studies have adopted Transformer-based pretraining techniques to learn network representations from massive traffic data. However, these methods pre-train on data-driven tasks but overlook network knowledge, such as masking partial digits of the indivisible network port numbers for prediction, thereby limiting semantic understanding. In addition, they struggle to extend classification to new classes during fine-tuning due to the distribution shift. Motivated by these limitations, we propose \Lens, a unified knowledge-guided foundation model for both network traffic classification and generation. In pretraining, we propose a Knowledge-Guided Mask Span Prediction method with textual context for learning knowledge-enriched representations. For extending to new classes in finetuning, we reframe the traffic classification as a closed-ended generation task and introduce context-aware finetuning to adapt to the distribution shift. Evaluation results across various benchmark datasets demonstrate that the proposed Lens~achieves superior performance on both classification and generation tasks. For traffic classification, Lens~outperforms competitive baselines substantially on 8 out of 12 tasks with an average accuracy of extbf{96.33\%} and extends to novel classes with significantly better performance. For traffic generation, Lens~generates better high-fidelity network traffic for network simulation, gaining up to extbf{30.46\%} and extbf{33.3\%} better accuracy and F1 in fuzzing tests. We will open-source the code upon publication.

Motivation & Objective

• To address the challenge of analyzing heterogeneous, encrypted, and semantically sparse network traffic data.
• To develop a unified foundation model that excels in both traffic understanding and generation, overcoming limitations of encoder-only or decoder-only architectures.
• To reduce dependency on large-scale labeled data for downstream tasks by leveraging effective pre-training on unlabeled traffic.
• To design a multi-task pre-training objective that captures global structure, sequence order, and relational patterns in network traffic.
• To enable strong zero-shot and few-shot generalization across diverse network traffic tasks using prompt-based fine-tuning.

Proposed method

• Tokenizes network traffic into hexadecimal sequences and evaluates multiple tokenization methods, including WordPiece with predefined vocabulary.
• Employs the T5 architecture as the backbone model to jointly learn representations for both understanding and generation via encoder-decoder attention.
• Designs a multi-task pre-training objective combining three tasks: Masked Span Prediction (MSP), Packet Order Prediction (POP), and Homologous Traffic Prediction (HTP).
• Uses a composite loss function that jointly optimizes MSP (denoising), POP (sequence order modeling), and HTP (cross-flow relation learning).
• Applies task-specific prompts during fine-tuning with minimal labeled data to adapt Lens to downstream tasks.
• Evaluates performance using standard metrics: accuracy/F1 for understanding, and JSD/TVD for generation.

Experimental results

Research questions

• RQ1 Can a foundation model based on the T5 architecture effectively learn latent representations from raw, unlabeled network traffic data?
• RQ2 Does a multi-task pre-training objective combining MSP, POP, and HTP improve both traffic understanding and generation performance?
• RQ3 To what extent does Lens reduce the need for labeled data in fine-tuning compared to existing pre-training methods?
• RQ4 How does the model generalize across diverse datasets and network traffic tasks, including those it was not pre-trained on?
• RQ5 What is the relative contribution of each pre-training task (MSP, POP, HTP) to the final model performance?

Key findings

• Lens achieves an average accuracy improvement of 10.75% on packet-level understanding tasks and 0.59% on flow-level tasks compared to baselines.
• For traffic generation, Lens reduces Jensen–Shannon Divergence (JSD) by 62% and Total Variation Distance (TVD) by 77% on average compared to NetShare.
• Lens outperforms NetShare in all generation tasks, including on the CrossPlatform dataset where it was not pre-trained.
• The model requires 50% to 95% less labeled data for fine-tuning than state-of-the-art methods like ET-BERT and YaTC.
• Ablation studies show that pre-training and MSP are critical for generation performance, with JSD increasing by 77% when MSP is removed.
• The full model with all three pre-training tasks achieves the highest performance, confirming the complementary nature of MSP, POP, and HTP.