Towards Compact and Robust DNNs via Compression-aware Sharpness Minimization

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.

💡 Research Summary

The paper addresses the tension between model compression (pruning) and robustness to natural input variations, a critical issue for on‑device deep neural networks. While Sharpness‑Aware Minimization (SAM) has proven effective at improving robustness by encouraging flat minima in the continuous parameter space, its direct application to pruned models is problematic: SAM’s flatness does not guarantee resilience to the discrete structural changes introduced by pruning, and pruning before SAM can lock the network into a sub‑optimal architecture that limits achievable robustness.

To bridge this gap, the authors propose Compression‑aware Sharpness‑Aware Minimization (C‑SAM), a novel framework that transfers the concept of sharpness‑aware learning from weight perturbations to mask perturbations. Instead of perturbing the dense weight vector, C‑SAM introduces a learnable soft mask C ∈