-
-

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

The LiteRT format (formerly known as Tensorflow Lite), saves models into an optimized representation designed for mobile and edge devices. LiteRT uses the flatbuffers serialization to reduce the model size and memory footprints. In this sense, LiteRT models are execution-focused rather than development-focused, containing only the necessary components for inference without the original model building logic.

The LiteRT format enhances deployment flexibility as it's optimized for resource-constrained environments. Its base capabilities support a subset of the original Tensorflow operators along with a set of LiteRT core operators, but other custom operators can be specified.

Models flagged for this threat meet the following criteria:

  1. The model format is detected as LiteRT.
  2. The model contains a potentially malicious operator which will execute code when the model is used for inference.

Key Points:

  1. Models saved using LiteRT should be deemed as running “packaged code”.
  2. The LiteRT format stores a flattened representation of the model's computation graph and trained parameters in an optimized FlatBuffer structure.
  3. LiteRT files can contain TensorFlow operators. Some of the standard/known Tensorflow operators can be exploited by attackers such as ReadFile and WriteFile to gain access to confidential information on a users machine.
  4. Only use/load models from trusted sources.

Background Information

LiteRT Format

The LiteRT format, short for Lite Runtime, was originally known as Tensorflow Lite (TFLite). It saves models into an optimized representation designed for mobile and edge devices. LiteRT uses the flatbuffers serialization to reduce the model size and memory footprints. In this sense, LiteRT models are execution-focused rather than development-focused, containing only the necessary components for inference without the original model building logic.

The LiteRT format enhances deployment flexibility since it's optimized for resource-constrained environments. LiteRT currently supports models authored in PyTorch, JAX, and Keras in addition to the original Tensorflow. Its base capabilities support a subset of the original Tensorflow operators along with a set of LiteRT core operators, but other custom operators can be specified through TensorFlow’s Custom Operator functionality.

Custom Operators

Guardian reports an issue if any of the operators in a LiteRT file is not a standard TensorFlow operator or a builtin LiteRT operator. This is because unknown TensorFlow operators can be used by attackers to inject malicious code which would execute when an unsuspecting user would load the model.

TensorFlow custom operators extend the framework's native functionality by allowing users to implement specialized operations not available in the standard library. These user-defined operators integrate seamlessly with TensorFlow's existing ecosystem, enabling the development of more efficient and tailored machine learning solutions.

Though custom operators provide additional flexibility when implementation of novel algorithms or integration of domain specific operations is required, it also gives attackers an opportunity to export malicious code to a users/victims ML system/pipeline.

Further reading:

  1. TensorFlow lazy execution using graphs
  2. LiteRT format
  3. Custom operator
  4. Standard TF operators
  5. Builtin LiteRT operators
  6. TensorFlow and LiteRT Operator compatibility

Impact

Subject to attacker’s proficiency but any of the following is possible:

  1. Collect user data and send it back using known TensorFlow operators ReadFile and WriteFile
  2. Steal credentials and passkeys by reading files
  3. Gain further access to the system by overwriting credential files
  4. Read and steal model data
  5. Poison other models or datasets by overwriting the files

Note: Malicious code execution using standard TF operators can be achieved without impacting a models performance - the user may never know that the attack has happened or is ongoing.

How The Attack Works

Remediation

If possible, avoid using the LiteRT format model since it contains code that will get executed when the model is loaded.

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