AI Models Without Malware: Pickle vs SafeTensors

An AI model file looks like data, but loading it may be as sensitive as installing a package. A checkpoint may use pickle, a repository may contain custom code, configuration can select a dynamically loaded class, and a parser handling huge tensors can exhaust memory or hit a native-code flaw. The model format is therefore part of supply-chain security, but never the whole story.

In one paragraph: prefer a format without executable semantics, such as SafeTensors, for weights—but never treat the .safetensors extension as a security certificate for the repository. Separate weights from code, pin an exact revision and digest, verify signatures and provenance, scan every file, perform the first load in a disposable environment without secrets or egress, and publish an approved artefact to your own registry.

The torch.load documentation warns never to load data from an untrusted source because the mechanism uses an unpickler. Hugging Face documents SafeTensors as a simple tensor-storage format designed as a safer alternative to pickle while supporting efficient access.

A model is more than a weights file

A typical repository may contain:

• tensor weights in one or many shards;
• an index mapping tensors to shards;
• architecture and generation configuration;
• a tokenizer, vocabulary, merges, and preprocessing models;
• custom Python or C++ code;
• conversion scripts, notebooks, and dependencies;
• LoRA adapters, quantisation settings, and kernel libraries;
• model cards, licences, and provenance metadata.

Safe weights do not neutralise setup.py, a malicious import, a vulnerable tokenizer, or a trust_remote_code decision. Conversely, a code-free repository can contain dimensions that exhaust resources or a model carrying a behavioural backdoor. The threat model must cover the complete package.

Why pickle can execute code

Python pickle stores instructions needed to reconstruct objects. It is not an interchange format for an untrusted sender. The protocol may reference functions and classes present in the environment, so a crafted stream can trigger unwanted behaviour during loading.

The PyTorch ecosystem has historically used .pt, .pth, .bin, and .ckpt. An extension does not unambiguously identify the format. Review how the artefact is actually loaded. Renaming model.pkl to weights.bin does not remove pickle semantics.

Starting with PyTorch 2.6, torch.load() defaults to weights_only=True when no custom pickle_module is supplied. The serialization semantics documentation explains that the restricted unpickler cannot dynamically import arbitrary functions and classes. PyTorch also documents limitations: the mode does not guarantee protection from denial of service and does not eliminate every potential memory-safety issue.

This is useful migration defence, not permission to load random checkpoints on a host with cloud credentials. When an error recommends switching to weights_only=False, do not do so automatically. Establish provenance and perform conversion in isolation first.

What SafeTensors provides

SafeTensors separates a small header containing tensor metadata from a contiguous data area. It is not designed to serialize arbitrary Python objects and has no mechanism for invoking functions while loading. It supports lazy loading and memory mapping, which is useful for large models.

Hugging Face, EleutherAI, and Stability AI commissioned an external review of the library. The SafeTensors security audit summary links the public report. The review found no critical arbitrary-code-execution issue in the assessed version and prompted improvements to the specification, validation, and test suite.

SafeTensors reduces an important risk class, but it does not guarantee:

• that weights came from the claimed author;
• absence of backdoors or harmful model behaviour;
• safety of other repository files;
• absence of flaws in the loader, runtime, driver, or kernel;
• that tensor shapes and types fit resource budgets;
• that the model’s licence and training data permit its use.

The format answers “does this weights file carry pickle’s executable semantics?” rather than “is this model trustworthy?”

ONNX, TorchScript, and executable artefacts

ONNX describes a computation graph and operators. That is a different risk model from pickle, but parsers and runtimes execute a complex graph, often using native code and optional custom operators. “Not pickle” is not equivalent to “safe to run from any source.”

TorchScript, compiled inference engines, CUDA libraries, and custom ops also expand the execution surface. A model may require a separately downloaded plugin. Every component needs an identified origin, version, digest, and update owner.

A practical classification is:

Artefact	Primary risk	Preferred control
pickle/checkpoint	object construction and execution	do not trust; isolated conversion
SafeTensors	parser, resources, malicious weights	limits + provenance + evaluation
ONNX/graph	parser, runtime, custom ops	sandbox + operator allowlist
model code	arbitrary execution	review, pin, build, and sign
tokenizer/parser	native bugs and DoS	limits, updates, and fuzzing
inference container	complete software supply chain	SBOM, scan, signature, least privilege

Threats in model repositories

Typosquatting and publisher impersonation

An attacker publishes a name similar to a popular model, copies its description, and adds a malicious artefact. Verify the organisation, repository history, links from official documentation, and a cryptographic digest. Download count is not evidence of authenticity.

Publisher account or token compromise

A legitimate repository may change after an account takeover. Pinning to main is not reproducible. Use an exact commit or revision and an approved internal mirror. Treat updates like new dependencies.

Malicious conversion

Converting pickle to SafeTensors requires reading the old file first. If conversion happens on a developer laptop, the risk materialises before the safer output exists. Run converters in disposable environments without network access, secrets, writable host mounts, or control-plane sockets.

Remote code and dependencies

Some models require a custom class. Enabling remote code is an explicit decision to execute supplier code. Pin it to a revision, review it, build it internally, and run it with least privilege. AI coding-agent supply-chain security demonstrates how automation can turn repository instructions into execution.

A reference model-import pipeline

1. Quarantine intake

Download through a dedicated service. Do not open the artefact automatically in a notebook. Record URL, namespace, revision, time, declared licence, size, and hash of every file. Block automatic installation scripts.

2. Provenance verification

Verify that the source is linked to the official publisher. Prefer a signed manifest or provenance statement. SHA-256 only identifies an artefact when the expected digest arrives through a separate trusted channel. A hash posted beside a replaced file proves nothing.

3. Format classification

Detect formats from structure rather than extension. Mark pickle, archives, native libraries, custom ops, and code. Policy can automatically advance selected formats while sending the rest for manual review.

4. Static scanning

Scan for malware, secrets, licences, vulnerable dependencies, and dangerous imports. Pickle opcodes can be inspected without execution, but “no known pattern found” is not a safety guarantee. Cap archive file count, expanded size, and path depth.

5. Isolated conversion and smoke test

A disposable worker reads from a read-only input and writes only to an empty output. It receives no keys, Docker socket, developer cache, or production data. Enforce CPU, RAM, GPU, time, process-count, and file-size limits. Destroy the worker after conversion.

6. Behavioural security evaluation

Test canary prompts, threat-relevant backdoors, unexpected connections, loaded libraries, and resource usage. AI red teaming evaluates behaviour that no storage format can guarantee.

7. Publish to an internal registry

Create an immutable signed manifest containing digests of weights, code, configuration, tokenizer, runtime image, and control results. Production should pull only approved versions from the registry, never directly from a public model hub.

Connecting SBOM, SLSA, and signatures

An SBOM describes software components, but a conventional SBOM may not fully capture weights, datasets, adapters, and conversion parameters. Add a model manifest that records relationships among them. Treat tokenizers, custom ops, and the inference container as first-class dependencies.

SLSA and Sigstore-style controls can bind an artefact to an identity and build process. A signature does not mean that the model is good; it allows policy to enforce that production runs exactly what passed the approved pipeline.

In cloud deployments, model-key release can also depend on TEE attestation. Provenance then says where the artefact came from, a signature says who approved it, and attestation says which code is running it now.

Hardening the inference environment

Even an approved model needs containment:

• dedicated workload identity without developer credentials;
• read-only filesystem outside controlled cache locations;
• no default egress, or a narrow destination allowlist;
• seccomp/AppArmor/SELinux and no privileged containers;
• memory, GPU, process, and request-time limits;
• signed minimal images;
• monitoring for new processes, libraries, connections, and file changes;
• isolation between tenants.

Container image security is part of the same chain. The model may be safe while the runtime is vulnerable—or the reverse.

Controlled updates and rollback

Do not automatically track the newest revision of a public repository. Every update repeats intake, diffing, scanning, conversion, evaluation, and signing. Compare behaviour on a fixed suite in addition to files, resource use, and dependency changes.

For rollback, retain an immutable manifest of the previous model and a compatible runtime. Never assume an upstream tag still points to the same commit. Audit records should answer who approved the artefact, when, from which source, using which toolchain, and with which results.

Import checklist

• Do we pin an exact revision and know every file digest?
• Was publisher identity verified through an independent channel?
• Does the pipeline detect content format instead of trusting filenames?
• Do pickle and remote code always require manual approval and a sandbox?
• Does conversion run without network, secrets, or host control?
• Are tokenizers, config, custom ops, and containers scanned along with weights?
• Does production use only an internal signed registry?
• Does the manifest bind model, code, configuration, runtime, and test results?
• Are updates and rollback reproducible?

SafeTensors is a sensible default for weights because it removes pickle’s executable semantics. It does not remove the need for provenance, isolation, and behavioural testing. An AI model is a compound software-and-data artefact and deserves the same control as a critical container image. Book an AI security assessment — we will review the entire path from public hub to production inference.