Keep Experts Diverse: A Task-Aware MoE for Multi-Task Traffic Analysis

Abstract

Network traffic analysis is crucial for maintaining the security of networks. Yet deploying accurate models on edge nodes remains challenging due to protocol diversity, complex traffic behaviors, and stringent resource constraints. Although recent deep learning models achieve strong performance, their dense architectures incur high inference latency, making them unsuitable for edge deployments.

Mixture-of-Experts offers a promising solution through conditional computation, yet existing methods overlook latent inter-task relations and suffer from severe expert load imbalance, leading to expert collapse and suboptimal efficiency.

To address these limitations, we propose TAMoE, a novel Task-Aware Mixture-of-Experts framework for multi-task network traffic analysis.

TAMoE integrates task semantics into both routing and representation learning through a novel Task-Aware Routing, enabling the router to dynamically select experts conditioned on both traffic features and task information.

Besides, to train stably usable and scalable multi-task models, we further design a collaborative multi-task training strategy that encompasses both a dataset mixed construction and a joint training process.

Experiments on six public benchmarks show that TAMoE achieves state-of-the-art results. Deployment on an NVIDIA® Jetson AGX Orin edge node demonstrates that TAMoE reduces inference latency by 60.6% compared to dense models and lowers the Gini Coefficient from 0.391 to 0.275 relative to Multi-Router MoE, achieving more balanced expert utilization.