Towards optimal sparse CNNs: sparsity-friendly knowledge distillation through feature decoupling

Network sparsity

Network sparsity is an effective technique for reducing parameter counts in deep neural networks, offering significant improvements in computational efficiency and memory usage. By creating sparse models with smaller memory footprints and lower computational demands, this approach enables efficient deployment on resource-constrained devices. Current methods for inducing sparsity include pruning techniques including static pruning, layer-wise sparsity allocation, rand pruning at initialization approaches (Hu et al., 2021; Chen et al., 2022b), which we briefly introduce as follows.

Static sparse training with random pruning (Mariet & Sra, 2016; He, Zhang & Sun, 2017; Suau, Zappella & Apostoloff, 2019; Gale, Elsen & Hooker, 2019) employs layer-wise random masking based on predefined sparsity ratios. Liu et al. (2022) demonstrated that simple random pruning serves as an effective baseline. While uniform pruning applies identical sparsity ratios across all layers, more sophisticated approaches have emerged to enhance sparse model performance. For example, non-uniform and scale-free topologies have shown improved performance compared to dense counterparts when applied to restricted Boltzmann machines (RBMs). Expander graphs have also been used to construct sparse convolutional neural networks (CNNs) that achieve comparable performance to dense CNNs.

While not originally developed for static sparse training, advanced layer-wise sparsity methods like Erase Random (ER) (Mocanu et al., 2018) and Erase Random wrt Kernel (ERK) (Evci et al., 2020) from graph theory have demonstrated strong performance. These approaches differ from traditional methods that pre-specify layer sparsity ratios. Instead of pre-choosing a sparsity ratio for each layer, some approaches utilize saliency criteria to learn layer-wise sparsity ratios before training. This approach, known as pruning at initialization (PaI), selects structurally important connections at initialization based on various criteria. Several efficient criteria have been proposed to improve the performance of non-random pruning at initialization. These include criteria based on gradient flow, synaptic strengths, neural tangent kernel, and iterative approaches. However, recent studies have revealed that existing PaI methods hardly exploit information from the training data and are robust to mask shuffling. In fact, magnitude pruning after training has been shown to learn both layer-wise sparsities and achieve comparable or better performance compared to PaI methods. Several sanity-check experiments have demonstrated that methods like Gradient Signal Preservation (GraSP); (Wang, Zhang & Grosse, 2020), synaptic strengths (SynFlow; (Tanaka et al., 2020)), neural tangent kernel (Liu & Zenke, 2020), and iterative Single-shot Network Pruning based on Connection Sensitivity (SNIP) (de Jorge et al., 2021; Verdenius, Stol & Forré, 2020) hardly utilize information from the training data and are robust to mask shuffling.

Random pruning at initialization with hand-designed layer-wise sparsity ratios has been shown to outperform or achieve similar performance compared to PaI methods. These findings suggest that what pruning at initialization methods discover are the layer-wise sparsities themselves rather than specific weights or values. This highlights a broader challenge inherent to pruning at initialization and emphasizes the need for further exploration and improvement in sparse training techniques.

Knowledge distillation

Knowledge distillation (Hinton, Vinyals & Dean, 2015; Liu et al., 2023; Shao et al., 2023; Shen et al., 2022) enables efficient knowledge transfer from a large teacher network to a compact student network while reducing computational and memory requirements. This technique trains the student network to replicate both the final outputs and intermediate representations of the teacher network through carefully designed loss functions. Current approaches can be classified into three main categories based on knowledge transfer mechanisms: (1) response-based distillation, (2) feature-based distillation, and (3) relationship-based distillation.

Response-based knowledge distillation usually refers to the use of responses from the final output layer in a network to obtain knowledge and migrate it. Feature-based knowledge distillation primarily utilizes the characterization of feature maps in the middle of the teacher’s network to guide the training of the student’s network. The intermediate feature maps of a network contain rich spatial and structural information regarding image content. Therefore, feature distillation methods (Romero et al., 2015; Yim et al., 2017; Zagoruyko & Komodakis, 2017; Chung et al., 2020) are proposed to encourage the student model to mimic the feature representations learned by the teacher model, which shows improved knowledge transfer performance. Yang et al. (2024) introduced a student-centered distillation paradigm inspired by human educational principles, while Huang et al. (2022) specifically addressed feature map distillation for low-resolution recognition in the compressed networks. As the feature maps from different layers of the student and teacher networks usually have different dimensions (e.g., widths, heights, and channels), existing feature distillation methods adopt various transformations to match their dimensions and different distance metrics to measure feature differences. Relationship-based knowledge distillation (Park et al., 2019; Xie et al., 2019; Yang et al., 2024) further explores the relationships between different layers or data samples and utilizes such relationships as knowledge to be migrated, and as the study progressed, researchers found that relationships between features also play a large role in network performance.

Preliminary

We begin by outlining the fundamental preliminaries of knowledge distillation. The principal concept behind knowledge distillation involves the integration of soft targets derived from the teacher network into the overall loss function. This integration facilitates the training of student networks, thereby enhancing their performance via effective knowledge transfer. In classification tasks, the soft target represents the output from the teacher network’s final layer, indicating the probability that the input image is classified under a particular category. These soft targets are mathematically formulated as follows:

(1) $$p(z_i,T)=\frac{\exp \left(\frac{z_i}{T}\right)}{{\sum }_j\exp \left(\frac{z_j}{T}\right)}.$$Here, $z_i$ denotes the logits corresponding to class $i$, and T is the temperature parameter used to modulate the relevance of each soft target.

The soft target encapsulates dark knowledge within the teacher’s network, which can be transferred to the student’s network by incorporating a knowledge migration loss, articulated as:

(2) $$L_{\mathrm{R}}\left(P(z_t,T),P(z_s,T)\right)=KL\left(P(z_t,T),P(z_s,T)\right).$$

Here, KL represents the Kullback-Leibler divergence. By calculating the Kullback-Leibler divergence between the teacher and student outputs, we can facilitate the student network’s approximation of the teacher network’s logarithmic output. Beyond the transfer of knowledge, the student network also incurs a cross-entropy loss relative to the actual labels, culminating in a total loss function expressed as:

(3) $$L=L_{CE}\left(y,P(z_s,T=1)\right)+L_{\mathrm{R}}\left(P(z_t,T),P(z_s,T)\right).$$

Background modeling of sparse networks

Sparse training entails the training of deep neural networks with sparse architectures, offering benefits such as enhanced computational efficiency and diminished memory demands. A significant challenge in sparse training, however, is the inadvertent neglect of critical foreground features, potentially diminishing the model’s overall effectiveness.

More particularly, there is a tendency for the sparse models during training to optimize background modeling capabilities at the expense of salient foreground features. This imbalanced focus may weaken the model’s ability to represent key foreground elements, thereby impairing performance. More critically, due to the absence of such foreground-background bias in dense models, employing the conventional distillation function, i.e., Eq. (3), introduces a bottleneck in the knowledge distillation process, which hinders the sparse student model from fully assimilating the comprehensive knowledge imparted by the dense model.

To mitigate this issue, it is imperative to maintain an equilibrium between sparsity and foreground feature modeling. This balance may be achieved through the implementation of distinct distillation processes tailored for both foreground and background features. By conscientiously distilling and integrating knowledge from both domains, the model can more accurately represent the comprehensive characteristics of the dataset, thereby enhancing its performance.

Sparsity-friendly knowledge distillation

While sparse training provides benefits in computational efficiency and reduced memory requirements, it is imperative to address potential performance declines caused by the inadvertent omission of crucial foreground features. To counter this, we propose sparsity-friendly knowledge distillation (SF-KD), the overall framework of SF-KD is illustrated in Fig. 1, which deploys distinct distillation techniques to strike a balance between sparsity and accurate foreground feature representation, thereby achieving superior performance overall.

In particular, we delineate foreground from background by extracting a central patch from the image, defining the central area as the foreground and treating the peripheral region as the background, foreground/background patches definition examples and feature decoupling schematic in Fig. 2. Such separate distillation of foreground and background features enables us to capture and preserve their distinctive qualities effectively. This approach not only allows for detailed comparison and analysis of these features within and across different classes but also facilitates the learning of decoupled, more informative, and discriminative features.

We extract foreground and background features by applying a spatial mask. The mask matrix $M\in {\mathbb{R}}^{N\times C\times H\times W}$ is defined as:

(4) $$M_{n,c,i,j}=\{\begin{array}{cc}0\hfill & \mathrm{i}\mathrm{f}i<l\mathrm{o}\mathrm{r}i\ge H\text{-}l\mathrm{o}\mathrm{r}j<l\mathrm{o}\mathrm{r}j\ge W\text{-l},\hfill \\ 1\hfill & \mathrm{o}\mathrm{t}\mathrm{h}\mathrm{e}\mathrm{r}\mathrm{w}\mathrm{i}\mathrm{s}\mathrm{e}\hfill \end{array}.$$

Here, $l=H//4$ is an empirically chosen value, which is further analyzed in our ablation studies. Using this mask, we obtain the foreground and background features as follows:

(5) $$\begin{array}{c}fg\mathrm{\_}s={\mathbf{F}}_s\odot M,\hfill \\ fg\mathrm{\_}t={\mathbf{F}}_t\odot M,\hfill \\ bg\mathrm{\_}s={\mathbf{F}}_s-fg\mathrm{\_}s,\hfill \\ bg\mathrm{\_}t={\mathbf{F}}_t-fg\mathrm{\_}t\hfill \end{array}$$where ${\mathbf{F}}_s$, ${\mathbf{F}}_t$ denote the intermediate feature maps of the student model and the teacher model, respectively.

To orchestrate the knowledge distillation process, we introduce a loss function predicated on decoupling of middle layer features, articulated as follows:

(6) $$L_{FKD}=MSE(fg\mathrm{\_}s,fg\mathrm{\_}t)+\eta MSE(bg\mathrm{\_}s,bg\mathrm{\_}t).$$

Here, $fg\mathrm{\_}s$ and $bg\mathrm{\_}s$ denote the foreground and background features of the student network, respectively, while $fg\mathrm{\_}t$ and $bg\mathrm{\_}t$ represent those of the teacher network. The parameter $\eta$, which is less than 1, moderates the focus on background features during distillation.

To enhance model efficacy further, we perform feature distillation on the decoupled features independently. This process involves training the student network to replicate the intermediate representations of the teacher network rather than merely its output. The incorporation of additional loss terms fosters learning from these intermediate representations, thereby improving the model’s performance while preserving its computational efficiency. The overall loss function for the student network is then defined as:

(7) $$L_{total}=\gamma L_{CE}+\alpha L_{KD}+\beta L_{FKD}.$$In this equation, $\gamma$ is a hyperparameter used to balance classification loss and distillation loss and is set to 1 by default; $\alpha$ and $\beta$ serve as hyperparameters that adjust the weighting of the loss components related to response and feature decoupling.

Datasets

(1) CIFAR-10 (Krizhevsky & Hinton, 2009) consists of 60,000 images of 32 $\times$ 32 pixels in color, which are divided into 10 different classes, with each class containing 6,000 images. These classes include: airplanes, automobiles, birds, cats, deer, dogs, frogs, horses, ships, and trucks. The dataset is split into 50,000 training images and 10,000 testing images. (2) CIFAR-100 (Krizhevsky & Hinton, 2009) contains 50 K training images with 0.5 K images per class and 10 K test images. (3) Mini-Imagenet (Deng et al., 2009) a subset of the ImageNet dataset, commonly referred to as Mini-ImageNet, which consists of 100 classes with 600 images per class.

Implementation

On CIFAR-100, we conduct experiments on various teacher-student models under same or different architecture style with Contrastive Representation Distillation’s (CRD’s) settings (Tian, Krishnan & Isola, 2020), whose training epochs are 240. We use a mini-batch size of 64 and a standard Stochastic Gradient Descent (SGD) optimizer with a weight decay of 0.0005. The multi-step learning rate is initialized to 0.05, decayed by 0.1 at 150, 180, and 210 epochs. For the comparison experiments with online KD methods, we adopt the same training settings with Online Knowledge Distillation with Diverse Peers (OKDDip) (Chen et al., 2020), whose training epochs are 300.

Main experiment

Different network structures

Table 3 shows the experimental results for different network architectures, including three backbones: standard resnet, vgg, wide resnet, with the dataset of CIFAR100 and the density uniformly set to 0.9. The experimental results show that the top-1 accuracy rate is improved after knowledge distillation, and the best result is achieved by our feature distillation method SF-KD. Compared to the state-of-the-art (SOTA) feature-based knowledge distillation methods, such as ReviewKD (Pengguang et al., 2021), SimKD (Chen et al., 2022a), and Class Attention Transfer (CAT)-KD (Guo et al., 2023), our method achieves competitive results even in models without sparsity.

The network sparsity experiments are shown in Table 4, the number of parameters of different structural models is halved after reducing the density from 1.0 to 0.5, and the sparsification process effectively reduces the model parameters. The inference time of a single image is reduced, and the sparsity effectively improves the computational efficiency of the model. However, the Top-1 accuracy generally decreases, ranging from −0.24% to −1.66%, indicating that sparsity has different degrees of negative impact on model performance.

Ablation

In the ablation experiments, following Gao et al. (2019), Zhang, Shu & Zhou (2018) we comparatively study the effects of different network densities, foreground v.s. background ratios, and different hyperparameter ratios on the performance of the network.

Visualization of feature map

We visualize the differences between the vanilla model, sparse model and distilled model by visualizing the feature maps, activation heatmaps as shown in Fig. 3. Relative to the dense resnet20 network, the sparsified resnet20+ERK model exhibits a focus on the background, thus weakening its ability to effectively represent foreground components. Our proposed SF-KD approach pulls attention back to foreground representation in sparse networks and outperforms dense vanilla resnet20.