Towards Compact and Robust DNNs via Compression-aware Structural Stability Optimization

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. Motivated by this mismatch between robustness-aware training and compression, we propose Compression-aware Structural Stability Optimization (C-SSO). Rather than applying SAM before or after pruning, C-SSO optimizes pruning masks under stochastic structural perturbations, favoring subnetworks whose predictions remain stable under nearby connectivity changes while preserving robustness to input variations. Extensive experiments on CelebA-HQ, Flowers-102, CIFAR-10-C, and Tiny-ImageNet, spanning ResNet-18, GoogLeNet, MobileNet-V2, and DeiT-Ti, show that C-SSO consistently achieves higher certified robustness than competing baselines, with improvements of up to 42%, while maintaining task accuracy comparable to the corresponding unpruned models.