Gradient pooling distillation network for lightweight single image super-resolution reconstruction

Network architecture

As depicted in Fig. 1, our proposed network architecture comprises three primary components: Shallow Feature Extraction (SFE) module, Nonlinear Mapping module incorporating Multi-Layer Stacked Feature Mixture Distillation module (MLS-FMDM), and lightweight image reconstruction module. Here, the SFE module is used to convert low-resolution images into high-dimensional features; MLS-FMDM refines and distills high-dimensional features; and the lightweight image reconstruction module reconstructs the previously distilled features into a high-resolution image.

In this network, the Shallow Feature Extraction (SFE) module is a 3 × 3 convolutional layer, which aims to convert the input low-resolution image into a higher-dimensional feature representation, denoted as $F_0$, which is essential for subsequent network operations. This transformation is facilitated through a local linking operation. The SFE module can be articulated as follows:

(1) $$F_0=f_{sfe}\left(I_{LR}\right).$$

Here, the function $f_{sfe}(\cdot )$ denotes a 3 × 3 convolutional layer, and $I_{LR}$ signifies the input low-resolution image.

Subsequently, we employ an innovative Multi-Layer Stacked Feature Mixture Distillation module (MLS-FMDM) to further refine and optimize the initial high-dimensional features ( $F_0$), thereby generating deep features suitable for network upsampling operations aimed at high-resolution image reconstruction. The FMDM is architecturally constituted by a Gradient Pooling Distillation module (GPD) and a Channel Attention module (CAM). Here, both GPD and CAM will be introduced later.

Notably, our design incorporates global residual connections throughout the entire network architecture, which can ensure seamless information flow between all layers and accelerate the training process. Moreover, we embed dense residual connections within the internal configuration of the MLS-FMDM, allowing the network to preserve local detail alongside feature extraction. This promotes feature sharing and propagation, enhancing both the performance and efficiency of the network. We denote the proposed MLS-FMDM as follows:

(2) $$F_n=f_{\mathrm{F}\mathrm{M}\mathrm{D}\mathrm{M}}^n\left(f_{\mathrm{F}\mathrm{M}\mathrm{D}\mathrm{M}}^{n-1}\left(\dots f_{\mathrm{F}\mathrm{M}\mathrm{D}\mathrm{M}}^0\left(F_0\right)\dots \right)\right),$$

(3) $$f_{FMDM}^0\left(F_0\right)=f_{CAB}\left(LN(f_{GPD}\left(F_0\right)+F_0)\right)+F_0.$$

In these two equations, $f_{FMDM}(\cdot )$ symbolizes the FMDM layer, $f_{GPD}(\cdot )$ denotes the GPD layer, $LN(\cdot )$ represents layer normalization, $f_{CAB}(\cdot )$ signifies the CAM layer, and $F_n$ denotes the feature map corresponding to the output of the $n$th FMDM, respectively.

Finally, to ensure model efficiency, we employ a streamlined upsampling technique that amalgamates the shallow feature map $F_0$ with the deep feature map $F_n$, subsequently resizing to the desired dimensions of the high-definition (HD) image. This component comprises a 3 × 3 convolutional layer followed by a pixel shuffling layer. We refer to the resultant output $I_{SR}$ as the super-resolution image.

(4) $$I_{SR}=f_{up}\left(F_0+F_n\right).$$

Here, $I_{SR}$ represents the final, reconstructed high-definition (HD) image, while $f_{up}$ denotes the upsampling module employed in the process.

In these whole processes, the low-resolution image stays the same size in the first two modules, and only through the third module the resultant image would be enlarged by a factor of 2, 3, or 4.

Gradient pooling distillation module

In order to enhance the model’s dynamic modeling ability and hierarchical distillation efficiency, especially in the face of complex datasets, the multi-head pooling strategy proposed in SAFM (e.g., Sun et al., 2023) is considered in our approach. This strategy can capture multi-scale information more efficiently and enhance the sensitivity of the model to fine details. At the same time, we also pay attention to the design principle of the hierarchical distillation architecture from RFDB (e.g., Liu, Tang & Wu, 2020), which enhances the feature representation of the model by decomposing and reorganizing the features. Based on the above two lightweight architectural choices, we designed the GPD module.

The GPD architecture is depicted in Fig. 2. This module is predicated on the concept of multi-head hierarchical pooling. It executes feature extraction through a hierarchical pooling approach, and processes the input data at various levels. At each stratum, features are handled independently, to ensure that the model attends to information across different scales. This hierarchical approach not only bolsters computational efficiency but also facilitates the model’s comprehension and learning of the image’s textural structure. The features, once processed, are consolidated through a 1 × 1 convolution to produce a distilled feature space map that encapsulates multi-scale information.

Concretely, we commence by performing a bisection operation on the normalized input feature channels, dividing them into two constituent parts. One portion is preserved intact, while the remaining portion undergoes further feature distillation procedures. The retained portion is subjected to a layer of deep convolution for direct feature manipulation. Conversely, the other portion is independently subjected to a pooling operation, followed by another bisection of the feature channels and repetition of the preceding steps. This iterative process can be outlined as follows:

(5) $$\begin{array}{rl}& F_{d1},F_{c1}=Spli{t}_1\left(\mathrm{L}\mathrm{N}\left(F_i\right)\right),\\ & F_{d2},F_{c2}=Spli{t}_2\left(P_{\frac{1}{k}\downarrow }\left(F_{c1}\right)\right),\\ & F_{d3},F_{c3}=Spli{t}_3\left(P_{\frac{1}{k}\downarrow }\left(F_{c2}\right)\right),\\ & F_{d4}=P_{\frac{1}{k}\downarrow }\left(F_{c3}\right),\\ & {\hat{F}}_1=DW\mathrm{\_}{\mathrm{C}\mathrm{o}\mathrm{n}\mathrm{v}}_{3\times 3}\left(F_{d1}\right),\\ & {\hat{F}}_i=I_{\uparrow }\left(DW\mathrm{\_}{\mathrm{C}\mathrm{o}\mathrm{n}\mathrm{v}}_{3\times 3}\left(F_{di}\right)\right),2\le i\le 4.\end{array}$$

In this context, $LN(\cdot )$ indicates layer normalization, while $Spli{t}_j$ denotes the operation for separating the feature channels at the $jth$ layer. $F_i$ is denoted as the ith layer feature. $F_{dj}$ and $F_{cj}$ correspond to the distillation and coarse features extracted at the $jth$ layer, respectively. The notation $P_{\frac{1}{k}\downarrow }$ indicates a pooling down-sampling operation with a factor of $k\left(k=2,4,8\right)$. Different $k$ corresponds to different GPDN versions. $D{W}_{-}\mathrm{C}\mathrm{o}\mathrm{n}{\mathrm{v}}_{3\times 3}(\cdot )$ refers to a 3 × 3 depthwise convolution, and $I_{\uparrow }(\cdot )$ represents the operation of upsampling the feature size back to that of the original image using linear interpolation.

Subsequently, the feature values yielded by each distillation operation are concatenated, yielding ${\hat{F}}_i\left(1\le i\le 4\right)$ following a layer of 1 × 1 convolution. Thereafter, we employ the GELU activation function to inject nonlinearity into the model, thereby facilitating the estimation of the attention distribution. The comprehensive process can be encapsulated as follows:

(6) $$\begin{array}{c}\hat{F}=\mathrm{C}\mathrm{o}\mathrm{n}{\mathrm{v}}_{1\times 1}\left(\mathrm{C}\mathrm{o}\mathrm{n}\mathrm{c}\mathrm{a}\mathrm{t}\left(\left[{\hat{F}}_1,{\hat{F}}_2,{\hat{F}}_3,{\hat{F}}_4\right]\right)\right)\\ \overline{F}=\varphi \left(\hat{F}\right)\odot F_0.\hfill \end{array}$$

In this framework, $\mathrm{C}\mathrm{o}\mathrm{n}\mathrm{c}\mathrm{a}\mathrm{t}(\cdot )$ denotes the concatenation of features, $\mathrm{C}\mathrm{o}\mathrm{n}{\mathrm{v}}_{1\times 1}(\cdot )$ represents the application of a 1 × 1 convolution, $\varphi (\cdot )$ refers to the employment of an activation function, and $\odot$ signifies element-wise multiplication.

Channel attention module

After being distilled and processed by the GPD module, these features shall enter the CAM module for extraction once again.

The CAM architecture is depicted in Fig. 3. As previously discussed, the fundamental role of the CAM in FMDM is to accentuate the significance of salient pixels across dimensions through meticulous filtering of key features. To augment the contextual integration capabilities of our model, we incorporate the CAM proposed by HAT (e.g., Chen et al., 2023). This CAM consists of a pair of 3 × 3 convolutions, two 1 × 1 convolutions, and a pooling layer. In the CAM’s initial stage, channel compression and expansion are achieved via the consecutive application of two 3 × 3 convolutional layers, followed by an initial screening of the feature channels. Subsequently, a pooling layer is utilized to condense the spatial dimensionality of the feature map, after which two layers of 1 × 1 convolutions are employed to compute channel weights. Ultimately, the sigmoid function is applied to finalize the refinement of the feature map.

Dataset and implementation details

Datasets and metrics. According to previous research (e.g., Zamfir, Conde & Timofte, 2023; Sun et al., 2023; Zhao et al., 2020), we selected 800 HD images in the DIV2K (e.g., Agustsson & Timofte, 2017) dataset and 2,650 HD images in the Flickr2K (e.g., Lim et al., 2017) dataset and merged them as the training dataset. In this process, we perform bicubic interpolation downsampling on these HD images to generate corresponding low-resolution (LR) images. In our tests, we selected four widely recognized benchmark datasets, namely Set5 (e.g., Bevilacqua et al., 2012), Set14 (e.g., Zeyde, Elad & Protter, 2012), BSD100 (e.g., Timofte, De Smet & Van Gool, 2015), and Urban100 (e.g., Huang, Singh & Ahuja, 2015). The general information of all used datasets is shown in Table 1.

To measure the effect of image restoration, we used peak signal-to-noise ratio (PSNR) and structural similarity index (SSIM) as evaluation indexes. Also, learned perceptual image patch similarity (LPIPS) (e.g., Zhang et al., 2018c) was used as a supplementary evaluation index.

PSNR is a quantitative index used to objectively evaluate image quality, which is used to measure the degree of similarity or difference between the original image and the reconstructed image. The basic calculation formula is as follows:

(7) $$MSE=\frac{1}{mn}{\sum }_{i=1}^m{\sum }_{j=1}^n{[I_{ij}-K_{ij}]}^2$$

(8) $$PSNR=10\cdot {\log }_{10}\left({\displaystyle \frac{MA{X}_I^2}{MSE}}\right).$$

Here $I$ means the original image, $K$ is the reconstructed image, $m$ and $n$ represent image height and width, respectively, $I_{ij}$ and $K_{ij}$ are the $\left(i,j\right)$ pixel values of images $I$ and $K$ respectively, $MA{X}_I$ refers to the possible maximum pixels in the image.

Similar to PSNR, SSIM also compares the difference between the processed image and the original image, however, it focuses more on simulating the perception of the image structure by the human eye. It is calculated as follows:

(9) $$SSIM\left(x,y\right)={\displaystyle \frac{\left(2{\mu }_x{\mu }_y+C_1\right)\left(2{\sigma }_{xy}+C_2\right)}{\left({\mu }_x^2+{\mu }_y^2+C_1\right)\left({\sigma }_x^2+{\sigma }_y^2+C_2\right)}.}$$

Among them, the ${\mu }_x,{\mu }_y$ are the average of $x$ and $y$ pixel, ${\sigma }_x^2,{\sigma }_y^2$ are their variance, ${\sigma }_{xy}$ are their covariance. The constants $C_1$ and $C_2$ are positive constants used to avoid instability caused by denominators approaching zero, and they are usually associated with the maximum pixel value of the image and a small constant.

To perform these evaluations, we first convert the image to YcbCr color space and then calculate the PSNR and SSIM in the Y channel.

Regarding the supplementary evaluation index LPIPS, it is a learning-based perceptual image patch similarity indicator, which is mainly used to evaluate the perceptual quality of images. This indicator uses convolutional neural networks (CNN) to perform feature extraction operations on local patches of images and calculate the similarity score between patches.

Detailed rules for implementation. Following the previous work (e.g., Bilecen & Ayazoglu, 2023), which is the winner of NTIRE RTSR Challange Track 2 (x3 SR), the number of GPD modules and the number of feature channels are set as 8 and 32, respectively. Moreover, we use the KneeLR scheduler in Bilecen & Ayazoglu (2023) to dynamically the learning rate. We set the base learning rate to $5\times {10}^{-4}$ and the minimum learning rate is $5\times {10}^{-6}$. The number of model training iterations is 1,000. To further refine the content of the experimental dataset, we adopt the slicing operation. To increase the diversity of the training data, we also used random horizontal flips and 90 rotations. All our experimental operations are performed in the hardware environment equipped with NVIDIA GeForce RTX 3090 GPU, and PyTorch deep learning framework is adopted to complete the experiment programming and model training work.

Model analysis

Distillation pooling rate experiment. The purpose of this experiment is to explore the most suitable distillation pooling rate for the GPD module in the most suitable GPDN network. We set four different pooling rates for comparison, which are two times, four times, eight times of pooling downsampling per layer, and an extreme case, that is, only two pixels downsampling per layer. We use GPDN_2, GPDN_4, GPDN_8, and GPDN_2p to represent the downsampled versions of the network with different multipliers, with the underlining after the underscore indicating the downsampling multipliers.

From the results in Table 2, we know that GPDN_8 and GPDN_2p perform better in Set5 and Set14 datasets. Especially, for GPDN_2p, PSNR index and SSIM index are 0.2527dB and 0.0044 higher than those for GPDN_2 in Set5. However, GPDN_2 has the best performance on the two benchmark datasets of B100 and Urban100, both of which show good performance in PSNR and SSIM indicators. Therefore, we speculate that the increase of the sampling rate will increase the degree of GPDN network to learn the characteristics of the data set, but will affect the generalization ability of GPDN.

Ablation experiments of the main module. To gain a deeper understanding and accurately evaluate the performance of our proposed GPDN model, we conduct an ablation study on the main components of the model. This study is to adopt the GPDN_8 version of the network architecture. This version is a GPDN, which adopts max pooling with 8 * 8 filters and stride 2. And the training is carried out on the untrimmed training dataset. We use the standard dataset with an upsampling factor of ×4 for testing, covering the two datasets Set14 and B100.

To verify the effectiveness of the GPD module, we first remove the GPD module. Based on the results in Table 3 we can observe that the performance of the model is significantly reduced in the absence of the GPD module. Specifically, on the Set14 and B100 datasets, the PSNR value is reduced by 0.21272 and 0.0524 dB, respectively, which highlights the importance of the GPD module in improving the performance. After removing the max pooling distillation step in the GPD module, we notice a noticeable degradation in the performance of super-resolution compared to the baseline model. This shows that the quality of super-resolution can be effectively improved by introducing a pooling distillation step, even without adding any trainable parameters. This is due to the fact that the design of this structure is aimed at reducing the data dimension layer by layer and highlighting the main feature information within the data. Each layer of the network is capable of extracting the most representative portion of the data. As a result, the model can concentrate on these crucial features and disregard some subtle changes and noise. Furthermore, the required Parameters and FLOPs are not increased. This point verifies the effectiveness of the layer-by-layer pooling distillation strategy we adopted.

After replacing the max-pooling distillation in the GDP module with convolutional distillation, we can clearly observe a significant increase in both model parameters and computational load. Although the enhancement in super-resolution quality is evident, the parameter count has escalated by almost four times. We argue that trading nearly quadruple storage space for marginal improvements in super-resolution performance does not offer compelling value. Furthermore, we have increased the number of feature channels in the model, gradually increasing from the original 32 to 128. Experimental results indicate that the super-resolution reconstruction quality of the model does not continuously improve with the increase in the number of feature channels. When the number of channels increased to 72, there was indeed an improvement in super-resolution quality, but the number of parameters also increased significantly. Furthermore, when the number of feature channels was increased to 128, the model exhibited overfitting during the training process, leading to a decrease in super-resolution quality and reduced generalization. Based on these findings, we believe that having 32 feature channels can strike a good balance among the number of parameters, computational load, and super-resolution quality.

Effect of local residual link mode on the network. After the main frame of the network is determined, we further explore the impact of different types of local residual connections on the final super-resolution accuracy. We experimented a total of four kinds of local residual as shown in Fig. 4. In the experiments, we selected four different local residual structures for comparison. They are dense residual connection (DRC), single point residual connection (SPRC), traditional residual connection (TLR), and no residual connection (NRC). These experiments are all performed under the GPDN_8 model version using the original untrimmed dataset.

It can be found from the data in Table 4 that on the Set14 benchmark data set, TLR performs the best. Its PSNR value is 0.3582 dB higher than the worst DRC, 0.1757 dB higher than the second highest SPRC, and 0.0007 higher than the SSIM index. However, in this dataset Validation set for DlV2K, SPRC slightly outperforms TLRC in terms of PSNR and SSIM metrics, although this advantage is not significant. In addition, the performance of NRC on the two datasets Set14 and DIV2K.val does not show any advantage. The results of this experiment also demonstrate the effectiveness of residual connections in the GPDN model.

Training dataset segmentation experiments. To speed up the model training efficiency, we consider the segmentation of the training dataset. In the training of the proposed GPDN model, the cropped versions of the traditional DIV2K and Flickr2K are used. In this version, after combining the above two data sets, the software is used to cut them into several parts in a certain proportion. From the data in Table 5, we can see that training with the cropped training data set will obtain better training results than training with the uncropped training data set. For example, there is a big gap between the results trained on Flickr2k data and Flickt2K_sub data set. The PSNR index on Set14 is 0.1026 dB lower, and the effect of training on DIV2K is also worse than that on DIV2K_sub. The PSNR index on B100 is 0.2661 lower. The experiments fully prove that the quality of super-resolution can be improved by cutting and expanding the dataset without increasing any learnable parameters in the network.

Comparison with state-of-the-art methods

It is obvious from Tables 6 and 7 that the number of learnable parameters of our proposed GPRN is significantly less than that of other SR methods we compared. At the same time, the demand for computing power is also close to or even lower than the algorithms we compared. Among them, our model can save up to 97.66% of the computing power on average. However, in terms of the number of learnable parameters, the other algorithms we compared have on average 5 times more parameters than our model.

It is obvious that the number of learning parameters contained in the GPDN_2 and GPDN_8 is only about one-tenth of that of the Cascading Residual Network (CARN) model (e.g., Ahn, Kang & Sohn, 2018), and its computational requirements are only equivalent to 0.015% of CARN. This is mainly due to the lightweight nature of the GPD module. The module is mainly composed of deep convolution, maximum pooling, and point convolution. The role of deep convolution and maximum pooling is to distill pooling and extract key features, while point convolution is responsible for integrating multi-channel features. Compared with standard convolution, deep convolution can significantly reduce the number of parameters of the model. Due to the reduction in the number of parameters, the computational amount of deep convolution is correspondingly reduced. Compared with the RT4KSR-XXXL (e.g., Zamfir, Conde & Timofte, 2023) algorithm, although the latter is lower in computational resource consumption than GPDN_2 and GPDN_8, it is worth noting that GPDN achieves higher PSNR performance on the two benchmark datasets Set5 and Set14.

It can be inferred from Table 8 that in some specific datasets, the performance of GPDN at the perceptual level is not inferior to that of CARN and SwinIR, such as Set5 and Set14. However, the performance on B100 is not very satisfactory, mainly due to the fact that the image size of B100 after 4× downsampling is too small ( $120\times 80$), and GPDN pursues the balance between extreme lightweight and performance. Images that are too small may not be suitable for GPDN.

Upon examining Fig. 5, it is evident that our model not only possesses the trait of being lightweight but also rivals the performance of esteemed super-resolution models such as CARN and SwinIR in certain aspects of super-resolution capabilities. In the “baby” image, it is observed that the pattern on the baby’s hat reconstructed by our model GPDN is comparable to that reconstructed by CARN, with GPDN’s reconstruction exhibiting less smearing than CARN. In the “ppt3” image, GPDN’s performance is inferior to that of CARN and SwinIR, with GPDN exhibiting a higher noise level and a less pristine output. This could be attributed to the model’s smaller parameter set, which results in a weaker capability to suppress noise. In the “Img016” image, GPDN’s reconstruction performance surpasses that of SwinIR and CARN. Upon close inspection of the cheetah’s belly pattern, the spots reconstructed by SwinIR and CARN show some merging, whereas GPDN’s reconstruction is relatively clearer with fewer merged areas.