Ml Security

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.
• MITRE ATLAS.
• CWE-502 — Deserialization of Untrusted Data.