Ml Security

---
id: ml-security
version: "2.0.0"
title: "ML Model Security"
description: "The model and data artifacts: checkpoint formats that execute code on load, provenance for a model you did not train, training-data poisoning and the ingestion controls that bound it, PII that survives into weights, and notebooks that commit their own output. Use when loading a model from disk, a Hub, or object storage, ingesting user content for training or fine-tuning, or writing training, evaluation, or notebook code."
category: prevention
severity: critical
applies_to:
  - "when generating code that loads an ML model from disk / a Hub / object storage"
  - "when generating data pipelines that ingest user or third-party content for training"
  - "when fine-tuning on data that contains personal or regulated information"
  - "when generating ML notebooks or training / evaluation scripts"
languages: ["python", "jupyter"]
token_budget:
  minimal: 1200
  compact: 1700
  full: 2050
related_skills: ["deserialization-security", "llm-app-security", "supply-chain-security", "secret-detection"]
last_updated: "2026-08-14"
sources:
  - "NIST AI 100-2 (Adversarial Machine Learning)"
  - "MITRE ATLAS (Adversarial Threat Landscape for AI Systems)"
  - "CWE-502 — Deserialization of Untrusted Data"
---

# ML Model Security

## Rules (for AI agents)

### ALWAYS
- Prefer a **tensor-only format** — safetensors — for anything you load. It stores
  arrays and metadata and has no mechanism for executing code, which is the property
  that matters. A `.pt`, `.pkl`, `.bin`, `.ckpt` or `joblib` artifact is a pickle:
  loading it runs whatever the author put in it, before you see a single weight.
- Pass `weights_only=True` to `torch.load` **explicitly**, every time. The default has
  changed across PyTorch releases, so code that relies on it behaves differently
  depending on the pin — and the flag narrows the surface without removing it, so it
  is a mitigation rather than a reason to load an untrusted file.
  `deserialization-security` owns the pickle boundary itself and the per-language
  detail; `references/artifacts-and-provenance.md` has the format-by-format table.
- Establish provenance before loading: the artifact's **hash matches one you recorded**
  from a source you decided to trust, at a revision you pinned. A hash the artifact
  came with proves only that the file is the file. `supply-chain-security` owns what
  makes a source trustworthy, and its rule holds here — a checksum served next to the
  download is not authentication, and an author's name is not a signature.
- Pin the Hub revision to an **immutable commit**, not a branch or tag. `main` moves,
  and `from_pretrained(..., revision="<sha>")` is the difference between a model you
  reviewed and whatever is there today. Where the loader offers `trust_remote_code`,
  understand that enabling it executes repository Python at load time, which is the
  same exposure as the pickle it was meant to avoid.
- Bound what user-contributed data can do to a model. Attribute each training record
  to its source, keep contributed data in a separate pool from curated data, gate
  promotion between them on review, and cap how much any single contributor can
  influence one training run. Poisoning does not need many records — it needs enough
  of them near one behaviour.
- Treat feedback signals as untrusted training input. Thumbs-up/down, ratings,
  corrections and RLHF preferences are user-controlled writes into the next model, and
  they usually arrive through an endpoint nobody thought of as a training pipeline.
- Record model and dataset versions together with the code that produced them, so a
  suspected poisoning can be scoped to a training run and rolled back to a known one.
  Traceability is not prevention; it is what makes the incident finite.
- Scrub personal data at **ingestion**, not only at storage. A model trained on
  personal data memorizes some of it, and no downstream deletion request reaches the
  weights — deleting the row leaves the trained artifact intact.
- Treat notebooks as code that carries its own output. Cell outputs commit credentials,
  data samples and connection strings that were never in the source; clear them before
  committing, and keep checkpoints, datasets and `.ipynb_checkpoints` out of the repo.
- Consult `llm-app-security` for an application that puts a model in a prompt loop —
  injection, tool authorization, what the output may reach — and `api-security` for an
  inference endpoint's own authentication and rate limits. This skill stops at the
  artifact and the data that made it.

### NEVER
- Load a pickle-backed artifact from a source you have not established provenance for
  — `pickle.loads`, `joblib.load`, `dill.loads`, `torch.load`, a `.pt` pulled at
  runtime from a URL. These reconstruct arbitrary objects by design; loading **is**
  execution.
- Convert, inspect, or "just check" a suspect artifact by loading it. Conversion to
  safetensors is a load. Examine it as bytes, in a sandbox with no network and no
  credentials, or not at all.
- Retain training examples containing personal data with no retention window and no
  deletion path — and do not treat the trained model as out of scope for that decision.
- Hard-code a model-provider API key in a notebook, a cell output, or a repo file.
  `secret-detection` owns the patterns and the placeholders.
- Commit generated or synthetic data without labelling it as such. Unlabelled model
  output in a training set silently becomes ground truth for the next run.

### KNOWN FALSE POSITIVES
- A **training checkpoint** that must carry optimizer state, scheduler state, RNG
  state or epoch metadata cannot be expressed in safetensors alone, which stores
  tensors. A pickle-backed checkpoint written and read by your own pipeline, inside
  your own storage, is that format being used for its purpose — the rule is about
  artifacts that arrive from elsewhere.
- A first-party model loaded from storage your deployment controls needs no Hub
  provenance check; the control is that the storage is yours and the path is not
  caller-supplied.
- Research and red-team pipelines deliberately handle poisoned datasets and adversarial
  checkpoints. They belong in an isolated environment with no production credentials,
  and the artifacts there are not findings.
- A public model card, config, tokenizer or `.json` fetched from a Hub is not a code
  artifact. The provenance rule is about what gets deserialized and what gets executed.

## Context (for humans)

Two very different risks live in this skill, and conflating them is the usual mistake.

The first is immediate and unglamorous: a model file is a program. `torch.load` on a
downloaded checkpoint is `exec` on a stranger's code, running as whoever runs training
— typically with cloud credentials, dataset access and a GPU. This has nothing to do
with machine learning; it is deserialization wearing a `.pt` extension, which is why
the boundary belongs to `deserialization-security` and only the artifact question
belongs here.

The second is slow: what the training data does to the model. Poisoning is hard to
detect after the fact because the artifact is valid, the metrics look normal, and the
behaviour is wrong only for inputs the attacker chose. There is no scan for it, so the
controls are all upstream — where the data came from, who could write it, how much any
one writer could contribute, and whether you can say which run produced which model.

Model **serving** is deliberately not covered here. Inference-endpoint authentication
and rate limits belong to `api-security`; an application that puts a language model in
a prompt loop belongs to `llm-app-security`. This skill stops at the artifact and the
data that made it.

## References

- `references/verifying-findings.md` — confirm or refute a finding, then lock it
- `references/artifacts-and-provenance.md` — format-by-format execution risk
  (safetensors, pickle family, Keras, ONNX, GGUF, SavedModel), how to check what your
  pinned PyTorch does by default, Hub loading with pinned revisions, and sandbox
  settings for inspecting an untrusted artifact
- [NIST AI 100-2](https://nvlpubs.nist.gov/nistpubs/ai/NIST.AI.100-2e2023.pdf).
- [MITRE ATLAS](https://atlas.mitre.org/).
- [CWE-502](https://cwe.mitre.org/data/definitions/502.html) — Deserialization of Untrusted Data.