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

arXiv: https://arxiv.org/abs/2605.17575 · HTML: https://mlsanity.com/paper/2605.17575/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 42 of its 72 references that are in the map.
- First posted on arXiv 17 May 2026.

## Builds on (most cited first, 30 of 42)

- [Rethinking the Inception Architecture for Computer Vision](https://mlsanity.com/paper/1512.00567/) (2015)
- [Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks](https://mlsanity.com/paper/1703.03400/) (2017)
- [Domain-Adversarial Training of Neural Networks](https://mlsanity.com/paper/1505.07818/) (2015)
- [UMAP: Uniform Manifold Approximation and Projection for Dimension Reduction](https://mlsanity.com/paper/1802.03426/) (2018)
- [A Baseline for Detecting Misclassified and Out-of-Distribution Examples in Neural Networks](https://mlsanity.com/paper/1610.02136/) (2016)
- [Sharpness-Aware Minimization for Efficiently Improving Generalization](https://mlsanity.com/paper/2010.01412/) (2020)
- [Multi-Task Learning Using Uncertainty to Weigh Losses for Scene Geometry and Semantics](https://mlsanity.com/paper/1705.07115/) (2017)
- [Enhancing The Reliability of Out-of-distribution Image Detection in Neural Networks](https://mlsanity.com/paper/1706.02690/) (2017)
- [Energy-based Out-of-distribution Detection](https://mlsanity.com/paper/2010.03759/) (2020)
- [Averaging Weights Leads to Wider Optima and Better Generalization](https://mlsanity.com/paper/1803.05407/) (2018)
- [Deep CORAL: Correlation Alignment for Deep Domain Adaptation](https://mlsanity.com/paper/1607.01719/) (2016)
- [When Does Label Smoothing Help?](https://mlsanity.com/paper/1906.02629/) (2019)
- [GradNorm: Gradient Normalization for Adaptive Loss Balancing in Deep Multitask Networks](https://mlsanity.com/paper/1711.02257/) (2017)
- [Domain Generalization: A Survey](https://mlsanity.com/paper/2103.02503/) (2021)
- [Learning to Generalize: Meta-Learning for Domain Generalization](https://mlsanity.com/paper/1710.03463/) (2017)
- [Out-of-Distribution Detection with Deep Nearest Neighbors](https://mlsanity.com/paper/2204.06507/) (2022)
- [Regularizing Neural Networks by Penalizing Confident Output Distributions](https://mlsanity.com/paper/1701.06548/) (2017)
- [SWAD: Domain Generalization by Seeking Flat Minima](https://mlsanity.com/paper/2102.08604/) (2021)
- [On the Importance of Gradients for Detecting Distributional Shifts in the Wild](https://mlsanity.com/paper/2110.00218/) (2021)
- [SSD: A Unified Framework for Self-Supervised Outlier Detection](https://mlsanity.com/paper/2103.12051/) (2021)
- [Surrogate Gap Minimization Improves Sharpness-Aware Training](https://mlsanity.com/paper/2203.08065/) (2022)
- [Domain Generalization via Conditional Invariant Representation](https://mlsanity.com/paper/1807.08479/) (2018)
- [Ensemble of Averages: Improving Model Selection and Boosting Performance in Domain Generalization](https://mlsanity.com/paper/2110.10832/) (2021)
- [Kitsune: An Ensemble of Autoencoders for Online Network Intrusion Detection](https://mlsanity.com/paper/1802.09089/) (2018)
- [Deep Packet: A Novel Approach For Encrypted Traffic Classification Using Deep Learning](https://mlsanity.com/paper/1709.02656/) (2017)
- [ET-BERT: A Contextualized Datagram Representation with Pre-training Transformers for Encrypted Traffic Classification](https://mlsanity.com/paper/2202.06335/) (2022)
- [Nearest Neighbor Guidance for Out-of-Distribution Detection](https://mlsanity.com/paper/2309.14888/) (2023)
- [A Survey on Deep Active Learning: Recent Advances and New Frontiers](https://mlsanity.com/paper/2405.00334/) (2024)
- [Deep Fingerprinting: Undermining Website Fingerprinting Defenses with Deep Learning](https://mlsanity.com/paper/1801.02265/) (2018)
- [k-fingerprinting: a Robust Scalable Website Fingerprinting Technique](https://mlsanity.com/paper/1509.00789/) (2015)

## Built on by (most cited first, 1 of 1)

- [Unknown-Traffic Detection, Calibration and Shortcut Reliance in Distilled Encrypted-Traffic Classifiers over One Year](https://mlsanity.com/paper/2609.31141/) (2026)

## Abstract (arXiv)

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 domain alignment fine-tuning, which encourages the learning of domain-invariant traffic representations across heterogeneous network conditions, with 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.