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.