Self-supervised Feature-Gate Coupling for Dynamic Network Pruning

We introduce the basic formulation of gated convolutional layers. For an instance in the dataset $\{x_{i},y_{i}\}_{n=1}^{N}$, its feature in the $l$-th layer of CNNs is denoted as:

where $x_{i}^{l-1}\in\mathbb{R}^{C^{l-1}\times W^{l-1}\times H^{l-1}}$ and $f_{i}^{l}\in\mathbb{R}^{C^{l}\times W^{l}\times H^{l}}$ are respectively the input and output of the $l$-th convolutional layer. $C^{l}$ is the channel number. $W^{l}$ and $H^{l}$ represent the width and height of the feature map. $w^{l}$ denotes the convolutional weight matrix. $\ast$ is the convolution operator, and $\rm{ReLU}$ is the activation function. Batch normalization is applied after each convolution layer.

Dynamic channel pruning methods temporarily remove channels to reduce related computation in both preceding and succeeding layers, where a binary gate vector, generated from the preceding feature $x_{i}^{l-1}$, is used to turn on or off corresponding channels. The gated feature in the $l$-th layer is denoted as:

where $g_{i}^{l}$ denotes the gate vector and $\odot$ denotes the channel-wise multiplication.

The gate vectors come from various gating modules with uniform functionality. Without loss of generality, we use the form of gating module [7], which is defined as:

The gating module first goes with the global average pooling layer $p$, which is used to generate the spatial descriptor of the input feature $x_{i}^{l-1}$. A lightweight network $M$ with two fully connected layers is then applied to generate the gating probability $\pi_{i}^{l}=Linear(p(x_{i}^{l-1}))$. Finally, the BinConcrete [31] activation function $B$ is used to sample from $\pi_{i}^{l}$ to get the binary mask $g_{i}^{l}$.

Neighborhood Relationship Exploration. Instance neighborhood relationships vary in feature spaces with different semantic levels, $e.g$, in a low-level feature space, instances with similar colors or texture may be closer, while in a high-level feature space, instances from the same class may gather together. Thus, human annotations could not provide adaptive supervision of instance neighborhood relationships across different network stages. We propose to utilize the $k$NN algorithm to adaptively explore instance neighborhood relationships in each feature space and extract self-supervision accordingly for feature-gate distribution alignment.

Note that the exploration could be adopted in arbitrary layers, $e.g.$, the $l$-th layer of CNNs. In practice, in the $l$-th layer, the feature $f_{i}^{l}$ of the $i$-th instance is first fed to a global average pooling layer, as $\hat{f}_{i}^{l}=p(f_{i}^{l})$. The dimension of pooled feature reduces from $C^{l}\times W^{l}\times H^{l}$ to $C^{l}\times 1\times 1$, which is more efficient for processing. We then compute its similarities with other instances in the training dataset as:

where $\mathcal{S}^{l}\in\mathbb{R}^{N\times N}$ denotes the similarity matrix and $\mathcal{S}^{l}[i,j]$ denotes similarity between the $i$-th instance and the $j$-th instance. $N$ denotes the number of total training instances. To identify nearest neighbors of $x_{i}$, we sort the $i$-th row of $\mathcal{S}_{l}$ and return the subscripts of the $k$ biggest elements, as

The set of subscripts $\mathcal{N}_{i}^{l}$, $i.e.$, indexes of neighbor instances in the training dataset, are treated as self-supervisory signals to regularizing gating modules.

To compute the similarity matrix $\mathcal{S}^{l}$ in Eq. (4), $\hat{f}_{j}^{l}$ for all instances are required. Instead of exhaustively computing these pooled features at each iteration, we maintain a feature memory bank $\mathcal{V}_{f}^{l}\in\mathbb{R}^{N\times D}$ to store them [8]. $D$ is the dimension of the pooled feature, $i.e.$, the number of channels $C^{l}$. Now for the $i$-th pooled feature $\hat{f}_{i}^{l}$, we compute its similarities with vectors in $\mathcal{V}_{f}^{l}$. After that, the memory bank is updated in a momentum way, as

where the momentum coefficient is empirically set to 0.5. All vectors in the memory bank are initialized as unit random vectors.

Feature-Gate Alignment. We identify the nearest neighbor set $\{x_{j},j\in\mathcal{N}_{i}^{l}\}$ of each instance $x_{i}$ in the feature space. Thereby, instances in each neighbor set are treated as positives to be pulled closer while others as negatives to be pushed away in the gate space. To this end, we utilize the contrastive loss in a self-supervised manner.

Specifically, by defining nearest neighbors of an instance as its positives, the probability of an input instance $x_{j}$ being recognized as nearest neighbors of $x_{i}$ in the gate space, as

where $\pi_{i}^{l}$ and $\pi_{j}^{l}$ are gating probabilities, the gating module outputs of $x_{i}$ and $x_{j}$ respectively. $\tau$ is a temperature hyper-parameter and set as 0.07 according to [8]. Then the loss function is derived after summing up $p(j|\pi_{i}^{l})$ of all $k$ neighbors, as

Since nearest neighbors are defined as positives and others as negatives, Eq. (8) is minimized to draw previous nearest neighbors close, $i.e.$, reproducing the instance neighborhood relationships within the gate space. We maintain a gate memory bank $\mathcal{V}_{g}^{l}\in\mathbb{R}^{N\times D}$ to store gating probabilities. For each $\pi_{i}^{l}$, its positives are taken from $\mathcal{V}_{g}^{l}$ according to $\mathcal{N}_{i}^{l}$. After computing Eq. (8), $\mathcal{V}_{g}^{l}$ is updated in a similar momentum way as updating the feature memory bank. All vectors in $\mathcal{V}_{g}^{l}$ are initialized as unit random vectors.

Overall Objective Loss. Besides the proposed contrastive loss, we utilize the cross-entropy loss ($\mathcal{L}_{ce}$) for the classification task and the $\mathcal{L}_{0}$ loss to introduce sparsity. $\mathcal{L}_{0}^{l}=\lVert\pi_{i}^{l}\rVert_{0}$ is used in each gated layer, where $\lVert\cdot\rVert_{0}$ is $\mathcal{L}_{0}$-norm.

The overall loss function of our method can be formulated as:

where $\mathcal{L}_{g}^{l}$ denotes the contrastive loss adopted in the $l$-th layer and $\eta$ the corresponding coefficient, $\Omega$ denotes the layers $\mathcal{L}_{g}^{l}$ applied in. The coefficient $\rho$ controls the sparsity of pruned models, $i.e.$, a bigger $\rho$ leads to a sparser model with lower accuracy, while a lower $\rho$ derives less sparsity and higher accuracy.

The learning procedure of FGC is summarized as Algorithm 1.

We analyze the plausibility of designed contrastive loss (Eq. (8)) from the perspective of mutual information maximization. According to [30], the InfoNCE loss is used to maximize the mutual information between learned representation and input signals. It is assumed that minimizing the contrastive loss leads to the improved mutual information between features and corresponding gates.

We define the mutual information between gates and features as:

where the density ratio based on $f_{j}^{l}$ and $g_{i}^{l}$ is proportional to a positive real score as $q(f_{j}^{l},g_{i}^{l})\propto p(f_{j}^{l}|g_{i}^{l})/p(f_{j}^{l})$ [30], which is $q(f_{j}^{l},g_{i}^{l})=\exp(\pi_{j}^{lT}\cdot\pi_{i}^{l}/\tau)$ in our case. $\pi_{i}^{l}$ is the gating probabilities for $g_{i}^{l}$ and $\pi_{j}^{l}$ belongs to its positives. The InfoNCE loss item could be rewritten as:

By taking the density ratio into the equation, we could derive the relationship between InfoNCE loss and mutual information, where the superscript $l$ is omitted for brevity,

where $N$ is the number of training samples, and $X_{neg}$ denotes the set containing negative instances. The inequality holds because $f_{j}$ and $g_{i}$ is not independent and $p(f_{j}|g_{i})\geq p(f_{j})$.

As $I(f_{j},g_{i})\geq\mathop{\log}(N)-\mathcal{L}_{NCE}$, we could conclude that minimizing the InfoNCE loss item would maximize a lower bound on mutual information, justifying the efficacy of proposed contrastive loss $\mathcal{L}_{g}^{l}=\sum_{j\in\mathcal{N}_{i}^{l}}\mathcal{L}_{NCE}^{l}$. Accordingly, the improvement of mutual information between features and corresponding gates facilities distribution alignment. (Please refer to Fig. 10.)

Datasets. Experiments on image classification are conducted on benchmarks including CIFAR10/100 [32] and ImageNet [33]. CIFAR10 and CIFAR100 consist of 10 and 100 categories, respectively. Both datasets contain 60K colored images, 50K and 10K of which are used for training and testing. ImageNet is an important large-scale dataset composed of 1000 object categories, which has 1.28 million images for training and 50K images for validation.

We conduct object detection on PASCAL VOC 2007 [34], the trainval and test sets of which are used for training and testing, respectively. Instance segmentation is conducted on Cityscapes [35]. We use the train set for training and the validation set for testing.

Implementation Details. For CIFAR10 and CIFAR100, training images are applied with standard data augmentation schemes [7]. We use the ResNet of various depths, $i.e.$, ResNet–{20, 32, 56}, to fully demonstrate the effectiveness of our method. On CIFAR10, models are trained for 400 epochs with a mini-batch of 256. A Nesterov’s accelerated gradient descent optimizer is used, where the momentum is set to 0.9 and the weight decay is set to 5e-4. It is noticed that no weight decay is applied to parameters in gating modules. As for the learning rate, we adopt a multiple-step scheduler, where the initial value is set to 0.1 and decays by 0.1 at milestone epochs [200, 275, 350]. We also evalutate FGC on WideResNet [36] and MobileNetV2 [37] (a type of lightweight network). Specifically, models are trained for 200 epochs with a mini-batch of 256. The initial learning rate is 0.1 and decays by 0.2 at epochs [60, 120, 160]. On CIFAR100, models are trained for 200 epochs with a mini-batch of 128. The same optimization settings as CIFAR10 are used. The initial learning rate is set to 0.1, then decays by 0.2 at epochs [60, 120, 160]. For ImageNet, the data augmentation scheme is adopted from [18]. We only use ResNet-18 to validate the effectiveness on the large-scale dataset due to time-consuming training procedures. We train models for 130 epochs with a mini-batch of 256. A similar optimizer is used except for the weight decay of 1e-4. The learning rate is initialized to 0.1 and decays by 0.1 at epochs [40, 70, 100]. FGC is employed in the last few residual blocks of ResNets, $i.e.$, the last two blocks of ResNet-{20,32}, the last four blocks of ResNet-56, and the last block of ResNet-18. Hyper-parameters $\eta$ and $k$ are set to 0.003 and 200 by default, respectively. We set the l0 loss coefficient $\rho$ as 0.4 to introduce considerable sparsity into gates.

For PASCAL VOC 2007, we use the detector and the data augmentation described in Faster R-CNN [38]. The Faster R-CNN model consists of the ResNet backbone and Feature Pyramid Network [39] (FPN). Specifically, We replace the backbone with the gated ResNet-{34, 50} pre-trained on ImageNet, leaving other components unchanged. We fine-tune the detector under the 1x schedule where the initial learning rate is 0.001, the momentum is 0.9, and the weight decay is 1e-4. No weight decay is applied in gating modules. FGC is employed in the last stage of the backbone. During the fine-tuning process, hyper-parameters $\eta$ and $k$ are set to 0.0001 and 1, respectively.

For Cityscapes, we follow the data augmentation and the method as PSPNet [40]. PSPNet with a pre-trained ResNet-50 backbone is trained with a mini-batch of 6 for 300 epochs. The optimizer is an SGD optimizer with the poly learning rate schedule, where the initial learning rate is 0.02 and the decay exponent is 0.9. The weight decay and momentum are 5e-4 and 0.9, respectively. The weight decay of gating modules is 2e-6. In practice, We employ FGC in the last stage of the backbone. We set $\eta$ and $k$ to 0.0001 and 4.

We conduct ablation studies on CIFAR10 to validate the effectiveness of our method by exploring hyper-parameters, residual blocks, neighborhood relationships, and accuracy-computation trade-offs.

Hyper-parameters of Contrastive Loss. In Table 1 and Fig. 3, we show the performance of pruned models given different $k$, $i.e.$, the number of nearest neighbors. For a fair comparison, the computation of each model is tuned to be roughly the same so that it can be evaluated by accuracy. As shown by Fig. 3, FGC is robust about $k$ in a wide range, $i.e.$, the accuracy reaches the highest value around $k=$200, and keeps stable when $k$ falls into the range (20, 1024).

In Table 2 and Fig. 4, we show the performance of pruned models under different $\eta$, $i.e.$, the loss coefficient. Similarly, the computation of each model is also tuned to be almost the same. The contrastive loss brings a significant accuracy improvement compared with models without FGC. With the increase of $\eta$, the accuracy improves from 91.4% to 92.09%. The best performance is achieved when $\eta=0.003$. While accuracy slightly drops when $\eta$ is larger than 0.01, it remains better than that at $\eta=0$.

Residual Blocks. In Table 3, we exploit whether FGC is more effective when employed in multiple residual blocks. The first row of Table 3 denotes models trained without FGC. After gradually removing shallow residual blocks from adopting FGC, experimental results show that FGC in shallow stages brings a limited performance gain despite more memory banks, $e.g.$, total $18$ memory banks used when employing FGC in all blocks. However, FGC in the last stage, especially from the $8$-$th$ to the $9$-$th$ blocks, contributes better to the final performance, $i.e.$, the lowest error (7.91%) with considerable computation reduction (55.1%). It is concluded that the residual blocks in deep stages benefit more from FGC, probably due to more discriminative features and more conditional gates. Therefore, taking both training efficiency and performance into consideration, we employ FGC in the deep stages of ResNets.

Neighborhood Relationship Sources. In Table 4, we validate the necessity of exploring instance neighborhood relationships in feature spaces by comparing the counterparts. Specifically, we randomly select $k$ instances with the same label as neighbors for each input instance, denoted as “Label”; we also utilize $k$NN in the gate space to generate nearest neighbors, denoted as “Gate”. Moreover, since FGC is applied in multiple layers separately, an alternative adopts neighborhood relationships from the same layer, $e.g.$, the last layer. We define this case as “Feature” with “Shared” instance neighborhood relationships.

Results in Table 4 demonstrate that our method, $i.e.$, “Feature” with independent neighbors, prevails over all the counterparts, which fail to explore proper instance neighborhood relationships for alignment. For example, labels are too coarse or discrete to describe neighborhood relationships. Exploring neighbors in the gate space enhances its misalignment with the feature space. Therefore the performance, either error or pruning ratio (8.50%, 51.6%), is even worse than the model without FGC. Moreover, “Shared” neighbors from the same feature space could not deal with variants of intermediate features, thereby limiting the improvement.

We conclude that the proposed FGC consistently outperforms the compared methods, which attributes to the alignment of instance neighborhood relationships between feature and gate spaces in multiple layers.

Computation-Accuracy Trade-offs. In Fig. 5, we show the computation-accuracy trade-offs of FGC and the baseline method by selecting $\mathcal{L}_{0}$ loss coefficient $\rho$ from values of {0.01, 0.02, 0.05, 0.1, 0.2, 0.4}. Generally speaking, the accuracy of pruned models increases as the total computational cost increases. Compared with pruned models without FGC, the pruned ResNet-32 with FGC obtains higher accuracy with lower computational cost, as shown in the top right of Fig. 5. At the left bottom of Fig. 5, the pruned ResNet-20 with FGC achieves  $0.5\%$ higher accuracy using the same computational cost. The efficacy of FGC is clarified that increasing performances compared with the baseline models can be observed at lower pruning ratios.

Visualization Analysis. In Fig. 6, we use UMAP [41] to visualize the correspondence of features and gates during the training procedure. Specifically, we show features and gates from the last residual block (conv3-3) of ResNet-20, where features are trained to be discriminative, $i.e.$, easy to be distinguished in the embedding space. It’s seen that features and gates are simultaneously optimized and evolve together during training. In the last epoch, as instance neighborhood relationships have been well reproduced in the gate space, we can observe gates with a similar distribution as features, $e.g.$, instances of the same clusters in the feature space are also assembled in the gate space. Intuitively, the coupling of features and gates promises better predictions of our pruned models with given pruning ratios.

In Fig. 7, we compare feature and gate embeddings with pruned models in BAS [7]. Specifically, they are drawn from the last residual block (conv3-3) of ResNet-20 and the second to last residual block (conv3-4) of ResNet-32. As shown, our method brings pruned models with coupled features and gates, compared with BAS. For example, “bird” (green), “frog” (grey) and “deer” (brown) instances entangle together (left bottom of Fig. 7a), while adopting FGC (right bottom of Fig. 7a), they are clearly distinguished as corresponding features. Similar observations are seen in Fig. 7b, validating the alignment effect of FGC, which facilitates the model to leverage specific inference paths for semantically consistent instances and thereby decrease representation redundancy.

In Fig. 8, we visualize the execution frequency of channels in pruned ResNet-20 for each category on CIFAR10. In shallow layers ($e.g.$, conv1-3), gates almost entirely turn on or off for arbitrary categories, indicating why FGC is unnecessary: a common channel compression scheme is enough. However, in deep layers ($e.g.$, conv3-3), dynamic pruning is performed according to the semantics of categories. It’s noticed that semantically similar categories, such as “dog” and “cat”, also share similar gate execution patterns. We conclude that FGC utilizes and enhances the correlation between instance semantics and pruning patterns to pursue better representation.

In Fig. 9, we show images in the ordering of gate similarities with the target ones (leftmost). One can see that the semantic similarities and gate similarities of image instances decrease consistently from left to right.

Mutual Information Analysis. In Fig. 10, we validate that FGC improves the mutual information between features and gates to facility distribution alignment, as discussed in section 3.3. For comparison, only partial deep layers are employed with FGC in our pruned models, indicated by “$\ast$”. Specifically, the normalized mutual information (NMI) between features and gates is roughly unchanged for layers not employed with FGC. However, we observe the consistent improvement of NMI between features and gates for layers employed with FGC compared with the pruned models in BAS [7].

We compare the proposed method with state-of-the-art methods, including static and dynamic methods. For a clear and comprehensive comparison, we classify these methods into different types $w.r.t.$ pruned objects. Specifically, most static methods rely on pruning channels (or pruning filters, which is equivalent), while some prune weights in filter kernels or connections between channels. Dynamic methods are classified into three types, layer pruning, spatial pruning, and channel pruning, as discussed in Sec. 2. The performance is evaluated by two metrics, classification errors and average pruning ratios of computation. Higher pruning ratios and lower errors indicate better-pruned models.

CIFAR10. In Table 5, comparisons with SOTA methods to prune ResNet-{20, 32, 56} are shown. FGC significantly outperforms static pruning methods. For example, our dynamic model for ResNet-56 reduces $66.2\%$ computation, $11.2\%$ more than RL-MCTS [46], given a rather similar accuracy ($6.41\%$ $vs$ $6.44\%$). The reason lies in that FGC further reduces the instance-wise redundancy of input samples.

It prevails various types of dynamic pruning techniques. Specifically, FGC significantly outperforms AIG [19] and Blockdrop [2], typical dynamic layer pruning methods that derive dynamic inference by skipping layers. AIG for ResNet-20 achieves a slightly higher error than FGC but with less than half of the computational reduction. BlockDrop for ResNet-32 prunes $54.0\%$ computation with $8.70\%$ error, $1.44\%$ higher than FGC for ResNet-32. Unlike layer pruning, FGC excavates redundancy of each gated layer in a fine-grained manner, maximally reducing irrelevant computation with minor representation degradation. FGC also consistently outperforms dynamic spatial pruning methods (SACT [16] and DynConv [43]), which typically ignore the distribution consistency between instances and pruned architectures, by a large margin. FGC prunes similar computational costs with lower errors than SACT for ResNet-20 ($53.7\%$ reduction and $10.02\%$ error). Compared with DynConv for ResNet-32 ($51.3\%$ reduction and $7.47\%$ error), FGC prunes much more computation with slightly lower error.

We further show that FGC achieves a superior accuracy-computation trade-off than BAS [7] and FBS [5], the very relevant dynamic channel pruning methods with gating modules. For ResNet-56, FGC prunes $3.3\%$ more computation than BAS with approximate classification error. For ResNet-20, FGC prunes $3.1\%$ more computation with lower error than BAS. FGC alleviates the inconsistency between distributions of features and gates in these methods. Our superior results validate the effectiveness of feature regularization toward gating modules.

When extending to WideResNet [36] and MobileNetV2 [37], FGC still achieves consistent improvements, as shown in Table 6. On WRN-28-10, FGC prunes 3.6% more computation while achieving 0.3% less error. FGC could even further reduce redundancy in the lightweight network. On MobileNetV2, FGC prunes 4.5% more computation and achieves 0.2% less error. The results reflect FGC’s generalizability on various network architectures.

CIFAR100. In Table 7, we present the pruning results of ResNet-{20, 32, 56}. For example, after pruning $26.9\%$ computational cost of ResNet-20, FGC even achieves $0.01\%$ less classification error. For ResNet-32 and ResNet-56, FGC outperforms SOTA methods CAC [47] and BAS [7] with significant margins, from perspectives of both test error and computation reduction. Specifically, FGC, compared with BAS, reduces $5.4\%$ and $3.6\%$ more computation of ResNet-32 and ResNet-56 without bringing more errors. Experimental results on CIFAR100 justify the generality of our method to various datasets.

ImageNet. In Table 8, we further compare FGC with methods on the large-scale dataset to prune ResNet-18. One can see that FGC outperforms most static pruning methods. FPGM [1] reduces $41.8\%$ computation with $31.59\%$ top-1 error, while FGC reduces $42.2\%$ computation with $30.67\%$ top-1 error. It is comparable with other dynamic pruning methods. For example, compared with BAS [7], FGC prunes more computation ($48.1\%$ $vs$ $47.1\%$) with lower top-1 error ($31.49\%$ $vs$ $31.66\%$). FGC outperforms AIG [19] and LCCN [15] significantly with much higher computation reduction and lower errors. Experimental results demonstrate our method’s effectiveness in solving the complex representation redundancy on the large-scale image dataset.

The proposed gated models can be generalized to object detection and semantic segmentation tasks.

Object Detection. Experiments are conducted on PASCAL VOC 2007 using the Faster R-CNN method. The mAP (mean average precision) and the average pruning ratio are used to evaluate the accuracy-computation trade-off. The average pruning ratio is calculated based on the backbone. As shown in Table 9, with a slight accuracy drop, FGC brings the detector with more than $50\%$ computation reduction. Note that FGC improves the trade-off compared with vanilla gated models. It achieves $0.48\%$ higher mAP with approximate computation of vanilla gated model on ResNet-50 backbone. And FGC achieves $1.52\%$ higher mAP with more reduction ($52.0\%$ vs $50.6\%$) on ResNet-34. These experiments indicate that our method could better utilize instance-wise redundancy to derive compact sub-networks for each input candidate.

Semantic Segmentation. We conduct semantic segmentation experiments on Cityscapes using PSPNet. We use mean IoU (Intersection over Union), pixel accuracy, and average pruning ratio to evaluate gated models. As shown in Table 10, the baseline model achieves $71.68\%$ IoU and $96.08\%$ pixel accuracy. Our gated model performs better segmentation results with significant computation reduction in comparison. Furthermore, FGC achieves $72.32\%$ IoU and $96.18\%$ pixel accuracy with $34.6\%$ computation reduction, outperforming the vanilla gated model, which achieves $71.70\%$ IoU with an approximate computational cost. It is concluded that our method achieves a better accuracy-computation trade-off.