
Artificial intelligence systems are no longer limited to a single machine learning model. A production AI system can include training datasets, data pipelines, model repositories, cloud infrastructure, APIs, third-party models, and connected tools. Each component can introduce security risks.
A weak API can expose sensitive data. A poisoned document can influence a RAG system. An AI agent with excessive permissions can turn a successful prompt injection into an unauthorized action.
The ML pipeline should therefore be treated as a connected security environment. It can include data sources, data processing, training, model development, model storage, deployment infrastructure, APIs and applications, third-party AI services, and monitoring and feedback loops.
A weakness at one stage can also affect another. For example, compromised training data may influence the model deployed months later. Similarly, a vulnerable model may become more dangerous when connected to sensitive enterprise systems.
So, how do you threat model an AI system? The practical approach is to map the entire AI environment, identify where trust changes, understand how an attacker could move through the system, prioritize the resulting risks, and connect each important threat to a security control.
AI system threat modeling is the process of identifying how an AI system could be attacked, misused, manipulated, or compromised and determining what controls are needed to reduce those risks.
Traditional application threat modeling usually focuses on application components, users, data flows, trust boundaries, and common software threats. Those principles still apply to AI. However, AI systems introduce additional components and behaviors that need to be considered.
For example, an AI threat model may need to account for:
The goal is to understand how an attacker could reach an asset, manipulate an AI component, and affect the wider business system.
Vulnerability scanning primarily looks for known technical weaknesses in systems, applications, infrastructure, or dependencies. It can identify issues such as outdated software, insecure configurations, and known vulnerabilities.
Threat modeling happens earlier and at a broader architectural level. It asks questions such as:
Penetration testing attempts to exploit security weaknesses in a real environment. It can validate whether an identified vulnerability is actually exploitable and determine its impact. Threat modeling helps determine what should be tested in the first place.
For example, if threat modeling identifies a risk that an AI agent could access a sensitive database after a prompt injection, security teams can design targeted tests around prompt manipulation, tool authorization, and database access controls.
AI red teaming takes the process further by simulating adversarial behavior against the AI system. It can test prompt injection, model manipulation, data leakage, unsafe tool use, model extraction, and other AI-specific attack scenarios.
In simple terms:
Threat modeling should therefore happen before testing and continue as the AI system changes. It provides the context needed to make security testing more focused and useful.
Threat modeling only the deployed model leaves important parts of the attack surface unexplored. An attacker may never need to compromise the model directly. They may target the data pipeline, API, cloud environment, retrieval layer, or an integration with excessive privileges.
Security risks can enter an AI system before a model reaches production. Training datasets may contain malicious or untrusted content. Data sources may be compromised. An attacker may attempt to poison data to influence future model behavior.
Poor preprocessing can also introduce unexpected weaknesses or expose sensitive information. Common pre-deployment risks include:
NIST's adversarial machine learning taxonomy includes poisoning and privacy attacks among the major attack categories that organizations should consider when assessing AI systems.
The model itself also becomes a security asset. Attackers may attempt to steal model behavior through repeated queries, introduce backdoors, or craft inputs that cause unwanted behavior.
Depending on the system, successful model attacks can affect intellectual property, decision integrity, privacy, or downstream applications. Relevant threats include:
A model should therefore be treated as an important component of the security architecture, not as an isolated algorithm.
Once an AI model is deployed, it is often exposed through APIs, web applications, internal services, or cloud infrastructure. This creates familiar application and infrastructure risks alongside AI-specific ones.
Examples include:
An attacker who cannot compromise the model directly may instead target the interface surrounding it.
AI systems can also create attack paths that only appear during runtime. A user may submit a prompt containing malicious instructions. A retrieved document may contain an indirect prompt injection. An AI agent may interpret malicious content and use a connected tool.
A poorly designed RAG system may retrieve information that the requesting user should not be allowed to access. Runtime threats can include:
The central lesson is simple: AI security must follow the system, not just the model.
A practical AI system threat modeling process can be broken into eight steps. The exact workflow may vary based on the architecture and risk level, but the following approach provides a strong starting point.
Start by defining exactly what you are threat modeling. Avoid beginning with only the model name. An AI system is usually much larger than the model itself. Document the:
For example, an enterprise customer-service assistant may include a web application, authentication service, LLM API, RAG pipeline, vector database, internal knowledge base, monitoring platform, and ticketing system.
All of these components may affect the system's security. At this stage, also identify what is inside and outside the threat model. Clearly document assumptions and exclusions. If a third-party model is treated as a trusted service, record that assumption so it can be reviewed later.
Next, create a clear view of how data and decisions move through the system. A simplified AI pipeline may look like:
Data Sources → Data Ingestion → Data Processing → Training → Model Storage → Deployment → API/Application → User → Monitoring → Feedback/Retraining
Not every AI system will follow this exact sequence. Some systems use pre-trained models instead of internal training. Others add fine-tuning, RAG, agent orchestration, or external AI APIs. For every stage, document:
This mapping helps security teams see where an attacker could interfere with the pipeline.
For example, if documents move from an external repository into a RAG knowledge base, the threat model should consider how those documents are validated, who can modify them, and whether malicious content could influence model behavior.
After mapping the system, identify the assets that need protection. An asset is not limited to sensitive customer information. In an AI environment, attackers may target data, models, credentials, infrastructure, instructions, or system behavior.
Important assets may include:
For example, a public model endpoint and a proprietary model weight file may both be important, but the impact of compromise can be very different. Similarly, an internal vector database containing customer records may require stronger access controls than a public knowledge base.
A trust boundary is a point where the level of trust, control, or authorization changes between components. These boundaries are especially important in AI systems because data often moves between users, applications, models, external services, and enterprise tools.
Potential entry points include:
Consider boundaries such as:
Users ↔ Applications
A user may be authenticated, but that does not mean every prompt or uploaded file should be trusted.
Applications ↔ Models
The application may need to control which users can access the model and what functions they can invoke.
Models ↔ Tools
This is especially important for AI agents. A model should not automatically receive unrestricted access to databases, email systems, cloud resources, or other sensitive tools.
Internal ↔ External Data
External documents and APIs should not automatically receive the same level of trust as internal sources.
Development ↔ Production
Training environments, model repositories, CI/CD systems, and production inference environments should have carefully controlled access paths.
.webp)
Once the architecture is mapped, identify threats that are relevant to each component. Traditional security threats should still be included. However, AI systems require additional consideration because attackers can target data, model behavior, prompts, retrieval context, and AI-enabled actions.
Relevant threats include:
MITRE ATLAS is useful at this stage because it organizes adversarial behavior against AI systems into tactics and techniques. Its current knowledge base covers predictive AI, generative AI, and agentic AI environments.
For example,
The goal is not to force every possible AI threat into the model. Threats should be selected based on the architecture, data, users, integrations, and business purpose of the system.
A threat list alone is not enough. The next step is to understand how an attacker could actually move through the system. For each high-value threat, ask:
Consider an enterprise RAG assistant connected to internal systems. A possible attack path could be:
Malicious document → RAG ingestion → Retrieval → Prompt context → Model response → Privileged tool → Sensitive system
The initial malicious document may not look like a conventional software vulnerability. The problem emerges when the document reaches the model as trusted context and influences an AI agent with access to a sensitive tool.
This illustrates why isolated component testing can miss chained AI attacks. The threat model should therefore consider the full attack path, not just individual weaknesses.
Not every identified threat requires the same level of attention. Prioritize risks using factors such as:
A simple starting point is:
Risk = Likelihood × Impact
Organizations can then expand the model when more detailed risk analysis is required. For example, a prompt injection against a public chatbot may have limited impact if the system has no sensitive data or connected tools.
The same type of attack could become critical if the AI agent can access customer records or execute financial transactions. High-priority threats often include those that could cause:
Risk prioritization keeps the threat model actionable. It helps teams focus security resources on the attack paths that matter most to the business.
The final step is to connect important threats with practical security controls. A threat model becomes much more valuable when it tells teams what they need to change.
For example, prompt filtering may reduce some prompt injection attempts, but it should not replace authorization controls around sensitive tools. Similarly, model access controls should be supported by API security, network controls, monitoring, and rate limiting.
The strongest AI security architecture uses multiple layers so that a failure in one control does not automatically lead to system compromise.
MITRE ATLAS is a knowledge base for understanding adversarial threats against AI-enabled systems. It helps security teams study attacker behavior and connect AI-specific attack techniques with realistic attack scenarios.
The current ATLAS matrix covers predictive AI, generative AI, and agentic AI environments. MITRE ATLAS complements traditional threat modeling because it adds an AI-focused view of adversary behavior. Teams can use ATLAS to:
Suppose an organization operates an AI agent that can retrieve internal information and invoke enterprise tools. The team can then use MITRE ATLAS to research relevant adversarial techniques and determine which attack behaviors should be tested.
This can lead to targeted security scenarios such as prompt manipulation, unauthorized tool use, data collection, or data exfiltration.
MITRE ATT&CK is a broader knowledge base for adversary tactics and techniques across enterprise environments. MITRE ATLAS applies a similar threat-informed approach specifically to AI systems. The two can complement each other.
An AI platform may contain conventional enterprise infrastructure and AI-specific components. ATT&CK can help model threats against the surrounding enterprise environment, while ATLAS can help analyze attacks targeting the AI components and their unique behavior.
This makes MITRE ATLAS particularly useful when an AI threat model needs to connect traditional cybersecurity risks with AI-specific attack techniques.
The NIST AI Risk Management Framework (AI RMF) and threat modeling address related but different needs. NIST AI RMF provides a broader structure for managing AI risks across the lifecycle. It is intended for organizations that design, develop, deploy, or use AI systems and is designed to be flexible across industries and use cases.
Threat modeling adds more technical detail by translating those broader risks into system-specific assets, threats, attack paths, and controls.
The AI RMF Core is organized around four functions:
NIST emphasizes that these functions are not simply a linear checklist. Risk management should be continuous and applied throughout the AI system lifecycle.
Threat modeling can make these risk-management activities more concrete. It can translate high-level AI risks into:
For example, an organization may identify sensitive information disclosure as an AI risk under its broader risk-management process. Threat modeling can then determine where that risk could occur.
It might identify a RAG database, model prompt, API, application session, or logging system as a potential exposure point. Security teams can then define controls and testing requirements for those specific components.
Do not limit the threat model to the production model or inference endpoint. Map the full AI lifecycle, including data collection, ingestion, processing, training, fine-tuning, model storage, deployment, inference, application integrations, monitoring, feedback, and retraining.
Each stage can introduce different security risks. For example, attackers may target training data before a model is deployed, compromise a model artifact in a repository, or exploit an exposed API after deployment.
A RAG application may also introduce risks at the document ingestion and retrieval stages that would not appear in a traditional ML pipeline. For each lifecycle stage, document what data enters the system, what processing occurs, which assets are created, who has access, and where the output goes next.
This makes it easier to identify threats that cross multiple stages instead of treating each component as an isolated system.
A detailed data-flow map is one of the most useful outputs of an AI threat modeling exercise. It should show how data moves between users, applications, models, databases, APIs, external services, and supporting infrastructure.
Pay particular attention to trust boundaries. A trust boundary can exist when data moves from an external source into an internal system, from a user-controlled application into a model, or from an AI model into a privileged enterprise tool.
For example, an employee may be authorized to use an AI assistant, but that does not mean the assistant should be allowed to retrieve every internal document. Similarly, a model may be allowed to generate an API request without being authorized to execute every possible action through that API.
An AI application may still contain web interfaces, APIs, databases, cloud resources, authentication mechanisms, CI/CD pipelines, containers, open-source dependencies, and identity systems.
These components remain exposed to conventional threats such as broken access control, injection, insecure configuration, credential compromise, and vulnerable dependencies.
At the same time, AI introduces threats that require additional consideration. These can include prompt injection, data poisoning, adversarial inputs, RAG poisoning, model manipulation, sensitive information disclosure, and unauthorized AI tool use.
A complete threat model should therefore combine both categories. For example, an AI chatbot may be vulnerable to prompt injection at the AI layer while its underlying API has a conventional authorization flaw. An attacker could potentially chain both weaknesses to access data that the AI application was never intended to expose.
MITRE ATLAS can help security teams move from generic threat descriptions to realistic AI attack scenarios. ATLAS is a knowledge base focused on adversarial tactics and techniques against AI-enabled systems and includes coverage for areas such as machine learning, generative AI, and agentic AI.
During threat modeling, use ATLAS to identify attacker behaviors that are relevant to the architecture. For example, a system using a public inference endpoint may require analysis of model extraction or inference-related attacks.
A RAG application may require scenarios involving malicious content and retrieval manipulation. An AI agent with access to enterprise tools may require analysis of techniques involving prompt manipulation, data access, or unauthorized actions.
.webp)
Threat modeling provides technical detail, but it should fit into a broader AI risk management process. The NIST AI Risk Management Framework (AI RMF) provides a useful structure for doing this through its four core functions: Govern, Map, Measure, and Manage.
The Govern function helps establish organizational policies, responsibilities, and accountability for AI risks. Map focuses on understanding the intended context, stakeholders, system characteristics, and potential risks.
Measure involves assessing and analyzing identified risks using appropriate methods. Manage focuses on prioritizing and addressing those risks.
AI system threat modeling fits particularly well into the technical work behind these activities. It can help teams identify assets, map data flows, document threats, analyze attack paths, define security controls, and establish testing requirements.
Prioritize threats based on factors such as exploitability, likelihood, data sensitivity, system exposure, model criticality, privilege level, and potential business impact.
Consider what could happen if an attacker successfully exploited the threat, not just how technically interesting the vulnerability appears.
For example, prompt injection against a public chatbot with no sensitive integrations may have limited impact. The same attack against an AI agent that can access customer records, modify business data, or initiate financial actions could be significantly more serious.
Consider prioritizing attack paths that could result in:
This risk-based approach helps organizations focus on attack paths with the greatest potential impact rather than simply counting vulnerabilities.
Controls may include:
For example, a prompt injection defense should not be the only protection around an AI agent that can access a sensitive database.
The database should also enforce authorization independently. This way, a successful prompt manipulation does not automatically result in unauthorized data access.
Third-party AI components can become part of an organization's effective attack surface. This includes external LLM APIs, hosted models, open-source models, vector databases, plugins, libraries, cloud AI services, and AI development platforms.
Threat modeling should identify what information is sent to each external service, what permissions it receives, where data is processed, and what security controls are available. For external AI services, review areas such as:
Third-party risk is particularly important when an external model processes sensitive information or when an AI service is given access to internal systems.
ioSENTRIX provides AI/ML and LLM security testing that covers areas such as data pipelines, model training processes, deployment environments, and API integrations. Our published capabilities also include AI-specific testing for adversarial attacks, data poisoning, API exploitation, prompt injection, data leakage, and model manipulation.
For organizations that need deeper adversarial validation, AI red teaming can evaluate how an AI system responds to intentional manipulation and misuse. This can include prompt injection, model abuse, data extraction, API exploitation, and risks created by AI-driven automation.
ioSENTRIX has also published an AI/ML threat modeling case study involving a Fortune 500 financial services organization. The assessment analyzed more than 30 components across AWS and Azure and identified security threats across RAG pipelines, LLM orchestration, and ML inference workflows.
Start by defining the AI system and its scope. Then map the architecture and data flows, identify critical assets and trust boundaries, identify AI-specific threats, analyze attack paths, prioritize risks, and map threats to security controls. High-risk scenarios should then be validated through penetration testing or AI red teaming.
AI systems have attack surfaces that extend beyond traditional software. Attackers can target training data, model files, prompts, APIs, retrieval systems, connected tools, and deployment infrastructure. Threat modeling helps organizations identify these risks before they become exploitable weaknesses.
The threats depend on the AI architecture, but common examples include data poisoning, model tampering, model extraction, adversarial manipulation, prompt injection, RAG poisoning, sensitive data exposure, supply-chain compromise, credential theft, API abuse, and unauthorized tool use.
Yes. AI threat modeling should cover the relevant lifecycle stages, from data collection and processing through model development, deployment, inference, integrations, monitoring, and retraining. The exact scope depends on the architecture, but limiting the assessment to the production model can leave significant attack paths undiscovered.
AI threat modeling identifies potential threats and attack paths based on the system architecture. AI red teaming actively simulates adversarial behavior to determine whether those attack paths can be exploited. Threat modeling helps define what to test, while red teaming helps validate how the system performs under attack.