MalRAG: A Retrieval-Augmented LLM Framework for Open-set Malicious Traffic Identification

2025-11-17

AI Summary

Open-set malicious traffic identification faces three key challenges: difficulty in discriminating known classes while simultaneously detecting novel malicious traffic, strong task coupling in existing deep learning models, and poor cross-dataset generalization. To address these, this paper proposes MalRAG, a large language model (LLM)-based retrieval-augmented framework. Its core innovation is the first application of retrieval-augmented generation (RAG) to this domain: it constructs a multi-perspective traffic knowledge base, designs a coverage-enhanced retrieval algorithm coupled with traffic-aware dynamic pruning, and employs task-oriented prompt engineering—enabling cross-dataset adaptation without LLM fine-tuning. Experiments on multiple real-world datasets demonstrate that MalRAG significantly improves both known-class classification accuracy and unknown malicious traffic detection rate. Ablation studies confirm the effectiveness of each component.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: Safety and Robustness
• Data Mining & Knowledge Management: Conversational Systems for Recommendation & Retrieval

Application Category

• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• User Modeling, Personalization and Recommendation: Attacks and countermeasures in recommendation systems

Abstract

Fine-grained identification of IDS-flagged suspicious traffic is crucial in cybersecurity. In practice, cyber threats evolve continuously, making the discovery of novel malicious traffic a critical necessity as well as the identification of known classes. Recent studies have advanced this goal with deep models, but they often rely on task-specific architectures that limit transferability and require per-dataset tuning. In this paper we introduce MalRAG, the first LLM driven retrieval-augmented framework for open-set malicious traffic identification. MalRAG freezes the LLM and operates via comprehensive traffic knowledge construction, adaptive retrieval, and prompt engineering. Concretely, we construct a multi-view traffic database by mining prior malicious traffic from content, structural, and temporal perspectives. Furthermore, we introduce a Coverage-Enhanced Retrieval Algorithm that queries across these views to assemble the most probable candidates, thereby improving the inclusion of correct evidence. We then employ Traffic-Aware Adaptive Pruning to select a variable subset of these candidates based on traffic-aware similarity scores, suppressing incorrect matches and yielding reliable retrieved evidence. Moreover, we develop a suite of guidance prompts where task instruction, evidence referencing, and decision guidance are integrated with the retrieved evidence to improve LLM performance. Across diverse real-world datasets and settings, MalRAG delivers state-of-the-art results in both fine-grained identification of known classes and novel malicious traffic discovery. Ablation and deep-dive analyses further show that MalRAG effective leverages LLM capabilities yet achieves open-set malicious traffic identification without relying on a specific LLM.

Problem

• Identifies novel malicious traffic in cybersecurity
• Overcomes task-specific model limitations in traffic analysis
• Uses retrieval-augmented LLM for open-set threat detection

Innovation

• Retrieval-augmented LLM framework for malicious traffic identification
• Multi-view traffic database construction from content and structure
• Coverage-enhanced retrieval with traffic-aware adaptive pruning