google-deepmind/batch-isolation-checker

google-deepmind/batch-isolation-checker

Language: Python

License: Apache-2.0

Stars: 5

Forks: 2

Open issues: 0

Created: 2025-05-23T15:01:44Z

Pushed: 2026-03-20T15:16:13Z

Default branch: main

Fork: no

Archived: no

README:

Batch Isolation Checker

*Paper: Architectural Backdoors for Within-Batch Data Stealing and Model Inference Manipulation*

The practice of batching requests from multiple users in machine learning inference poses significant security and privacy risks. These risks include breaches of confidentiality, potentially revealing information about other users' requests, and breaches of integrity, enabling the manipulation of other users' results. Fundamentally, these vulnerabilities arise from the lack of strict isolation guarantees between batched user requests, which creates the potential for both malicious and unintentional intra-batch side channels. At the same time, not batching is often not an option, since it greatly decreases deployment efficiency. While this issue is relevant to any deployment scenario employing request batching, it poses a particular threat to private inference, where it can completely compromise the confidentiality and integrity protections.

This prototype illustrates the potential vulnerabilities of models by using a maliciously modified version of the Gemma LLM. We subsequently introduce our model checker, which employs information flow control on the model graph to formally verify the absence of intra-batch side channels, thus ensuring the model's safe operation in batching scenarios. While this particular attack implementation is designed for the Gemma model, our model checker is more general and can be used with any ONNX model, provided it meets the checker's specified requirements.

We recommend downloading the Gemma ONNX model as an example. First, install and log in to the Hugging Face CLI. You can do this by following the instructions. Then, use the following command, where `` corresponds to the model directory used in the other commands:

poetry run huggingface-cli download aless2212/gemma-2b-it-fp16-onnx \
--revision 6c8b4d5173dcc1969cd733afff424e02ea173eeb \
--local-dir data/gemma-2b-it-fp16-onnx

huggingface-cli download aless2212/gemma-2b-it-fp16-onnx \
--revision 6c8b4d5173dcc1969cd733afff424e02ea173eeb \
--local-dir /gemma-2b-it-fp16-onnx

Check Interference

Use the following command to run an interference check on your CPU:

poetry run python batching_security_checker/check_interference.py \
--dir /home/ubuntu/batching-security-checker/data \
--family gemma \
--model gemma-2b-it-fp16-onnx

The command expects your model to be located in the --dir you provide. There are two ways to structure your model:

1. Directory: You can have a directory named gemma-2b-it-fp16-onnx inside your --dir. This directory should contain the actual ONNX model file (e.g., model.onnx).

2. File: Alternatively, you can have a file named gemma-2b-it-fp16-onnx.onnx directly within the --dir.

The --family flag (e.g., gemma) determines how dynamic input parameters are handled and how inputs are labeled.

Check Interference Requirements

We provide a tool to analyze ONNX models to determine if they meet the requirements for the batching security checker:

1. Operator Support: The model must consist of operators for which a label propagation implementation exists.

2. Static Tensor Shapes: All tensors in the model graph must have statically determined shapes (independent of concrete inputs). Models with dynamic input parameters can be converted to use fixed input shapes before running the batching security checker.

This tool offers two primary functionalities:

1. Model Analysis: This step generates a report for each ONNX model in the directory. It identifies all operators used by the model. To check whether all tensors have statically determined shapes, dynamic parameters are first replaced with fixed default values. Then, shape inference is performed on the model graph. The resulting information for each model is saved as a JSON file. This step can be skipped using the --skip flag, in which case existing JSON reports in the directory are used for the second step. 2. Report CLI This component provides a command-line interface to aggregate and analyze the JSON reports generated for all processed models. It reports any missing operators and suggests priorities for maximizing model coverage.

To analyze models and launch the interactive report CLI:

poetry run python batching_security_checker/report_cli.py --database /home/ubuntu/batching-security-checker/data/report_db.json --models /home/ubuntu/batching-security-checker/data/hf-onnx-community

See docs/ for the most recent requirements analysis of the ONNX models (from the onnx/models repository).

Implementing New Operators

The label propagation operators reside in the onnx_ops/ directory, mirroring the structure of the jaxonnxruntime project.

While most operators have individual files within onnx_ops/, unary and binary elementwise operators are grouped into onnx_ops/elementwise_unary.py and onnx_ops/elementwise_binary.py, respectively. This grouping is due to their shared label propagation logic. A tool (tools/generate_elementwise_ops.py) automates the generation of boilerplate code for these elementwise operators.

To add a new elementwise operator:

1. Add the operator to the operator lists in the beginning of the file: tools/generate_elementwise_ops.py

2. Generate the boilerplate code for all elementwise operators:

poetry run python tools/generate_elementwise_ops.py -- --root_dir $(pwd)

1. Include corresponding test cases in test/taint_ops_test.py, according to the structure in the jaxonnxruntime project's tests\onnx_ops_test.py file.

To manually add a new non-elementwise operator:

1. Create a file named .py in the onnx_ops/ directory, using existing operators as a template.

2. Include corresponding test cases in test/taint_ops_test.py, according to the structure in the jaxonnxruntime project's tests\onnx_ops_test.py file.

Implementing Placeholder Operators

Label propagation requires a corresponding…