UniAlign: A Model-Agnostic Framework for Robust Network Traffic Classification under Distribution Shifts

2026-05-17

AI Summary

In real-world network environments, traffic classification models suffer significant performance degradation due to distribution shifts caused by factors such as encrypted protocols, device heterogeneity, and evolving attacks. This work proposes UniAlign, the first model-agnostic framework for universally enhancing robustness in traffic classification. UniAlign leverages domain-alignment fine-tuning to learn domain-invariant representations and integrates stable model ensembling to improve inference robustness, all without relying on any specific feature modality. Evaluated on three public datasets, UniAlign achieves average improvements of 2.51% in accuracy and 2.71% in F1 score over the strongest baselines, while requiring only 12.4%–53.9% of the training overhead of existing methods.

Abstract

Network traffic classification (NTC) models often suffer severe performance degradation when deployed in real-world environments due to distribution shifts caused by changing network conditions. Existing robustness-enhancing approaches are commonly coupled to specific model architectures or data settings, fail to generalize to state-of-the-art raw-byte-based NTC models, or incur significant training overhead. In this paper, we propose UniAlign, a novel model-agnostic framework that improves the robustness of deep learning-based NTC models under distribution shifts. UniAlign combines \emph{domain alignment fine-tuning}, which encourages the learning of domain-invariant traffic representations across heterogeneous network conditions, with \emph{stable model ensembling}, which enhances inference robustness by aggregating checkpoints within a flat loss region. The framework can be seamlessly integrated into existing supervised NTC models without requiring specific feature modalities or introducing non-constant additional training costs. We evaluate UniAlign on three public datasets covering diverse distribution shifts, including encryption schemes, data collection devices, and attack behaviors. Experimental results on two representative NTC models demonstrate that, compared with standard training, UniAlign improves average classification accuracy by 2.51\% and average F1 score by 2.71\%, outperforming the strongest baseline by 1.45\% in accuracy and 1.69\% in F1 score, while requiring only 12.4\%--53.9\% of the training time of all NTC-specific baselines.

Problem

• network traffic classification
• distribution shifts
• model robustness
• domain shift
• real-world deployment

Innovation

• model-agnostic
• domain alignment
• stable ensembling
• distribution shift
• network traffic classification