Concept-Aware Pruning for Robust Deep Neural Networks

Abstract

Deep Neural Networks (DNNs) are increasingly used in high-stakes settings, such as healthcare, education, hiring, lending, and autonomous vehicles. These models consume massive amounts of environmental and economic resources, and are largely inaccessible to researchers and users without access to expensive high-performance compute. Therefore, it has become essential to compress DNNs to reduce costs in storage, compute, and energy. However, compression comes with a hidden cost: existing methods degrade performance for minority subgroups and out-of-distribution data, creating unfair and non-robust models. Our novel framework, Concept-aware Network Pruning (CNP), mitigates this shortcut learning by identifying and removing human-understandable concepts within a model during pruning. CNP augments a pretrained DNN with a virtual concept layer in which nodes represent semantic concepts, ablates human-identified target concepts, and then applies node-level pruning so that nodes serving the ablated concepts are naturally irrelevant. We demonstrate on binary image classification tasks that CNP can mitigate model reliance on targeted concepts, leading to improved performance over vanilla pruning techniques for under-represented subgroups while maintaining overall accuracy. We evaluate applications in robustness and fairness in compression.