Exploring Sparsity in Image Super-Resolution for Efficient Inference

1) Training Phase

Spatial Mask. The goal of spatial mask is to identify “important” regions in feature maps (i.e., 0 for “unimportant” regions and 1 for “important” ones). To make the binary spatial mask learnable, we use Gumbel softmax distribution to approximate the one-hot distribution [18]. Specifically, input feature $F\!\in\!\mathbb{R}^{C\times{H}\times{W}}$ is first fed to an hourglass block to produce $F^{spa}\!\in\!\mathbb{R}^{2\times{H}\times{W}}$, as shown in Fig. 5(a). Then, the Gumbel softmax trick is used to obtain a softened spatial mask $M_{k}^{spa}\!\in\!\mathbb{R}^{{H}\times{W}}$:

where $x,y$ are vertical and horizontal indices, $G_{k}^{spa}\!\in\!\mathbb{R}^{2\!\times\!{H}\times\!{W}}$ is a Gumbel noise tensor with all elements following ${\rm Gumbel}(0,1)$ distribution and $\tau$ is a temperature parameter. When $\tau\!\rightarrow\!\infty$, samples from Gumbel softmax distribution become uniform. That is, all elements in $M_{k}^{spa}$ are 0.5. When $\tau\!\rightarrow\!0$, samples from Gumbel softmax distribution become one-hot. That is, $M_{k}^{spa}$ becomes binary. In practice, we start at a high temperature and anneal to a small one to obtain binary spatial masks.

Channel Mask. In addition to spatial masks, channel masks are used to mark redundant channels in those “unimportant” regions (i.e., 0 for redundant channels and 1 for preserved ones). Here, we also use Gumbel softmax trick to produce binary channel masks. For the $l^{\rm th}$ convolutional layer in the $k^{\rm th}$ SMM, we feed auxiliary parameter $S_{k,l}\in\mathbb{R}^{2\times{C}}$ to a Gumbel softmax layer to generate softened channel masks $M_{k,l}^{ch}\in\mathbb{R}^{{C}}$:

where $c$ is the channel index and $G_{k,l}^{ch}\in\mathbb{R}^{2\times{C}}$ is a Gumbel noise tensor. In our experiments, $S_{k,l}$ is initialized using random values drawn from a Gaussian distribution ${\rm N}(0,1)$.

Sparsity Regularization. Based on spatial and channel masks, we define a sparsity term $\eta_{k,l}$:

where $I\!\in\!{\mathbb{R}^{{H}\times{W}}}$ is a tensor with all ones. Note that, $\eta_{k,l}$ represents the ratio of activated locations in the output feature maps. To encourage the output features to be more sparse with fewer locations being activated, we further introduce a sparsity regularization loss:

where $K$ is the number of SMMs and $L$ is the number of sparse mask convolutional layers within each SMM.

Training Strategy. During the training phase, the temperature parameter $\tau$ in Gumbel softmax layers is annealed using the schedule $\tau\!=\!{\rm max}(0.4,~{}~{}1\!-\!\frac{t}{T_{temp}})$, where $t$ is the number of epochs and $T_{temp}$ is empirically set to 500 in our experiments. As $\tau$ gradually decreases, Gumbel softmax distribution is forced to approach an one-hot distribution to produce binary spatial and channel masks.

2) Inference Phase During training, Gumbel softmax distributions are forced to approach one-hot distributions as $\tau$ decreases. Therefore, we replace the Gumbel softmax layers with argmax layers after training to obtain binary spatial and channel masks, as shown in Fig. 5(c).

1) Training Phase

2) Inference Phase During the inference phase, sparse convolution is performed based on the predicted spatial and channel masks, as shown in Fig. 5(d). Take the $l^{\rm th}$ layer in the $k^{\rm th}$ SMM as an example, its kernel is first splitted into four sub-kernels according to $M_{k,l-1}^{ch}$ and $M_{k,l}^{ch}$ to obtain four convolutions. Meanwhile, input feature $F$ is splitted into $F^{D}$ and $F^{S}$ based on $M_{k,l-1}^{ch}$. Then, $F^{D}$ is fed to convolutions ➀ and ➁ to produce $F^{D2D}$ and $F^{D2S}$, while $F^{S}$ is fed to convolutions ➂ and ➃ to produce $F^{S2D}$ and $F^{S2S}$. Note that, $F^{D2D}$ is produced by a vanilla “dense” convolution while $F^{D2S}$, $F^{S2D}$ and $F^{S2S}$ are generated by sparse convolutions with only “important” regions (marked by $M_{k}^{spa}$) being computed. Finally, features obtained from these four branches are summed and concatenated to produce the output feature $F^{out}$. Using sparse mask convolution, computation for redundant channels within those “unimportant” regions can be skipped for efficient inference.

Effectiveness of Sparse Masks. To demonstrate the effectiveness of our sparse masks, we first introduced variant 1 by removing both spatial and channel masks. Then, we developed variants 2 and 3 by adding channel masks and spatial masks, respectively. Comparative results are shown in Table 1. Without spatial and channel masks, all locations and all channels are processed equally. Therefore, variant 1 has a high computational cost. Using channel masks, redundant channels are pruned at all spatial locations. Therefore, variant 2 can be considered as a pruned version of variant 1. Although variant 2 has fewer parameters and FLOPs, it suffers a notable performance drop (33.53 vs. 33.65) since beneficial information in “important” regions of these pruned channels are discarded. With only spatial masks, variant 3 suffers from a conflict between efficiency and performance since redundant computation in channel dimension cannot be well handled. Consequently, its FLOPs is reduced with a performance drop (33.60 vs. 33.65). Using both spatial and channel masks, our SMSR can effectively localize and skip redundant computation at a finer-grained level to reduce FLOPs by $39\%$ while maintaining comparable performance (33.64 vs. 33.65).

Effect of Sparsity. To investigate the effect of sparsity, we retrained our SMSR with large $\lambda_{0}$ to encourage high sparsity. Nvidia RTX 2080Ti, Intel I9-9900K and Kirin 990/810 were used as platforms of GPU, CPU and mobile processor for evaluation. For fair comparison of memory consumption and inference time, all convolutional layers in the backbone of different networks were implemented using im2col [5] based convolutions since different implementation methods (e.g., Winograd [23] and FFT [41]) have different computational costs. Comparative results are presented in Table 2.

Visualization of Sparse Masks. We visualize the sparse masks generated in the first SMM for $\times 2$ SR in Fig. 8. More results are provided in the supplemental material. It can be seen that locations around edges and textures in $M^{spa}$ are considered as “important” ones, which is consistent with our observations in Sec. 3. Moreover, we can see that there are more sparse channels (i.e., green regions in $M^{ch}$) in deep layers than shallow layers. This means that a subset of channels in shallow layers are informative enough for “unimportant” regions and our network progressively focuses more on “important” regions as the depth increases. Overall, our spatial and channel masks work jointly for fine-grained localization of redundant computation.

Learning-based Masks vs. Heuristic Masks. As regions of edges are usually identified as important ones in our spatial masks (Fig. 8), another straightforward choice is to use heuristic masks. KernelGAN [3] follows this idea to identify regions with large gradients as important ones when applying ZSSR [38] and uses a masked loss to focus on these regions. To demonstrate the effectiveness of learning-based masks in our SMSR, we introduced a variant with gradient-induced masks. Specifically, we consider locations with gradients larger than a threshold $\alpha$ as important ones and keep the spatial mask fixed within the network. The performance of this variant is compared to our SMSR in Table 3. Compared to learning-based masks, the variant with gradient-based masks suffers a notable performance drop with comparable sparsity (e.g., 33.52 vs. 33.33/33.30). Further, we can see from Fig. 8 that learning-based masks facilitate our SMSR to achieve better trade-off between SR performance and computational efficiency. Using fixed heuristic masks, it is difficult to obtain fine-grained localization of redundant computation. In contrast, learning-based masks enable our SMSR to accurately localize redundant computation to produce better results.

Quantitative Results. As shown in Table 4, our SMSR outperforms the state-of-the-art methods on most datasets. For example, our SMSR achieves much better performance than CARN for $\times 2$ SR, with the number of parameters and FLOPs being reduced by 38% and 41%, respectively. With a comparable model size, our SMSR performs favorably against FALSR-A and achieves better inference efficiency in terms of FLOPs (131.6G vs. 234.7G). With comparable computational complexity in terms of FLOPs (131.6G vs. 127.7G), our SMSR achieves much higher PSNR values than IDN. Using sparse masks to skip redundant computation, our SMSR reduces $41\%/33\%/27\%$ FLOPs for $\times 2/3/4$ SR while maintaining the state-of-the-art performance. We further show the trade-off between performance, number of parameters and FLOPs in Fig. 1. We can see that our SMSR achieves the best PSNR performance with low computational cost.

Qualitative Results. Figure 9 compares the qualitative results achieved on Urban100. Compared to other methods, our SMSR produces better visual results with fewer artifacts, such as the lattices in $img\_004$ and the stripes on the building in $img\_033$. We further tested our SMSR on a real-world image to demonstrate its effectiveness. As shown in Fig. 10, our SMSR achieves better perceptual quality while other methods suffer notable artifacts.