Self-supervised Feature-Gate Coupling for Dynamic Network Pruning

TL;DR

The paper addresses distortion in dynamic channel pruning caused by misaligned feature and gate distributions. It introduces Self-supervised Feature-Gate Coupling (FGC), a plug-and-play two-step loop that (i) discovers instance neighborhoods in feature space via $k$NN and (ii) regularizes gating modules with a contrastive self-supervised loss (InfoNCE-based) to align neighborhoods in the gate space, supported by memory banks. The method optimizes a multi-term objective $\mathcal{L} = \mathcal{L}_{ce} + \eta \sum_l \mathcal{L}_g^l + \rho \sum_l \mathcal{L}_0^l$, promoting sparse gates while preserving discriminative power; theoretical analysis links $\mathcal{L}_g^l$ to mutual information $I(f^l, g^l)$. Empirically, FGC improves accuracy-computation trade-offs across CIFAR-10/100, ImageNet, and transferring tasks (VOC detection, Cityscapes segmentation), demonstrating robust feature-gate alignment and strong generalization to diverse architectures.

Abstract

Gating modules have been widely explored in dynamic network pruning to reduce the run-time computational cost of deep neural networks while preserving the representation of features. Despite the substantial progress, existing methods remain ignoring the consistency between feature and gate distributions, which may lead to distortion of gated features. In this paper, we propose a feature-gate coupling (FGC) approach aiming to align distributions of features and gates. FGC is a plug-and-play module, which consists of two steps carried out in an iterative self-supervised manner. In the first step, FGC utilizes the $k$-Nearest Neighbor method in the feature space to explore instance neighborhood relationships, which are treated as self-supervisory signals. In the second step, FGC exploits contrastive learning to regularize gating modules with generated self-supervisory signals, leading to the alignment of instance neighborhood relationships within the feature and gate spaces. Experimental results validate that the proposed FGC method improves the baseline approach with significant margins, outperforming the state-of-the-arts with better accuracy-computation trade-off. Code is publicly available.

Figures (10)

• Figure 1: Up: inconsistent distributions of features and gates without feature-gate coupling (FGC). Down: aligned distributions of features and gates with FGC, where instances with similar semantics would also have similar gates and vice versa. The alignment indicates that FGC allocates consistent sub-networks with specific representation capacities for semantically similar instances, which reduces the representation redundancy and thereby improves the pruning effects.
• Figure 2: Flowchart of the proposed feature-gate coupling (FGC) method for DNP. FGC consists of two iterative steps: neighborhood relationship exploration for instance-wise neighborhood relationship modeling and feature-gate alignment for gating module regularization.
• Figure 3: Test accuracies and computation of the pruned ResNet-20 w.r.t the number of nearest neighbors $k$.
• Figure 4: Test accuracies and computation of the pruned ResNet-20 w.r.t coefficient $\eta$.
• Figure 5: Computation-accuracy trade-offs of pruned ResNet-{20, 32}. The "w/" and "w/o" refer to "with" and "without", respectively, indicating whether the FGC method is used or not. (Ditto for other tables or figures.)