Towards a graph-based foundation model for network traffic analysis

TL;DR

This paper introduces a graph-based foundation model for network traffic analysis that operates on a dynamic flow-level spatio-temporal graph. It uses a lightweight, GraphSAGE-based GNN with a line-graph representation (flow node plus source/destination IP nodes) and temporal edges, pretrained via self-supervised link prediction on unlabeled data. Finetuning on three downstream tasks—intrusion detection, traffic classification, and botnet detection—shows an average improvement of 6.87% over training from scratch, demonstrating effective capture of general network dynamics and transferability to unseen settings. The work highlights substantial efficiency gains and points toward scaling toward large-scale, operational network traffic foundation models with additional pretraining tasks and data.

Abstract

Foundation models have shown great promise in various fields of study. A potential application of such models is in computer network traffic analysis, where these models can grasp the complexities of network traffic dynamics and adapt to any specific task or network environment with minimal fine-tuning. Previous approaches have used tokenized hex-level packet data and the model architecture of large language transformer models. We propose a new, efficient graph-based alternative at the flow-level. Our approach represents network traffic as a dynamic spatio-temporal graph, employing a self-supervised link prediction pretraining task to capture the spatial and temporal dynamics in this network graph framework. To evaluate the effectiveness of our approach, we conduct a few-shot learning experiment for three distinct downstream network tasks: intrusion detection, traffic classification, and botnet classification. Models finetuned from our pretrained base achieve an average performance increase of 6.87\% over training from scratch, demonstrating their ability to effectively learn general network traffic dynamics during pretraining. This success suggests the potential for a large-scale version to serve as an operational foundational model.

Figures (3)

• Figure 1: Overview of proposed model and workflow. A graph neural network model is pretrained using an unsupervised link-prediction task. The resulting model is used for finetuning to different downstream tasks. The Graph Neural Network model incorporates spatio-temporal information and uses both structural and feature-based information to update node representations.
• Figure 2: Few shot learning experiment. For each of the downstream tasks (columns), one model is trained from scratch (orange) and one is finetuned from the pretrained model (blue). Each data point in the graph represents a different model. Training epochs are limited to 50. The results are expressed in percentage loss of optimal model performance of that metric (row) on the test set.
• Figure 3: Average normalized loss curves. For each of the downstream tasks (columns), the training and validation curves of the training routines of the few-shot experiment are normalized and then averaged at each epoch. This allows for a direct comparison of convergence speed and generalization capabilities between the finetuned models and the ones trained from scratch.