Robustness Verification and the effect of Post Training Quantization#

Robustness to slight input variations is a serious topic for edge AI as the input considered by the embedded neural network can be noisy due to the measurement and sensibility of the sensors.

In this tutorial, we propose to address this topic, by evaluating the robustness to input variation of a neural network before and after its quantization with the post training quantization method in order to study the effect of the quantization on the robustness. We will consider a neural network trained classically and a neural network trained robustly via interval training as explained in the robust training tutorial.

Given a neural network \(N\), an input \(x\) and a perturbation threshold \(\epsilon\), then \(N\) is \(\epsilon\text{-robust}\) around \(x\) if

\[\forall x' \in B_{\|.\|_\infty}(x, \epsilon), \,\, arg\max\, N(x) = arg\max\, N(x').\]

We can extend the \(\epsilon\text{-robustness}\) of \(N\) with respect to an input set.

loc_robustness

Setup#

Let’s install the required aidge packages.

[ ]:

!pip install aidge-core aidge-onnx aidge-backend-cpu aidge-learning aidge-quantization torchvision

import aidge_core, aidge_onnx

Loading Dataset and Model#

First, we create utility functions:

  • to transform dataset from torchvision to an AIDGE dataprovider

  • to load a model from onnx and setup it in AIDGE for inference

[ ]:

import torchvision, torch
from aidge_learning.confiance.training.dataset import AidgeDataset


def create_mnist_dataprovider(
    train: bool, transform, size: int = None, batch_size: int = 1
):
    """Transform each image of the MNIST dataset and create an AIDGE dataprovider batched."""
    dataset = torchvision.datasets.MNIST(
        root="./data", train=train, download=True, transform=transform
    )
    if size is not None:
        dataset = torch.utils.data.Subset(dataset, torch.arange(size))

    aidge_dataset = AidgeDataset(data=dataset, input_shape=input_shape)
    return aidge_core.DataProvider(
        aidge_dataset, batch_size=batch_size, shuffle=False, drop_last=True
    )


def load_and_setup_model(model_path: str, input_shape):
    # Load the model
    aidge_model = aidge_onnx.load_onnx(model_path)

    # Set up the backend
    aidge_model.set_datatype(aidge_core.dtype.float32)
    aidge_model.set_backend("cpu")

    # Compile model
    aidge_model.compile("cpu", aidge_core.dtype.float32, dims=[input_shape])

    return aidge_model

We then load the MNIST dataset and both classic and robust network.

[ ]:

# Dataset information
size_test = 200
input_shape = (1, 1, 28, 28)

[ ]:

# The transformation applied to the MNIST dataset
transform = torchvision.transforms.Compose([torchvision.transforms.ToTensor()])

# Create the aidge dataprovider with 1000 images from the test set
test_dataprovider = create_mnist_dataprovider(
    train=False, transform=transform, size=size_test, batch_size=32
)

[ ]:

# Define path of the models
model_path = "small_FNN.onnx"
robust_model_path = "small_FNN_rob.onnx"

aidge_model = load_and_setup_model(model_path, input_shape)
aidge_robust_model = load_and_setup_model(robust_model_path, input_shape)

Training#

The architecture of both networks is the same: 4 MatMul layers of size 100 with each MatMul separated by ReLUs except for the first one which is a Sigmoid. The robust network was trained first with 5 warm-up epochs and 25 epochs where the perbutation size increased linearly up until 4 / 256. On the training set, it achieves roughly an accuracy of 97.5%.

Thus, for comparison purposes, the classic network was trained until it reaches an accuracy of roughly 97.5% on the training set which took 11 epochs.

Both network achieves 98.96 % on the test dataset.

The script to train both network is available in training_script.py.

Robustness evaluation#

Relying on the method from the robust training tutorial which builds an Aidge GraphView to do interval propagation in the network using interval arithmetic, we want to prove local robustness for \(\epsilon = 4\,/\,256\) by looking at the output intervals.

[ ]:

from aidge_learning.confiance.training.utils import evaluate_rob

perturbation_size = 4 / 256

[ ]:

print(f"Certified accuracy of the classical model with ε={perturbation_size}:")
_ = evaluate_rob(test_dataprovider, aidge_model, perturbation_size, input_shape)

print(f"Certified accuracy of the robust model with ε={perturbation_size}:")
_ = evaluate_rob(test_dataprovider, aidge_robust_model, perturbation_size, input_shape)

The classically trained neural network has a robust accuracy of 0 % while the robust network achieves a robust accuracy of 95 %.

This means that the classification of more than 95% of the input tested on the robust network remains the same even after any type of perturbation of magnitude \(\leq 4 \,/\, 256\). For the classical network, none of the inputs managed to be proven robust to perturbations.

Yet, this method didn’t prove that it was not robust to perturbation, i.e., we don’t know if there exists a perturbed input that changes the classification.

Robustness Evaluation Using PyRAT#

PyRAT stands for Python Reachability Assessment Tool and is based on abstract interpretation techniques with the aim of determining whether in a given neural network a certain state can be reach by the network from some initial parameters. To do so, it propagates different types of abstract domains along the different layers of the neural network up until the output layer. This approach will overapproximate some of the operations of the layers and thus the output space reached will always be bigger than the exact output space. If the output space reached is inside a certain space defining the property we want, then we can say that the property is verified. Otherwise, we cannot conclude.

One of the simplest over-approximation techniques that PyRAT implements is the interval arithmetic which is the same technique used by AIDGE for the robustness evaluation. PyRAT implements more complex methods allowing significant gains regarding the precision of the over-approximation, e.g. the zonotopes domain. But the more precise the method is, the longer the verification takes.

At the end of the verification process, PyRAT either returns:

  • True if the input is robust

  • False if a counterexample was found

  • Unknown if PyRAT could not conclude

PyRAT supports multiple type of network format such as AIDGE’s network or ONNX.

More information about PyRAT can be found in this paper: Verifying neural networks with PyRAT

PyRAT is proprietary CEA but is made available in Aidge for the purpose of this tutorial only for research purposes.

For more information contact:

Installing the PyRAT library#

PyRAT is made available through a compiled python version which requires Python 3.10, 3.11 or 3.12.

[ ]:

from get_pyrat_branch import get_pyrat_branch

branch_name = get_pyrat_branch()
!uv pip install git+https://git.frama-c.com/pub/pyrat.git@{branch_name}

We define a utility function to verify the aidge’s dataprovider in pyrat_script.py and set the verification parameters needed for PyRAT.

[ ]:

from pyrat_script import verif_pyrat_dataprovider

We use the zonotopes domain to evaluate the robustness of both networks.

[ ]:

verif_pyrat_dataprovider(
    model=aidge_model,
    dataprovider=test_dataprovider,
    eps=perturbation_size,
    desc="Classic",
    domains="zonotopes",
)

verif_pyrat_dataprovider(
    model=aidge_robust_model,
    dataprovider=test_dataprovider,
    eps=perturbation_size,
    desc="Robust",
    domains="zonotopes",
)

  • Classic network:

    PyRAT’s zonotopes domains allowed us to find that 28 % of the inputs were robust to input perturbations while none could be found using intervals. It also found that 28 % of the inputs were not robust to perturbations.

  • Robust network:

    PyRAT found that 95.8 % of the inputs were robust, which only an increase of 0.5 % compared to the interval domain. It also found that 3 % of the inputs were not robust.

With roughly 1 % of the inputs that were not conclusive, we clearly see that the robust network is more easily verifiable compared to the 43 % unknown results for the classic network.

Performing Robustness Evaluation on Quantified Networks#

Now, we want to study the impact of the Post Training Quantization (PTQ) on the robustness of the networks.

Quantization of both models#

First, we define the number of bits for quantizing the models and we modify our perturbation range to fit the quantized input.

We also create the calibration set needed for the PTQ method.

[ ]:

nb_bits = 8
size_calib_set = 1_000
perturb_size_quant = perturbation_size * (2**nb_bits)
fake_quantize = True

[ ]:

transform = torchvision.transforms.Compose([torchvision.transforms.ToTensor()])

calib_dataprovider = create_mnist_dataprovider(
    train=True, transform=transform, size=size_calib_set, batch_size=1
)

Then, we quantize both networks with a script inspired from the PTQ tutorial which can be found in quantify_network.py.

[ ]:

from quantify_network import quantify_network

aidge_quant = quantify_network(aidge_model, nb_bits, calib_dataprovider, fake_quantize)
aidge_quant_rob = quantify_network(
    aidge_robust_model, nb_bits, calib_dataprovider, fake_quantize
)

Evaluation of the robustness of the quantized networks#

First, we create a dataset with quantized input to assess the verified robustness of both quantized models.

[ ]:

# define the quantization scaling factor
scaling = 2 ** (nb_bits - 1) - 1

[ ]:

# each image is transformed into tensors and quantized
transform = torchvision.transforms.Compose(
    [
        torchvision.transforms.ToTensor(),
        torchvision.transforms.Lambda(lambda img: torch.round(img * scaling)),
    ]
)

test_quant_dataprovider = create_mnist_dataprovider(
    train=False, transform=transform, size=size_test, batch_size=32
)

Then, we evaluate their robustness using aidge’s interval.

[ ]:

print("Robustness evaluation of the non-Robust Network Quantized")
_ = evaluate_rob(test_quant_dataprovider, aidge_quant, perturb_size_quant, input_shape)

print("Robustness evaluation of the Quantized Robust Network Quantized")
_ = evaluate_rob(
    test_quant_dataprovider, aidge_quant_rob, perturb_size_quant, input_shape
)

As shown in the PTQ tutorial with 8 bits quantization, the quantized networks conserve almost the same accuracy:

  • 98.44 % for the classic network compared to 98.96 % for its non-quantized version.

  • 98.96 % for the robust network compared to 98.96 % for its non-quantized version.

Regarding the robustness, for the robust network, it yields a score of 84% which is a drop of only 11 % compared to its non-quantized version, meaning that we observe a conservation of the robustness property. Meanwhile the classic network remains at 0%.

This results suggest one of two things, either the quantization process makes the model less robust to perturbations or it makes the verification harder.

Now, we evaluate their robustness with PyRAT’s zonotopes to see if we also find an increase in robustness for both networks and how much of the input set is considered not safe under perturbation.

[ ]:

print(
    f"Inputs are on {nb_bits} bits, so ranging from 0 to {2**(nb_bits-1) - 1} and the perturbation size allowed is {perturb_size_quant}"
)

verif_pyrat_dataprovider(
    model=aidge_quant,
    dataprovider=test_quant_dataprovider,
    eps=perturb_size_quant,
    desc="Classic Quantized",
    domains="zonotopes",
)

verif_pyrat_dataprovider(
    model=aidge_quant_rob,
    dataprovider=test_quant_dataprovider,
    eps=perturb_size_quant,
    desc="Robust Quantized",
    domains="zonotopes",
)

The robustness of the quantized classic network remains at 0% even with a more precise over-approximation using PyRAT. Very few counter-examples could be found and more than 97 % of the inputs could not be proven robust nor unsafe.

Whereas, for the quantized robust network, more than 85 % of the inputs are robust to perturbations and only 13 % of the inputs remains unknown.

Conclusion#

Network

Accuracy

Robust

Unsafe

Unknown

Classic

98.96 %

28.65 %

28.65 %

42.70 %

Classic Quantized

98.44 %

0.0 %

2.08 %

97.92 %

Robust

98.96 %

95.83 %

3.12 %

1.04 %

Robust Quantized

98.96 %

85.94 %

1.04 %

13.02 %

Looking at the result for both networks, we clearly see an increase in the number of unknown instances meaning that the quantization process makes the networks harder to verify.

Even if the verification of quantized networks is harder, we still observe that the robustness is almost fully retained for the robustly trained network.