Teams that train, fine-tune or serve their own AI models often threat model the application around the model and leave the model pipeline itself, its training data, weights, serving API and upstream artefacts, without a recorded owner, threat list or control. This template gives you that record and the policy that requires it. It covers the model lifecycle from data collection to retirement: a pipeline asset register, a threat register mapped to MITRE ATLAS technique IDs, and a control register mapped to ATLAS mitigations, with owners and tests. It is written for ML engineers, platform owners, security architects and the risk leads who approve model releases. AI agents are covered by the AI Agent Threat Model Template and retrieval-augmented generation by the RAG Threat Modelling Guide; the general method is in the Threat Model Template.
How to use this template
- Replace text in square brackets, such as [Organisation], [ML platform team] and [Model risk committee], with your own terms.
- "Must" is mandatory; "should" is recommended practice where a documented exception may apply.
- Sections 4.3 to 4.5 hold the three registers. Copy them into a spreadsheet or threat modelling tool and keep one record per model or model family.
- ATLAS IDs and names in this template were checked against the MITRE ATLAS data release of October 2026. ATLAS is updated regularly, so confirm IDs on the ATLAS site when you review the record.
- Pair this template with the AI Red Team Test Plan Template to test the controls, and the AI Model Theft Response Playbook for response.
1. Purpose
This document sets the requirements for threat modelling AI models at [Organisation]. It ensures that, before a model is released to production or shared outside its development environment, [Organisation] has recorded the assets in its pipeline, the threats to each lifecycle stage, the controls that reduce them, who owns those controls and how they are tested. It supports Principle 3 of the UK Code of Practice for the Cyber Security of AI, which expects threat modelling to cover AI-specific attacks such as data poisoning, model inversion and membership inference, and to be repeated when new configurations are introduced.
2. Scope
This document applies to:
- Models trained from scratch by [Organisation], including predictive and generative models.
- Pre-trained or open-weight models that [Organisation] fine-tunes, distils, quantises or hosts.
- Training, evaluation and fine-tuning datasets, labelling processes and feature stores.
- Model registries, artefact stores, training and serving infrastructure, and inference APIs.
- Third-party models consumed only through a supplier API, for the threats [Organisation] controls (keys, inputs, outputs, cost).
Agent tools, memory and delegated credentials, and RAG retrieval stores, are threat modelled with their own templates and referenced from this record.
3. Roles and responsibilities
| Role | Responsibility |
|---|---|
| [Model owner] | Requests the threat model, states intended use and misuse, accepts residual risk within their authority, keeps the record current |
| [ML engineering lead] | Supplies pipeline architecture, data sources, artefact formats and serving design; implements engineering controls |
| [Data custodian] | Confirms dataset provenance, licensing, classification and personal data content |
| [Security architect] | Facilitates, completes the threat register, rates risks, proposes controls, signs the record as complete |
| [Platform or NHI team] | Validates the service accounts, API keys and workload identities used by pipelines and serving |
| [Test lead] | Turns each control into a test, including adversarial and robustness tests |
| [Model risk committee] | Approves records for high-impact models, accepts high residual risk, approves exceptions |
4. Policy statements
4.1 When a threat model is required
- Every in-scope model must have an approved threat model before production release, external sharing of its weights, or use of production data in training.
- The record must be updated, and re-approved where risk increases, when any of these change: training or fine-tuning data source; base model or supplier; artefact format or serialisation; serving exposure (internal, partner, public); output detail returned by the API; hosting environment; intended use.
- High-impact models must be reviewed at least [every six months], others at least [annually].
- A security incident, red team finding or ATLAS update affecting the model must trigger a review within [10 working days].
4.2 Minimum content
- A one-sentence statement of intended use and a list of outcomes the model must never produce or enable.
- A pipeline diagram covering data sources, labelling, training, evaluation, registry, packaging, serving and monitoring, with trust boundaries and every external input marked.
- The asset register (4.3), threat register (4.4) and control register (4.5).
- A lifecycle coverage statement confirming each stage in 4.4 was considered, with a reason for any marked not applicable.
- Residual risk, the person who accepted it, and the date.
4.3 Pipeline asset register
| Field | What to record | Example |
|---|---|---|
| Asset ID | Unique reference | M-03 |
| Asset type | Dataset, label set, base model, fine-tuned weights, checkpoint, registry, training cluster, inference endpoint, system prompt, pipeline credential, AI-BOM | Fine-tuned weights |
| Lifecycle stage | Data, training, evaluation, packaging, serving, retirement | Packaging |
| Source and provenance | Who produced it, where from, hash or signature | Internal training run 2026-10-12; SHA-256 recorded in registry |
| Format | Serialisation and whether it can execute code on load | safetensors (no code on load) |
| Who can read or write | Identities, roles, keys | Training service account (write); serving workload identity (read) |
| Confidentiality and integrity impact | Worst result if stolen or altered | Loss of competitive IP; altered outputs to customers |
| Owner | Named person or team | [ML platform team] |
Record pipeline credentials as assets in their own right. A registry write token or a long-lived cloud key on a training cluster can alter or exfiltrate every model it touches. See the Model Supply Chain Guide for the upstream side.
4.4 Threat register
Record each threat with its entry point, affected assets, likelihood, impact and ATLAS mapping. The rows below are the minimum starting set by lifecycle stage; add threats specific to the model.
| Stage | Threat | MITRE ATLAS | Typical assets affected |
|---|---|---|---|
| Data | Poisoned public or scraped data enters training | AML.T0020 Training Data Poisoning; AML.T0115 Publish Poisoned AI Artifacts | Datasets, all downstream weights |
| Data | Gradual corruption of a maintained dataset or labels | AML.T0059 Erode Dataset Integrity | Datasets, label sets |
| Data | Hidden trigger planted so the model misbehaves on specific inputs | AML.T0043.004 Insert Backdoor Trigger; AML.T0018.000 Poison AI Model | Weights, outputs |
| Supply chain | Compromised base model, dataset, library or container image | AML.T0010 Supply Chain Compromise (.001 Software, .002 Data Repository, .003 Model Repository, .004 Container Registry); AML.T0115 Publish Poisoned AI Artifacts | Base model, training environment |
| Supply chain | Unsafe serialised artefact runs code when loaded | AML.T0011.000 Unsafe AI Artifacts | Training and serving hosts, their credentials |
| Training and registry | Weights or architecture altered in the registry or pipeline | AML.T0018.001 Modify AI Model Architecture; AML.T0076 Corrupt AI Model | Checkpoints, registry |
| Training and registry | Weights, datasets or checkpoints stolen through infrastructure access | AML.T0035 AI Artifact Collection; AML.T0025 Exfiltration via Cyber Means; AML.T0012 Valid Accounts | Weights, datasets, pipeline credentials |
| Serving | Model functionality replicated through repeated API queries | AML.T0005.001 Replicate Behavior from Model Outputs; AML.T0024 Exfiltration via AI Inference API | Weights (as IP), inference API |
| Serving | Training data membership inferred or inputs reconstructed from outputs | AML.T0137.000 Infer Membership; AML.T0137.002 Invert AI Model | Personal or confidential training data |
| Serving | Crafted inputs cause misclassification or bypass a model-based control | AML.T0015 Evade AI Model; AML.T0043 Craft Adversarial Data | Decisions based on outputs |
| Serving (LLM) | Jailbreaks, prompt injection, system prompt extraction or training data leakage | AML.T0054 LLM Jailbreak; AML.T0051 LLM Prompt Injection; AML.T0069.002 System Prompt (Discover LLM System Information); AML.T0137.003 Extract Memorized Data | System prompt, training data, outputs |
| Serving | Flooding or costly queries degrade service or run up compute bills | AML.T0029 Denial of AI Service; AML.T0034 Cost Harvesting | Inference capacity, budget |
| Serving | Stolen API keys or accounts used to query the model | AML.T0040 AI Model Inference API Access; AML.T0012 Valid Accounts | Inference API, model service keys |
| Operation | Inputs crafted over time to degrade model performance | AML.T0031 Erode AI Model Integrity | Model accuracy, retraining data |
Cross-references: AI Data Poisoning and Backdoors Guide, Model Extraction and Distillation Defence Guide, Membership Inference Attacks. NIST AI 100-2 E2025 gives the terminology for these attack classes.
4.5 Control register
Each threat must have at least one control with an owner, a test and evidence. Map controls to ATLAS mitigations so coverage can be reviewed across models.
| Control | ATLAS mitigation | Threats addressed | Test or evidence |
|---|---|---|---|
| Record dataset sources, licences, dates and hashes; prefer documented sources | AML.M0025 Maintain AI Dataset Provenance | T0020, T0115, T0059 | Provenance record per dataset version |
| Filter and review training data for anomalies and outliers before use | AML.M0007 Sanitize Training Data | T0020, T0043.004 | Data validation report |
| Verify hashes or signatures of every model and dataset before use; sign released models | AML.M0014 Verify AI Artifacts; AML.M0013 Code Signing | T0010, T0115, T0076, T0018.001 | Signature verification log at load time |
| Load only non-executable formats such as safetensors; inspect imported models in an isolated zone | AML.M0011 Restrict Library Loading; AML.M0016 Vulnerability Scanning | T0011.000, T0010.003 | Format policy check in pipeline; scan results |
| Maintain an AI bill of materials for each release | AML.M0023 Maintain an AI Bill of Materials | T0010 | AI-BOM stored with the release |
| Store weights in restricted storage with separate encryption keys; least-privilege pipeline identities | AML.M0005 Control Access to AI Models and Data at Rest; AML.M0012 Encrypt Sensitive Information | T0035, T0025, T0012 | Access review; key custody record |
| Authenticate and authorise every inference call; short-lived credentials, no shared keys | AML.M0019 Control Access to AI Models and Data in Production | T0040, T0012 | Endpoint configuration review |
| Rate and spend limits per caller; return only the output detail needed (labels, not full scores) | AML.M0004 Limit AI Service Query Volume and Rate; AML.M0002 Predictive AI Output Obfuscation | T0024, T0005.001, T0137, T0029, T0034 | Load test; API response review |
| Adversarial robustness testing and validation after every change | AML.M0008 Validate AI Model; AML.M0003 Predictive AI Model Hardening | T0015, T0043, T0043.004 | Test report against release criteria |
| Input and output guardrails for generative models | AML.M0020 Generative AI Guardrails | T0054, T0051, T0069.002, T0137.003 | Red team results |
| Log inputs, outputs, errors and configuration changes; alert on extraction patterns and drift | AML.M0024 AI Telemetry Logging and Monitoring | T0024, T0137, T0031, T0029 | Alert rules; sample investigation |
5. Exceptions
- Exceptions must be requested by the [Model owner], state the control not met, the compensating control and an expiry date no later than [six months].
- Exceptions for high-impact models must be approved by the [Model risk committee] and recorded in the risk register.
- An exception never permits loading an unverified artefact in an executable format into production.
6. Compliance and monitoring
- [Security architecture team] must report quarterly: share of in-scope models with a current approved record; records overdue for review; high residual risks accepted; open exceptions.
- The model registry should block promotion to production of any model without a linked approved record.
- Internal audit may sample records against this document.
7. Review and approval
| Item | Value |
|---|---|
| Document owner | [Head of AI security] |
| Approved by | [Model risk committee] |
| Version and date | [1.0, date] |
| Next review | [Annually, or on a significant ATLAS or regulatory change] |
Implementation guidance
- Start with one model: pick a production model with external exposure and complete the asset register first. Missing provenance or unknown artefact formats are common first findings.
- Walk the pipeline left to right: run the workshop stage by stage using 4.4, asking for each asset who can write to it and how you would know it changed.
- Fix artefact integrity early: hashing, signing and non-executable formats close several threats at once. Pickle-based files can run arbitrary code when loaded, and Hugging Face's own documentation advises not to unpickle data from untrusted sources and notes that its pickle scanning is not foolproof. The OpenSSF model signing project (sigstore/model-transparency) signs model files with Sigstore, keys or certificates.
- Treat pipeline and serving credentials as NHIs: give training jobs and endpoints workload identities, scope registry access, and rotate model service keys. Stolen AI keys have been widely abused (see LLMjacking 2024 to 2026).
- Turn controls into tests: hand the control register to the test lead; each row needs a pass condition before release.
- Measure: track coverage of in-scope models, verified-artefact rate at load, extraction alerts investigated, and time from ATLAS update to record review.
| Framework | Relevant part |
|---|---|
| UK Code of Practice for the Cyber Security of AI | Principles 3 (threats), 5 (assets), 7 (supply chain), 8 (document data, models and prompts), 9 (testing), 12 (monitoring), 13 (disposal) |
| NCSC Guidelines for Secure AI System Development | Secure design (threat modelling), development (supply chain), deployment (protect models), operation (logging) |
| Deploying AI Systems Securely (NSA, CISA and partners) | Isolated weight storage, artefact hashing and signatures, inspection of imported models, API authentication, logging |
| NIST SP 800-218A | SSDF community profile for generative AI and dual-use foundation model development |
| NIST AI 100-2 E2025 | Adversarial machine learning taxonomy and terminology |
References: MITRE ATLAS · NIST AI 100-2 E2025 · NIST SP 800-218A · UK Code of Practice for the Cyber Security of AI · NCSC Guidelines for Secure AI System Development · Deploying AI Systems Securely · Hugging Face: Pickle Scanning · OpenSSF: Model Signing
Related NHI Mgmt Group resources: AI Model Security Guide · MITRE ATLAS Guide · AI Supply Chain and AI-BOM Guide · Model Supply Chain Guide