Learning Attention from Attention: Efficient Self-Refinement Transformer for Face Super-Resolution

1 Introduction

Face Super-resolution (FSR) is a specific super-resolution problem which needs to consider the unique textures on the face, such as eyes, mouth, nose, etc. The reconstruction of these structures are critical to distinguishing identity information. In recent years, FSR technology has been widely used and has attracted much attention. With the development of Convolutional Neural Networks (CNN), many researchers have designed various CNN networks [Huang et al., 2017; Zhang et al., 2018a; Zhang et al., 2018c; Shi and Zhao, 2019; Li et al., 2021] to improve FSR performance. Moreover, face priors, such as facial landmarks and heatmaps, were also incorporated into some methods [Chen et al., 2018; Bulat and Tzimiropoulos, 2018] in order to improve global face contour recovery.

∗Corresponding author

The CNN-based network has limited receptive field since convolution is a local operation. Due to the symmetry of the face, we also need to establish long-term dependencies in the FSR task in order to reconstruct the complex details on facial components. Different from CNN, Transformer [Vaswani et al., 2017] proposes a self-attention mechanism to establish global dependencies. In the field of computer vision, the pioneering work ViT [Dosovitskiy et al., 2021] separates the image into patches of equal size and then calculates the selfattention between the patches in the whole image, yielding better results compared to CNN in the classification task. Following this, various vision Transformers [Wang et al., 2021; Chu et al., 2021; Yuan et al., 2021; Liu et al., 2021] were proposed to address different visual tasks. Subsequently, Transformers were also introduced into the SR task. Generally, recent Transformers mainly conduct self-attention in the global searching region or pre-defined local window. However, it may encounter practical difficulties when deals with FSR. For the self-attention performed on global image region, the reconstructed texture details are calculated by the combination of all input tokens, which fails to focus the attention on the most relevant ones. The integration of irrelevant textures could cause blurry artifacts on the reconstructed results. For self-attention conducted on local window, it may fail to produce high-frequency details on complex facial components (e.g., eyes) if flat textures dominate the pre-defined rectangular window.

To solve these problems, we propose an efficient selfrefinement mechanism for Transformer, called Region Selection Attention (RSA), which first produces the coarse attention map for conducting self-attention and then learns finegrained attention map from the coarse one for further refinement. In the coarse-grained self-attention, the attention map is calculated on the down-sampled scale, which is effective for reconstructing flat facial regions and has the advantage of saving computational cost. To enhance the features of key patches in the coarse attention map, we adaptively divide several regions with the strongest attention as refined searching field to calculate fine-grained self-attention. In this way, we could further restore the detailed features to compensate the reconstructed coarse textures on the key regions. The advantages of the proposed RSA can be summarized in two aspects. First, it conducts a content-aware feature reconstruction that treats the coarse structure and detailed texture in different manners. Second, it could explore fine-grained self-attention in the receptive field with irregular shape, which is more robust compared to the traditional rectangular regions.

Another drawback in self-attention is that the operation takes the channel information of tokens equally without considering the importance. To privilege the important channels in feature extraction, we propose Feature Integration Module (FIM) which consists of alternating channel attention module and depth-wise convolution in dual-branch. It could promote a further step to achieve cross-spatial and cross-channel integration simultaneously in Transformer.

Furthermore, we design the Wavelet Fusion Module (WFM) which could modulate the global facial structure information and local detailed texture in frequency-domain. As we know, the shallow layer in the deep network contains the structural information (e.g., facial contour) while the deep layer has the advantage of extracting complex local details (e.g., eyelids). Different from previous methods which conduct simple concatenation or summation in the temporal domain, the proposed WFM achieves reconstruction by exploring the frequency property. We employ Wavelet Transformation (WT) to separate low-frequency and middle-frequency parts from shallow-layer features, while obtaining the highfrequency parts from deep-layer features. The modulation of frequency-specific feature maps is then conducted in WFM for better restoration.

Overall, our main contributions are summarized as follows:

• We propose an efficient self-refinement Transformerbased architecture for FSR task. It could adaptively conduct texture-aware reconstruction in a coarse-to-fine manner.
• The Feature Integration Module (FIM) is employed to consider cross-channel difference in Transformer, which promotes the integration of spatial-wise and channelwise information for feature extraction.
• We design the Wavelet Fusion Module (WFM) to achieve the modulation of shallow-layer and deep-layer features through frequency decomposition and recombination.
• Our method achieves state-of-the-art quantitative metrics and visualizations. It obtains the advantages of more than 0.32dB for PSNR values on the Helen datasets.

2 RelatedWorks

2.1 Face Super-Resolution

In recent years, deep learning has achieved great developments in the field of computer vision. The FSR task has attracted many researchers because of its wide application prospects. GLN [Tuzel et al., 2016] designs the Global Upsampling Network to reconstruct the overall face and the Local Refinement Network to enhance the local details of the face. Attention-FH [Cao et al., 2017] utilizes the context dependencies between facial components to recursively restore the details. RCAN [Zhang et al., 2018c] presents a very deep residual channel attention network, which significantly enhances the learning capacity of CNN. FSRNet [Chen et al., 2018] adds facial geometry priors such as facial landmarks, heatmaps and parsing maps to the network and achieves excellent results. SPARNet [Chen et al., 2020] proposes a network built from Face Attention Units that can efficiently capture key features in very low-resolution face images. DIC [Ma et al., 2020] adopts a novel iterative collaboration network to gradually obtain accurate facial landmarks and super-resolution images. IGAN [Li et al., 2021] considers SR as the information-growth process and recovers HR images by exploring information differences in images of different resolutions. SRDD [Maeda, 2022] proposes a high-resolution (HR) dictionary that can be learned explicitly, which reduces the information that the network needs to process in the HR space. HGSRCNN [Tian et al., 2022] adopts a heterogeneous structure to enhance the internal and external correlations of channels in parallel, which promotes the recovery of images.

2.2 Vision Transformer

The Transformer is originally used for sequence processing in natural language tasks. The proposal of ViT [Dosovitskiy et al., 2021] proves that Transformer can achieve state-of-theart performance in the image classification task. By recursively aggregating tokens of neighboring items, T2T [Yuan et al., 2021] greatly decreases tokens length for practical application. [Wu et al., 2021] builds a Pyramid Pooling Transformer to achieve better performance in various downstream vision tasks such as semantic segmentation and object detection. Swin Transformer [Liu et al., 2021] decomposes the image into non-overlapping windows, calculates multi-head self-attention (MHSA) within the window, and introduces a shifted window mechanism to establish the cross-window connection. Recently, various vision Transformers [Wang et al., 2021; Chu et al., 2021] are proposed to address different visual tasks. Meanwhile, Transformers are also applied in the field of super-resolution. [Liang et al., 2021] builds SwinIR based on Swin Transformer and achieves excellent results in SR tasks. VSR [Cao et al., 2021] turns MHSA into a spatialtemporal convolutional self-attention to achieve state-of-theart performance in video super-resolution tasks. [Shi et al., 2022] proposes a pyramid encoder/decoder Transformer architecture to extract and restore feature textures in different spaces through a hierarchical structure.

3 Method

3.1 Overview

As shown in Figure 1(a), we describe three core stages of the proposed network: Feature Extraction Stage, Feature Transformation Stage, and Feature Recovery Stage. Let I L ∈

R3×H×W be an input low-resolution RGB image, where H and W represent the height and width, respectively. We first expand the channel of input image I L to C as:

where HEF (·) consists of a 3 × 3 convolution and a LeakyReLU activation. The convolutional layer expands the feature to a higher dimensional feature space, which is beneficial for recovering details in different feature channels.

Feature Extraction Stage. We extract features through a hierarchical feature pyramid structure containing 4 steps. Each step consists of 2 Region Selection Attention (RSA) Blocks, a down-sampling operation, a Wavelet Transformation (WT), and a 3×3 convolution. In particular, the 4-th step only contains 2 RSA Blocks. We get the output features F i es of i-th step as:

where HRB(·) denotes 2 consecutive RSA Blocks, and Down(·) is the down-sampling operation. For the downsampling operation, we apply a 4×4 convolutional layer with stride 2 to double the channel number and reduce the height and width to 1/2 of their original size, respectively. In each step, the features are also decomposed into frequency-domain features by Wavelet Transformation (WT). Our proposed WT is composed of a low-pass filter and a high-pass filter. The features sequentially pass through the combination of two filters, which can be converted into frequency-domain features:

LL, LH, HL, and HH (see Sec. 3.3 for details). In the ith step, for the output ˆ F i es of 2 RSA Blocks, we extract the low-frequency and middle-frequency features to present the global facial structure information:

where HWT (·) denotes WT and Conv3×3(·) is a 3 × 3 convolutional layer.

Feature Transformation Stage. At this stage, the output of the previous step F 4 es is fed into 2 RSA Blocks and concatenated with itself as:

where Concat[·] denotes concatenation operation.

Feature Recovery Stage. This stage also contains 4 steps to gradually restore high-resolution images. Each step consists of 2 RSA Blocks, and a Wavelet Fusion Module (WFM). The 1-st step doesn't contain the WFM. WFM can integrate the general facial features extracted from the Feature Extraction Stage and the corresponding high-frequency features in the Feature Recovery Stage through Wavelet Inverse Transformation (WIT).

We will describe the WFM in detail in Sec 3.4. More specifically, the output of j-th step can be formulated as:

where i = 5-j, and HWFM(·) represents the WFM. Finally, the features F 4 rs ∈ R2C×H×W are refined by a 3×3 convolution and a LeakyReLU activation to obtain the residual image I R ∈ R3×H×W . The final recovered image ˆ I can be obtained by ˆ I = I L + I R.

3.2 Efficient Region Search Attention (RSA) Block

As shown in Figure 1(b), we replace the traditional multihead self-attention (MHSA) with our proposed Region Selection Attention (RSA). Given the input as Fs, our proposed RSA Block can be formulated as:

where LN(·) is Layer Normalization and RSA(·) is RSA, which selects multiple regions with rich high-frequency information and computes their self-attention.

RSA is a data-driven approach to adaptively implement texture-aware reconstruction in a coarse-to-fine procedure. Specifically, we get a global coarse attention map via Coarse Self-Attention (CSA). This coarse attention map can reflect the high-frequency information density of each region of the image. Under the guidance of the coarse attention map, we can find the largest k attention values in this feature map, representing the most important k regions in the original features. These regions are extracted from the original features and calculated MHSA to refine the results. Key regions with rich high-frequency information are thus reconstructed by the refinement procedure.

As illustrated in Figure 2, given the input Finput ∈ RC′×H′×W ′ , Finput is first downsampled into a smaller feature map Fcoarse as:

where Downsample(·) uses a convolutional layer with kernel size of λ × λ and stride λ. Then the features are fed into the Feature Integration Module (FIM) to privilege the channel-wise importance. This can be formulated as:

Then we transpose the coarse feature ˆ Fcoarse ∈ RC′× H′ λ × W ′ λ into Fc ∈ Rnc×C′ , where nc = H′ λ × W ′ λ . We feed it into MHSA with M heads as:

where Wq, Wk, Wv, Wo are learnable parameters, and Dh = C ′ /M is the number of dimensions for one head. SA(h) , AM (h) , Q(h) c , K(h) c and V (h) c represent the output of selfattention, attention map, query embedding, key embedding, and value embedding from the h-th attention head, respectively. Then we select the key regions according to the coarse attention map. It is can be defined as:

where AM ∈ RM×nc×nc means coarse attention map and Sum(·) represents accumulation along the penultimate dimension of AM, which turns region-to-region attention into region-to-global attention. Topk(·) means to get the largest k elements and GetIndex(·) denotes getting the indices of elements in feature ˆ Finput. Here we let k = ⌊ √H′W ′ λ ⌋ to balance the computational cost and effect. Thus we get the indices of the key regions in ˆ Finput. Next, we extract these regions to calculate the MHSA and add them back to ˆ Finput by indices. The output of this process can be obtained by:

where SelectByIndex(·) means to select the regions in ˆ Finput according to Index and AddByIndex(·) denotes adding Fregions after MHSA to Finput according to Index. Finally, we upsample Outputc and add it to Outputf , which is then fed into a depth-wise convolutional layer. It could be formulated as:

where Upsample(·) denotes the up-sampling operation, which is a 4 × 4 transposed convolution with stride 2. DWConv(·) means 3 × 3 depth-wise convolution.

Feature Integration Module (FIM). As shown in Figure 3, our proposed FIM consists of depth-wise convolutions, channel attention module, and residual connection. Depth-wise convolution can effectively model the spatial information of features. The channel attention module allows the network to focus on channels which should be paid more attention. We change the combination of the depth-wise convolution and the channel attention for better interaction. Moreover, we utilize the residual connections to merge dual branches and the original input. This design complements the features with rich cross-spatial and cross-channel information.

In FIM, the kernel size of depth-wise convolution is 3 × 3, and Adaptive AvgPool means global average pooling operation.

3.3 Wavelet Transformation (WT) andWavelet Inverse Transformation (WIT)

Our proposed Wavelet Transformation (WT) and Wavelet Inverse Transformation (WIT) consist of a convolutional block and a transposed convolution block, respectively. To decompose features into several frequency components, we adopt a high-efficiency wavelet transformation, namely the Haar wavelet.

where L and H represent low-pass and high-pass filters, respectively. Low-pass filters can capture general information like global contour, and facial structure. In contrast, a highpass filter extracts local details such as texture, eyes, facial components, etc.

With the combination of two filters, we can achieve four kernels of Haar wavelet:

The above four kernels could be utilized to decompose the feature map into frequency-domain components: LL, LH, HL, and HH. Given an arbitrary feature map F , the WT can be defined as:

where ConvLLT (·), ConvLHT (·), ConvHLT (·), ConvHHT (·) represent group convolution with stride 2 and weights LLT , LHT , HLT and HHT , respectively.

We use WIT to integrate the general information (LL, LH, and HL) and the high-frequency details HH to reconstruct the face image. The WIT can be defined as:

where DeconvLLT (·), DeconvLHT (·), DeconvHLT (·), DeconvHHT (·) represent four separate transposed convolutions with the weights as Eq. 14.

The WT operation achieves the modulation of shallowlayer and deep-layer feature maps from the perspective of the frequency domain. We employ skip connections to perform WIT in the Feature Recovery Stage, which can stabilize the generation of detailed information.

3.4 Wavelet FusionModule (WFM)

In this section, we describe our proposed Wavelet Fusion Module (WFM). As shown in Figure 1(c), for the feature Frs, it is split evenly into two features Frs L and Frs R as:

Then, Frs L and Frs R are fed into a Wavelet Inverse Transformation (WIT) and an up-sampling block, respectively, which are defined as:

where HWIT (·) and Upsample(·) separately represents the WIT and a 4 × 4 transposed convolution with stride 2. Deeplayer feature Frs L and the corresponding shallow-layer features LL, LH, HL are modulated to refine the detail in the WIT.

Then we feed the combination of two features into a convolutional layer to get the output of this module, which can be formulated as:

where Concat[·] denotes concatenation operation. Conv1×1(·) is an 1 × 1 convolution, which reduces the channel number of the features to half.

4 Experiments

4.1 Datasets

The CelebA [Liu et al., 2015] and the Helen [Le et al., 2012] are two publicly available face image datasets. First, we use MTCNN [Zhang et al., 2016] to crop the face region. After excluding images with a resolution smaller than 128 × 128, the image is resized to 128×128. So we obtained about 178k images from CelebA, of which 177k images were used as HR training images. Using bicubic interpolation, the HR images are downsampled to 16 × 16 to generate LR images. In the testing phase, we extract the remaining 1000 images from the cropped CelebA dataset and randomly extract 100 images from the cropped Helen dataset. For evaluation, we employ the following metrics: Peak Signal-to-Noise Ratio (PSNR), Structural Similarity (SSIM) [Wang et al., 2004] computed on the Y channel of the image YCbCr space, and Learned Perceptual Image Patch Similarity (LPIPS) [Zhang et al., 2018b].

4.2 Implementation Details

Before the Feature Extraction Stage, we first extend the number of feature channels to 32. In the Feature Extraction Stage, the channel number of feature map is 2i-1 × 32 and the attention head number of RSA Block is 2i-1 in the i-th step. In the Feature Transformation Stage, the number of feature channels is set to 256. The number of attention heads is set to 16. In the Feature Recovery Stage at the j-th step, the channel number of feature map is 25-j ×32 and the attention head number of RSA Block is 24-j . When the height × width of the feature map is less than or equal to 32 × 32, we only employ CSA in the RSA with λ = 1. Otherwise, we apply the full RSA and λ = 4. The AdamW optimizer is used to train our model with β1 = 0.9, β2 = 0.99, and weight decay is set to 0.02. The learning rate is initially set to 4 × 10-4 and dropped by half every 20 epochs. Meanwhile, L1 loss is used for training. The batchsize is set to 32. Our model is implemented in PyTorch and trained for 100 epochs using 2 NVIDIA GeForce RTX 3090 GPUs.

4.3 Comparisons with State-of-the-Art Methods

We compare the proposed FSR method with several state-ofthe-art methods: RCAN [Zhang et al., 2018c], DIC [Ma et al., 2020], SPARNet [Chen et al., 2020], IGAN [Li et al., 2021], SwinIR [Liang et al., 2021], SRDD [Maeda, 2022], and HGSRCNN [Tian et al., 2022]. Our quantitative results for Helen and CelebA dataset are shown in Table 1. The results demonstrate that our method achieves the best PSNR, SSIM, and LPIPS performance on both datasets, outperforming other state-of-the-art approaches by a significant margin. Specifically, for the PSNR value, our method outperforms the best other methods by 0.32dB and 0.31dB on the Helen and CelebA datasets, respectively. Qualitative results are shown in Figure 4. Our method can recover more realistic details in the blurry region, and the reconstructed result is closer to the real image. The pupils and nose of the face are restored more clearly. Other methods tend to produce blurring artifacts in regions with complex textures.

4.4 Ablation Study

Effectiveness of the Region Selection Attention (RSA). The core part of RSA is the region selection strategy. Relying on this design, our RSA can effectively focus on regions with long-range dependencies and rich high-frequency information. To demonstrate the effectiveness of this strategy, we remove the part that selects regions to compute key regions selfattention and only retain the Coarse Self-Attention (CSA). As shown in Table 2, our RSA improves the PSNR value by 0.12dB on the Helen test set. The above experiments demonstrate the effectiveness of the proposed RSA.

Effectiveness of the Feature Integration Module (FIM). Our proposed FIM is used to enhance cross-channel awareness and cross-spatial awareness of features. To verify the effectiveness of FIM, we use learnable position parameters, depth-wise convolution or single-branch FIM to replace full FIM. As can be seen from Table 3, compared to learnable position parameters and depth-wise convolution, our firstbranch-only FIM and second-branch-only FIM improve the PSNR values by 0.04dB and 0.03dB, respectively. The full FIM achieved a significant PSNR value improvement of 0.11dB, which illustrates the superiority of our proposed innovative FIM.

Effectiveness of the Wavelet Fusion Module (WFM). WFM is used in the Feature Recovery Stage to recombine the frequency-domain features. In WFM, we retain half of the features that do not participate in the WIT. Preserved features help stabilize feature recovery. Therefore, we designed two ablation experiments: In the first one, we simply replaced WFMwith a concatenation operation. In the second one, we let all the features in the WFM participate in the WIT without splitting. Table 4 demonstrates that our proposed WFM achieves 0.05dB and 0.04dB improvements, respectively. Experiments prove that the designed WFM is efficient.

5 Conclusion

In this paper, we propose a self-refinement Transformer for FSR. It could conduct the coarse-grained self-attention and further compensate for the details by fine-grained selfattention on key regions. To consider the importance of channel information in Transformer, we also employ the FIM to achieve cross-spatial and cross-channel integration simultaneously. Furthermore, WFM is designed to modulate the shallow and deep feature maps in frequency domain for restoration. Extensive experiments demonstrate the effectiveness of the proposed method.

Acknowledgments

This work was supported by the National Natural Science Foundation of China (Grant No. 62002283, 62311530046, U2003207, 61921004, U21B2048, 62102307, 61972016, 62071380), the Key Research and Development Program of Shaanxi (No.2021GXLH-Z-021) and the Fundamental Research Funds for the Central Universities.