-
-

Understanding Model Threats

This resource is designed to provide detailed information on various threat categories, helping you understand and mitigate potential risks in AI and machine learning systems.

Runtime Threats

Like a deserialization threat, runtime threats occur when untrusted data or code is used to reconstruct objects, leading to potential exploitation. The specific difference occurs in how the malicious code is triggered to execute. With a basic deserialization threat, this happens at model load time. A runtime threat is triggered when the model is used for inference or any form of execution. In AI and machine learning systems, this can result in malicious actors injecting harmful code during the deserialization process, exploiting vulnerabilities to gain unauthorized access or manipulate your systems behavior. Understanding deserialization threats is crucial for securing data integrity and preventing unauthorized code execution in your AI models.

Overview

Guardian scans for all layers in a Keras model and checks against Keras standard layer names. In case a layer in a model is found that does not match the Keras standard layers, Guardian would add an issue for it.

Keras is an open-source, deep learning library for developing machine learning models. Keras is a high level API that can work with other machine learning libraries seamlessly such as Tensorflow, PyTorch, and JAX.

Keras allows for custom object serialization (Ref). These custom objects can be custom operators, as well as custom layers or a combination of both custom layers with custom objects. Though Keras support custom object serialization, the .keras model file does not store the Python code for custom objects.

If a model reportedly has this issue it means that:

  1. The model format is detected as Keras v3 .keras.
  2. The model contains potentially unknown/custom layer which will execute code when the model is used for inference.

Key Points

  • Keras is a high level API for developing machine learning models.
  • Keras custom object serialization allows for custom layer and custom operator usage in a .keras model.
  • Attackers can use custom object serialization to inject malicious code in custom object definitions.
  • Loading compromised Keras model will execute malicious code.
  • Only load models from trusted sources.

Further reading:

  1. Keras custom object serialization

Impact

Since the code for custom objects does not become part of a Keras model, a user would need to load the custom object in memory before loading the Keras model with custom objects. This is why to make model loading a seamless experience for developers, model hubs such as Hugging Face give users the flexibility to trust remote code to load any accompanied code.

An attacker can exploit the flexibility of serializing custom object in a Keras model in addition with Hugging Face functionality to trust remote code to ship malicious code to a victims machine.

How The Attack Works

Best Practices

You should:

  1. Only load and execute models from trusted sources
  2. Implement a vetting process for third-party models before use
  3. Use sandboxing techniques when loading untrusted models
  4. Regularly update Keras and related libraries to benefit from security patches
  5. Use model formats that don't allow arbitrary code execution, such as SafeTensors, which provides a safe alternative to traditional serialization formats

Remediation

If possible, use a different model format like SafeTensors in order to remove this type of code injection attack from impacting your work entirely.

If not possible, reach out to the model creator and alert them that the model has failed our scan. You can even link to the specific page on our Insights Database to provide our most up to date findings.

The model provider should also report what they did to correct this issue as part of their release notes.

Protect AI's security scanner detects threats in model files
With Protect AI's Guardian you can scan models for threats before ML developers download them for use, and apply policies based on your risk tolerance.
Learn more