DivQAT: Enhancing Robustness of Quantized Convolutional Neural Networks against Model Extraction Attacks

TL;DR

DivQAT presents a divergence-based Quantization Aware Training approach that integrates a model extraction defense directly into the training of quantized CNNs. By adding a KL-divergence term between the quantized and full-precision models, it deliberately shifts output distributions to impede adversaries relying on probability-based queries, while maintaining accuracy. Empirical results across multiple vision datasets and attacks show improved robustness against extraction, with additional gains when combined with other defenses and across server/mobile quantization configurations. This work demonstrates that embedding security into the quantization process can yield stronger, multi-layered protection for edge-deployed models without sacrificing practical performance.

Abstract

Convolutional Neural Networks (CNNs) and their quantized counterparts are vulnerable to extraction attacks, posing a significant threat of IP theft. Yet, the robustness of quantized models against these attacks is little studied compared to large models. Previous defenses propose to inject calculated noise into the prediction probabilities. However, these defenses are limited since they are not incorporated during the model design and are only added as an afterthought after training. Additionally, most defense techniques are computationally expensive and often have unrealistic assumptions about the victim model that are not feasible in edge device implementations and do not apply to quantized models. In this paper, we propose DivQAT, a novel algorithm to train quantized CNNs based on Quantization Aware Training (QAT) aiming to enhance their robustness against extraction attacks. To the best of our knowledge, our technique is the first to modify the quantization process to integrate a model extraction defense into the training process. Through empirical validation on benchmark vision datasets, we demonstrate the efficacy of our technique in defending against model extraction attacks without compromising model accuracy. Furthermore, combining our quantization technique with other defense mechanisms improves their effectiveness compared to traditional QAT.

Figures (6)

• Figure 1: In model extraction attacks, the adversary leverages the prediction probabilities of their inputs to build a labeled dataset. Then the attacker uses this dataset to train their stolen model.
• Figure 2: Our Div-QAT approach: First, we initialize a quantized neural network based on a trained non-quantized model (red dashed arrow). Then, the quantized model is trained on the same training dataset as the original model. In each training step, input images are fed to the quantized model (forward pass) to provide prediction probabilities which will be used along with the true labels to compute the cross-entropy loss (green arrows). Additionally, we feed the training examples to the original (non-quantized) model and get their predictions to compute the KL-divergence between them and the quantized model's predictions (blue bold arrows). After that, the loss (Eq. \ref{['eq:loss_div_qat']}) is computed. The quantized neural network weights are updated through backpropagation and the training process continues.
• Figure 3: Results of Knockoffnets attack using different adversary architectures to extract quantized models obtained with QAT and DivQAT. (a) and (b) respectively show the adversary's classification error and the disagreement between the victim and the adversary. Each column represents the results for a dataset. Architectures are on the x-axis.
• Figure 4: Impact of $\alpha$ variation in DivQAT on the classification error of both the victim and the stolen model (adversary).
• Figure 5: Impact of $\alpha$ variation in DivQAT on the disagreement between the victim and the adversary.