Towards Compact and Robust DNNs via Compression-aware Sharpness Minimization

TL;DR

This work tackles the problem of obtaining compact yet certifiably robust DNNs for edge devices by linking pruning with robustness. It introduces Compression-aware ShArpness Minimization (C-SAM), which shifts sharpness minimization from weight space to pruning mask space, optimizing a composite objective with stability, ratio, and consistency terms to promote a flat, robust mask landscape. Through extensive experiments on CelebA-HQ, Flowers-102, and CIFAR-10-C across ResNet-18, GoogLeNet, and MobileNet-V2, C-SAM achieves substantial increases in Probabilistically Certified Accuracy (PCA) up to 42% while preserving task accuracy close to unpruned baselines. The approach generalizes to both unstructured and structured pruning, and ablation analyses confirm that each component of the objective contributes to robustness and compression, making C-SAM a practical pathway for reliable edge AI under memory constraints.

Abstract

Sharpness-Aware Minimization (SAM) has recently emerged as an effective technique for improving DNN robustness to input variations. However, its interplay with the compactness requirements of on-device DNN deployments remains less explored. Simply pruning a SAM-trained model can undermine robustness, since flatness in the continuous parameter space does not necessarily translate to robustness under the discrete structural changes induced by pruning. Conversely, applying SAM after pruning may be fundamentally constrained by architectural limitations imposed by an early, robustness-agnostic pruning pattern. To address this gap, we propose Compression-aware ShArpness Minimization (C-SAM), a framework that shifts sharpness-aware learning from parameter perturbations to mask perturbations. By explicitly perturbing pruning masks during training, C-SAM promotes a flatter loss landscape with respect to model structure, enabling the discovery of pruning patterns that simultaneously optimize model compactness and robustness to input variations. Extensive experiments on CelebA-HQ, Flowers-102, and CIFAR-10-C across ResNet-18, GoogLeNet, and MobileNet-V2 show that C-SAM consistently achieves higher certified robustness than strong baselines, with improvements of up to 42%, while maintaining task accuracy comparable to the corresponding unpruned models.

Figures (4)

• Figure 1: An illustration of the C-SAM framework, where pruning masks are optimized for stability under mask perturbations ($\mathcal{L}_{\text{stab}}$), robustness against semantic input variations ($\mathcal{L}_{\text{ratio}}$), soft-hard mask consistency ($\mathcal{L}_{\text{consis}}$), and $L_1$ regularization.
• Figure 2: Loss landscape visualization for models trained with (a) adversarial loss and (b) C-SAM loss. The orange star indicates the location of the initial model.
• Figure 3: The parameter distributions for fully connected layer and layer4.conv2.
• Figure 4: The demonstration of the transformation space of $\mathbb{S}_T(x)$, where $x$ is the original image, labeled 0.

Theorems & Definitions (3)

• proof
• proof
• proof