Towards Compact and Robust DNNs via Compression-aware Sharpness Minimization

TL;DR

This paper introduces C-SAM, a novel training framework that enhances the robustness and compactness of deep neural networks by making the loss landscape flatter with respect to pruning masks, leading to better robustness and efficiency.

Contribution

C-SAM shifts sharpness-aware minimization from parameters to pruning masks, enabling joint optimization of model robustness and compactness during training.

Findings

C-SAM outperforms baselines with up to 42% higher certified robustness.

C-SAM maintains accuracy comparable to unpruned models across datasets.

C-SAM discovers pruning patterns that improve robustness and compactness simultaneously.

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.