Multi-Agent Framework for Threat Mitigation and Resilience in AI-Based Systems

TL;DR

This work proposes a lifecycle-aware, multi-agent threat framework for AI-based systems, integrating MITRE ATLAS, ATT&CK, and the AI Incident Database with real-world vulnerabilities from ML repositories. A five-agent retrieval-augmented system builds an ontology-driven threat graph that links TTPs, vulnerabilities, and ML lifecycle stages, enabling scalable, graph-based severity estimation and real-time monitoring. Key contributions include a state-of-the-art model analysis revealing high vulnerability in foundation and multimodal models, a cross-layer vulnerability taxonomy spanning data, software, and infrastructure surfaces, and an automated CVE-to-ML-tool mapping that exposes supply-chain risks. The framework yields actionable mitigations across data, software, storage, system, and network layers, validated by external correlations to CVSS and incident costs and supported by operational SOC testing. Overall, adaptive, ML-specific security practices—coupling dependency hygiene, threat intelligence, and monitoring—are essential to mitigate evolving supply-chain and inference-time risks in modern ML ecosystems.

Abstract

Machine learning (ML) underpins foundation models in finance, healthcare, and critical infrastructure, making them targets for data poisoning, model extraction, prompt injection, automated jailbreaking, and preference-guided black-box attacks that exploit model comparisons. Larger models can be more vulnerable to introspection-driven jailbreaks and cross-modal manipulation. Traditional cybersecurity lacks ML-specific threat modeling for foundation, multimodal, and RAG systems. Objective: Characterize ML security risks by identifying dominant TTPs, vulnerabilities, and targeted lifecycle stages. Methods: We extract 93 threats from MITRE ATLAS (26), AI Incident Database (12), and literature (55), and analyze 854 GitHub/Python repositories. A multi-agent RAG system (ChatGPT-4o, temp 0.4) mines 300+ articles to build an ontology-driven threat graph linking TTPs, vulnerabilities, and stages. Results: We identify unreported threats including commercial LLM API model stealing, parameter memorization leakage, and preference-guided text-only jailbreaks. Dominant TTPs include MASTERKEY-style jailbreaking, federated poisoning, diffusion backdoors, and preference optimization leakage, mainly impacting pre-training and inference. Graph analysis reveals dense vulnerability clusters in libraries with poor patch propagation. Conclusion: Adaptive, ML-specific security frameworks, combining dependency hygiene, threat intelligence, and monitoring, are essential to mitigate supply-chain and inference risks across the ML lifecycle.

Figures (18)

• Figure 1: Assets and vulnerabilities across ML/AI system layers. Mapping five infrastructure layers—data, software, storage, system, and network, to their representative assets and prevalent vulnerabilities, drawing from CWE Top 25, OWASP Top 10, NVD, and CVE taxonomies.
• Figure 2: Study methodology using agentic workflow to address RQ1–RQ3.(a) Agentic–RAG pipeline that ingests literature and incident data, retrieves and reranks evidence, and maps TTPs to ML lifecycle phases, model types, and attack scenarios, integrating results into a knowledge graph. (b) Agent workflow that mines GitHub/PyPI repositories for CVE IDs, CPEs and tool names, then maps vulnerabilities to specific ML tools for visualization and analysis.
• Figure 3: Multi-agent LLM with RAG. A coordinated set of AI agents executes four roles: (a)Query & Search: retrieving relevant literature and threat intelligence from Academic Search Engines/Databases; (b)RAG: extracting and ranking relevant TTPs, vulnerabilities, and ML lifecycle stages; (c)Graph Representation: encoding extracted relationships into a heterogeneous knowledge graph; (d)Codebase & MITRE Mapping: Linking GitHub/PyPI CVEs to ATT&CK/ATLAS.
• Figure 4: Global SHAP attribution for the sentences in Listing \ref{['lst:cia-sentences']}. Each row shows the CIA probabilities (left) and the ten tokens with the largest Shapley values for the dominant class (right): green bars push the score up, red bars pull it down.
• Figure 5: LIME explanation for Sentence 1 in Listing \ref{['lst:cia-sentences']}. Left: class probabilities. Centre: token weights— orange = positive, teal = negative. Right: heat-map overlay on the original text.