Threat Modelling an AI Feature Before You Write Code

Threat modelling is a discipline borrowed from security engineering but essential to AI safety. Before a developer or team writes a single line of code for an AI feature, they should articulate what could go wrong, who could cause it, and what the consequence would be for users, the organisation and the model itself.

A threat model for an AI feature has different structure than a traditional software threat model because the risk surface is different. Traditional software threats are usually about confidentiality (data theft), integrity (unauthorised modification) and availability (denial of service). AI threats include those three, but also add model failure modes: hallucination and fabrication, output bias, prompt injection, data poisoning, model theft, and silent drift.

Start threat modelling by defining the boundary of your system

An AI-powered customer support chatbot has a clear boundary: it reads customer messages, calls a language model, applies guardrails to the output, and returns a response. Everything outside that boundary (the customer database, payment systems, ticket routing) is in scope for data flows but not in scope for the model itself.

Threat Modelling an AI Feature Before You Write Code Concept Diagram
Figure: Conceptual architecture and workflow for Threat Modelling an AI Feature Before You Write Code

Understanding boundaries prevents developers from building guardrails that only protect the model when they also need to protect downstream systems.

Next, list the assets inside and outside the boundary

For an AI-powered loan assessment system, assets include the training data (which contains applicant personal information), the model weights (which represent the bank’s underwriting logic), the live predictions (which affect people’s financial lives), and the audit logs (which defend the bank if a decision is challenged).

Each asset has different value and different threats.

Enumerate the threat actors

Who would want to harm your system? For a loan assessment model, threat actors include applicants who want to obtain credit fraudulently, employees with elevated access who might probe the model to find bias, competitors who want to copy your underwriting logic, and regulators performing stress tests.

Employees represent an insider threat; they have legitimate access but might abuse it. Applicants represent an external threat; they have no legitimate access and will probe for weaknesses.

Against each asset and threat actor pair, ask: what could go wrong? An applicant could submit a prompt-injected message designed to make the model recommend approval regardless of creditworthiness. An employee could request a batch of loan assessments with systematically biased candidate data to see if the model outputs discriminatory decisions.

Video: DEF CON 32 – Threat Modeling in the Age of AI – Adam Shostack

A regulator could request the model weights and ask why certain protected characteristics correlate with approval rates. Each of these is a different threat with different severity.

For each threat, estimate the likelihood and impact

A prompt injection attack on a chat interface is plausible and has high impact if successful: a misguided response damages user trust and creates liability. Silent model drift (where performance degrades over time without human awareness) is also plausible and has high impact: decisions become systematically wrong before anyone notices.

A model theft attack is lower likelihood (requires either source code access or reverse engineering), but if successful, has high impact: your proprietary logic is now your competitor’s advantage.

Map each threat to mitigations

For prompt injection, mitigations include input validation, guardrails that filter suspicious patterns, limiting tool access that the model can invoke, and structured output schemas that constrain what the model can return. For silent drift, mitigations include automated performance monitoring on held-out test data, alerting when metrics degrade by a defined threshold, and automated rollback to a previous model version.

Document the threat model

This is not optional extra work; it is the starting point for security and safety testing, for architecture review, and for evidence that your organisation took reasonable care in deployment. When a regulator asks why you shipped a system that caused harm, a documented threat model that shows you foresaw the risk and implemented mitigations is the difference between “negligence” and “reasonable safeguard”.

Threat modelling is not a one-time activity. Threat models should be revisited when the system changes: when you add new features, integrate with new systems, move to new infrastructure, or discover new attack vectors. Keep the threat model current.

When you discover a new attack or failure mode in production, add it to the threat model and design a mitigation. This iterative approach keeps your defences aligned with evolving threats.