26.2.4 Robustness-enhancing wrappers

The Wrapper Pattern in AI Security

Think of a wrapper as a security guard for your model. It doesn’t change how the model thinks, but it controls what comes in and what goes out. This pattern provides a clean separation of concerns: the data science team can focus on model performance, while the security team can build and maintain these defensive layers independently.

This approach is particularly valuable for protecting pre-trained or third-party models where you have limited or no ability to modify the internal workings. The wrapper becomes your primary point of control and defense.

Implementation: A Generic Wrapper Class

A good wrapper design is abstract and reusable. You can create a base class that handles the core logic of wrapping a model and then extend it with specific defensive techniques. This promotes modularity and makes it easy to chain multiple defenses together.

# A base class for creating model wrappers in Python.
# This abstract structure enforces a clean separation of concerns.
class ModelWrapper:
    def __init__(self, model):
        # The wrapper holds an instance of the model it protects.
        self._model = model

    def predict(self, data):
        # 1. Intercept and process the input.
        processed_input = self._preprocess(data)
        
        # 2. Pass the sanitized input to the actual model.
        model_output = self._model.predict(processed_input)
        
        # 3. Intercept and process the output.
        final_output = self._postprocess(model_output)
        
        return final_output

    def _preprocess(self, data):
        # Default behavior: do nothing. Subclasses will override this.
        return data

    def _postprocess(self, output):
        # Default behavior: do nothing. Subclasses will override this.
        return output

Example: Input Smoothing Wrapper

Let’s implement a specific wrapper that applies a Gaussian blur to input images. This is a simple defense that can degrade the effectiveness of finely-tuned adversarial noise.

# Assumes usage of libraries like OpenCV (cv2) and NumPy.
import cv2
import numpy as np

class GaussianSmoothingWrapper(ModelWrapper):
    def __init__(self, model, kernel_size=(5, 5), sigma=1.5):
        super().__init__(model)
        self.kernel_size = kernel_size
        self.sigma = sigma
        print(f"Initialized GaussianSmoothingWrapper with kernel {kernel_size}")

    def _preprocess(self, image_batch):
        # Apply Gaussian blur to each image in the batch.
        smoothed_batch = []
        for img in image_batch:
            smoothed_img = cv2.GaussianBlur(img, self.kernel_size, self.sigma)
            smoothed_batch.append(smoothed_img)
        return np.array(smoothed_batch)

To use this, you would simply instantiate your model and then wrap it:

original_model = load_my_image_classifier()
secured_model = GaussianSmoothingWrapper(original_model)
prediction = secured_model.predict(user_supplied_image)

Common Wrapper-Based Defenses

The wrapper pattern can host a variety of defenses, often combining techniques from input sanitization and anomaly detection. Here are a few common types:

• Perturbation Filtering: Techniques like blurring, JPEG compression, or spatial smoothing that aim to destroy or reduce the impact of adversarial noise before the model sees the input.
• Feature Squeezing: Reducing the complexity of the input space. For example, reducing the color bit depth of an image (e.g., from 24-bit to 8-bit color). If the model’s prediction changes drastically after squeezing, the input may be adversarial.
• Detector Integration: The _preprocess method can house an anomaly or adversarial detector (see Chapter 26.2.3). If an input is flagged as malicious, the wrapper can reject it outright instead of passing it to the model.
• Confidence Thresholding: An output-side wrapper (in _postprocess) can check the model’s confidence scores. If the top prediction’s confidence is below a certain threshold, the wrapper can return a generic “uncertain” response instead of a potentially incorrect, low-confidence prediction.
• Output Stabilization: The wrapper runs the model on the original input and a slightly perturbed version. If the predictions are inconsistent, it signals a potential lack of robustness or an adversarial attack, flagging the result for review.

Advantages and Limitations

While powerful, wrappers are not a silver bullet. You must understand their trade-offs to deploy them effectively.

Ultimately, robustness-enhancing wrappers are a crucial tool for practical AI security. They provide a flexible and immediate way to harden models, acting as a first line of defense in a layered security strategy.