PickleBall: Secure Deserialization of Pickle-based Machine Learning Models

2025-08-21

AI Summary

Unsafe pickle-based model serialization in Hugging Face and similar model repositories poses severe deserialization security risks—44.9% of popular models still rely on insecure pickle, while existing scanners suffer from high false-positive and false-negative rates. Method: We propose PickleGuard, a secure model loading framework integrating static source-code analysis with dynamic policy enforcement. It automatically infers fine-grained, human-interpretable safety policies by analyzing model source code and replaces the native pickle loader via an embedded, dependency-free injection mechanism. Contribution/Results: PickleGuard achieves zero-dependency, plug-and-play deployment while maintaining 79.8% compatibility with benign models and providing 100% detection and prevention of malicious models—including complete mitigation of arbitrary code execution attacks via unsafe function calls. Extensive evaluation demonstrates that PickleGuard significantly outperforms state-of-the-art tools in both security guarantees and practical usability.

Technology Category

• Natural Language Processing: Safety and Robustness
• Computer Vision: Large Vision Models
• Machine Learning: Hardware-aware ML

Application Category

• User Modeling, Personalization and Recommendation: Attacks and countermeasures in recommendation systems
• Security and Privacy: Large-scale security measurements
• Web Mining and Content Analysis: Large pretrained models with web data

Abstract

Machine learning model repositories such as the Hugging Face Model Hub facilitate model exchanges. However, bad actors can deliver malware through compromised models. Existing defenses such as safer model formats, restrictive (but inflexible) loading policies, and model scanners have shortcomings: 44.9% of popular models on Hugging Face still use the insecure pickle format, 15% of these cannot be loaded by restrictive loading policies, and model scanners have both false positives and false negatives. Pickle remains the de facto standard for model exchange, and the ML community lacks a tool that offers transparent safe loading. We present PickleBall to help machine learning engineers load pickle-based models safely. PickleBall statically analyzes the source code of a given machine learning library and computes a custom policy that specifies a safe load-time behavior for benign models. PickleBall then dynamically enforces the policy during load time as a drop-in replacement for the pickle module. PickleBall generates policies that correctly load 79.8% of benign pickle-based models in our dataset, while rejecting all (100%) malicious examples in our dataset. In comparison, evaluated model scanners fail to identify known malicious models, and the state-of-art loader loads 22% fewer benign models than PickleBall. PickleBall removes the threat of arbitrary function invocation from malicious pickle-based models, raising the bar for attackers to depend on code reuse techniques.

Problem

• Secure deserialization of pickle-based ML models
• Preventing malware delivery through compromised models
• Addressing shortcomings in existing model loading defenses

Innovation

• Static analysis for custom policy generation
• Dynamic enforcement during load time
• Drop-in replacement for pickle module