Hierarchical Bilevel Learning with Architecture and Loss Search for Hadamard-based Image Restoration

1 Introduction

Image restoration with deep learning methods has developed rapidly in recent years, and Convolutional Neural Networks (CNNs) have been used widely in many image restoration tasks such as low-light image enhancement (LLIE), single image haze removal (SIHR) and underwater image enhancement (UIE) [Ma et al., 2022; Ma et al., 2021; Qin et al., 2020; Fu and Cao, 2020]. The performance of CNNs largely depends on the network architecture and training loss. In order to obtain networks with high performance, the automatic search of these two components has been explored recently.

∗Corresponding author.

1.1 RelatedWork

Network Architecture Search. Great progress has been made in designing Neural Architecture Search (NAS) [Liu et al., 2021a; Liu et al., 2021e]. Early attempts [Zoph and Le, 2016] used evolutionary algorithms to optimize the architectures and parameters of the networks. Another approach [Tan et al., 2020] used reinforcement learning to train a meta-controller to find architectures. However, using evolutionary algorithms and reinforcement learning required a large amount of computations and made the search process more inefficient. Recently, differentiable search methods [Liu et al., 2018] were proposed. They converted the super-network into a differentiable form by introducing architectural parameters which could be optimized by gradient descent and the search cost could be reduced largely. The above algorithms have achieved competitive performance in various tasks. However, these existing methods tend to obtain unsatisfactory visual results in complicated scenarios. The main reason is that although these methods make use of NAS technique, they lacked physical constraints in principle. As a result, the architectures searched by the aforementioned methods often failed when meeting with complex degraded images because they couldn't capture the inherent image structure without the help of principle prior knowledge.

Training Loss Search. Training loss is a critical component affecting performance. And most existing methods employed hand-craft losses with human experience. ℓ2 loss is a popular training loss in image restoration problem. However, it often brings about blurry in the images. Perceptual loss [Johnson et al., 2016] calculates ℓ2 distance between the activations of the hidden layers of a pre-trained network. However, using this loss could result in a large misalignment between training loss and evaluation metrics. So, many works combine multiple losses with carefully designed trade-off weights for complementation [Lv et al., 2018]. To find a better loss combination or form, training loss search has also aroused great interest very recently [Liu et al., 2021b; Li et al., 2019b; Li et al., 2022]. A cross-entropy based loss space was proposed in face recognition [Li et al., 2019b]. It conducted the search process including a set of cross-entropy losses for face recognition. But the types of candidate losses were too simple and of the same type. And it brought about a large misalignment between training loss and target metrics. The work in [Li et al., 2022] designed an elementary search space which was composed of primitive mathematical operators to accommodate the heterogeneous tasks and evaluation metrics. However, it abandoned the current existing loss evaluation system completely. More importantly, the searched combination result of these operators changed with randomness and the search space was too sparse, which may lead to unpleasing results. By contrast, our loss space is image quality-oriented and more relevant to the evaluation metrics. As a result, it is suitable for a lot of image restoration tasks.

1.2 Our Contributions

To address the above issues, we propose a novel hierarchical bilevel learning scheme with architecture and loss search for Hadamard-based image restoration. A Hadamard-inspired neural unit is firstly built to convert the domain knowledge into the network layer. Then we construct a novel triple-level modeling for architecture and loss search by characterizing the latent connection among loss, architecture and parameters. Considering the difficulty of solving the triple-level model, we define a new hierarchical bilevel learning scheme with an approximate solving process. A series of experiments verify our effectiveness on different tasks and the result is shown in Figure 1. In summary, our contributions can be concluded as

• We develop a new Hadamard-inspired neural unit to integrate task-related domain knowledge into the network design to provide performance support for different Hadamard-based image restoration tasks.
• We supply a general triple-level modeling perspective for joint architecture and loss search. It can accurately depict the latent correspondence between loss, architecture, and parameters.
• To solve the triple-level model, we propose a hierarchical bilevel learning scheme that can alternatively estimate the architecture and loss to progressively generate the desired solutions.

2 The ProposedMethod

In this section, we first construct our Hadamard-inspired neural unit. And then, we establish a new triple-level modeling for learning architecture and training loss. After that, we propose a stage-wise learning strategy for solving the built triplelevel model. Finally, we introduce the architecture search space and image quality-oriented loss search space.

2.1 Hadamard-inspired Neural Unit

A set of image degradations can be represented as clear images with a degradation factor in Hadamard product form. And it can be formulated as y = u ⊗ v, where y denotes the degraded image, u denotes the original image and v denotes the degradation factor. ⊗ represents the element-wise multiplication. And u = y ⊘ v can be derived by performing inverse process, where ⊘ represents the element-wise division. By applying this rule to feature layer, we establish a Hadamard-inspired Neural Unit (HNU) to perform taskrelated domain knowledge. The designed HNU is consists of two feature enhancement blocks and is used for estimating u and v to realize degradation removal progressively. We cascade multiple HNUs and embed them into feature convert modules to form the optimization process of the model. The l-th intermediate HNU can be formulated as

for the l-th block HNUl, the feature F lu obtained from the previous block (the F 1u(l = 1) is the feature generated by FCM) is inputted into two feature enhancement blocks Ru and Rv which will be searched to get F l+1 ˆu and F l+1 ˆv . Then the intermediate enhanced feature F l+1 u is generated by executing the element-wise division ⊘ on F l+1 ˆu andF l+1 ˆv . Themultiple cascaded HNUs framework is built by applying this equation. We design a Feature Convert Module (FCM) for conversion between image and feature domain. The whole framework is generated by embedding them into FCMs. The component of feature convert module also needs to be searched. The search space of the above modules can be found in Section 2.4. The overall framework is shown in Figure 2.

2.2 Triple-Level Modeling

The crucial factors of performance for learning-based methods lie in how to define the architecture and training loss. As described above, we have clearly introduced a new block for Hadamard-based image restoration. Here we provide a general perspective for depicting the overall learning process. The formulation can be written as

where

where β represents the hyper-parameters (including trade-off weights of the losses and inner parameters when calculating a specific loss) that needs to be defined in the loss space B. The variable α represents the desired architecture determined in the defined search space A. The variable ω denotes the network parameters related to α. The loss function LC val represents the classification-driven loss used on validation dataset. The loss LQ val and LQ tr denote the image quality-oriented loss 1 used on validation and training dataset, respectively.

The above triple-level modeling actually builds an explicit relationship among training loss, network architecture, and network parameters. To be concrete, both of the loss and the architecture play a decisive role in the network parameters, i.e., the network parameters ω is the variable for the loss and architecture search. When searching the architecture, the adopted loss function directly dictates the form of the searched architecture, that is, the architecture α becomes the variable for the optimization of training loss.

From above, the selection of the validation loss is the key to the success of searching training loss and should be designed carefully. However, the image quality-oriented losses (such as VGG and LPIPS loss) are employed in our loss space and can't be used as validation loss for search. So which loss to use as the validation loss is the key point. Considering how to evaluate whether an image is clear or not, we find that classification network is able to distinguish the image from its degraded form, which is related to the image restoration tasks. As a result, we employ a classification loss as our validation loss to estimate the performance of the searched training loss more accurately. What's more, by setting the degraded image as the negative sample and setting the clear image as the positive sample, the network can extract abstract representation and high level information of natural images. This information is helpful for generating visually pleasing results. More specifically, for each Hadamard-based image restoration task, we train an image classifier to distinguish whether the image is the clear one or not. And then we use the classification loss as the validation loss to search training loss. The specially trained classifier can capture more task-related information, which could evaluate the searched losses more accurately.

1 The specific forms about it can be found in Section 2.4.

2.3 Hierarchical Bilevel Learning Scheme

Actually, solving the above built triple-level model is a challenging issue because of the complex functional relationship. Here we propose a hierarchical bilevel solution scheme to convert the triple-level model into two bilevel models. Because bilevel optimization is more simple and has been researched many years compared to the triple-level problem [Liu et al., 2022b; Liu et al., 2022a; Liu et al., 2021d; Liu et al., 2021c]. By fixing the variable α when optimizing the β, we can obtain the following solution scheme

where t represents the alternative number. The first equation indicates the optimization process of β. The classification loss LC val is used for validation. α is fixed as its initial value, and the solving process of β is converted into a bilevel problem involving the trainable variable β, ω and the fixed variable α. The second equation indicates the optimization process of α. The LQ val is obtained via the solution procedure of β. So the solving process of α is converted into another bilevel problem involving the trainable variable α and ω. What's more, in the solving process of the two bilevel problems, the gradients of weight parameters and architecture are computed using first-order approximation following standard differential NAS techniques [Liu et al., 2018].

2.4 Search Space of Architecture and Loss

Training Loss Search Space. We design an image qualityoriented losses set as the training loss search space. The candidate losses include ℓ1, ℓ2, color [Wang et al., 2019], SSIM, VGG [Johnson et al., 2016], LPIPS [Zhang et al., 2018], TV [Osher et al., 2005] and Laplacian [Rosenfeld, 1976] loss. For each loss, we define the weighting coefficient as the parameter to be optimized. What's more, the VGG loss has another kind of variable which is the selection of the hidden layers of the VGG network. We divide the layers of the VGG network into four groups in terms of feature scale, and the goal is to select one layer from each group. By employing a new variable for each group, we added the selection of VGG layers into loss search space. The detailed definition of each loss is shown in Table 1.

Architecture Search Space. According to different characteristics, we define search spaces for HNU and FCM respectively. The candidate operations include 1×1, 3×3 and 5×5 Convolution (1-C, 3-C and 5-C), 1×1 and 3×3 Residual Convolution (1-RC and 3-RC),3×3 Dilation Convolution with dilation rate of 2 (3-2-DC), 3×3 Residual Dilation Convolution with dilation rate of 2 (3-2-RDC), 1×1 and 3×3 Dense Convolution (1-DC and 3-DC), and Skip Connection (SC). The division info for HNU and FCM is shown in Figure 2.

3 Experimental Results

3.1 Implementation Details

For LLIE task, we used MIT-Adobe 5K [Bychkovsky et al., 2011] and LOL Dataset [Wei et al., 2018]. We sampled 600 low-light images randomly from the MIT-Adobe 5K dataset for searching and training, and 100 images for testing. As for the LOL Dataset, we randomly sampled 690 image pairs for searching and training, and the remaining 99 pairs for testing. For SIHR task, we used the dataset proposed by [Pan et al., 2020]. And for UIE task, we sampled 712 image pairs from UIEB dataset [Li et al., 2019a] for searching and training and the remaining 178 image pairs for testing. All the experiments were performed on a PC with a single TITAN X GPU. We used ResNet-50 to train our classifier. In searching phase, the maximum epoch was 30, the batch size was 1 and the initial learning rate was 0.0005. The momentum was 0.9 and the weight decay was 0.0003. The initiate weight of each loss in loss search space was 1.0. As for the training phase (with the searched architecture and searched training loss), we set the learning rate as 0.0001 and the batch size was 1.

3.2 Comparison with State-of-the-arts

To evaluate our method roundly, we adopted PSNR, SSIM and LPIPS as our evaluation metrics for three tasks. To fully demonstrate the superiority of our method, we compared our method with recently proposed state-of-the-art approaches on LLIE, SIHR and UIE tasks respectively.

Evaluations on LLIE. The compared methods include MBLLEN [Lv et al., 2018], RetinexNet [Wei et al., 2018], DeepUPE [Wang et al., 2019], KinD [Zhang et al., 2019], EnGAN [Jiang et al., 2021], DRBN [Yang et al., 2020], FIDE [Xu et al., 2020] and RUAS [Liu et al., 2021e]. As shown in Figure 3, the compared methods such as DeepUPE and RUAS either couldn't remove noise or caused color shift. By contrast, our method obtained the highest performance in removing noise and enhancing the image and was superior to others in both the evaluation metrics and visual results.

Evaluations on SIHR. The compared methods contain DAD [Shao et al., 2020], EPDN [Qu et al., 2019], FFA [Qin et al., 2020], GridDehazeNet [Liu et al., 2019], MSBDN [Dong et al., 2020], PSD [Chen et al., 2021] and PhysicsGAN [Pan et al., 2020]. The quantitative result was in Table 3 and the qualitative result was in Figure 4. As shown, most methods failed on the real hazy image. The reason is that these methods can't learn the inherent characteristics of the real hazy image. By comparison, our method was superior to others in both haze removal and detail recovery.

Evaluations on UIE. The compared methods include UDCP [Drews et al., 2016], OCM [Li et al., 2016], TSA [Fu et al., 2017], Fusion [Ancuti et al., 2012], AIO [Uplavikar et al., 2019], WaterNet [Li et al., 2019a], UWCNN [Li et al., 2020], FGAN [Islam et al., 2020] and GLN [Fu and Cao, 2020]. As showed in Table 4 and Figure 5, our method gained the highest scores and could enhance the challenging underwater image well while others failed to recover clear details.

4 Algorithmic Analyses

4.1 The Effect of Training Loss Search

To demonstrate the effect of training loss search. We conducted three groups of contrast experiments on LLIE task. (1): Employ MBLLEN as the basic model and its proposed loss as loss space. Use the experienced loss setting proposed by the author (MBLLEN) and compare it with the searched version (MBLLEN+). (2): Employ UNET as the basic model and our proposed loss space as loss space. Set the weights as 1.0 for each candidate (UNET) and compare it with the searched version (UNET+). (3): Using our searched architectures, set the weights as 1.0 for each candidate (Ours) and compare it with the searched version (Ours+). The result is shown in Figure 6. As showed, the searched version obtained a higher performance on both PSNR and SSIM than another. And our proposed method obtained the highest performance.

4.2 The Effect of Architecture Search

We analyzed the performance among our model and different heuristically-designed architectures on LLIE task and they were trained using the searched loss. The result was shown in Table 5. As can be seen, the performance of the supernet and heuristically designed architectures was either unideal or couldn't reach the searched architectures. Briefly, these architectures may not effective enough. It is because these architectures did not integrate the task cues/principles. By comparison, our searched architecture realized the highest performance in all metrics. In a word, this experiment indicated the necessity of searching the architecture and the superiority of our searched architecture.

4.3 Analyzing the Search Strategy

Actually, the search strategy is a decisive factor for the searched architecture. Here, we made an evaluation on the search strategy. Oursa: search training loss and architecture synchronously using the same classification loss LC val as the upper loss. In this way, architecture (α) and training loss (β) are searched by the same optimal goal at the same time. Ours: divide the triple-level problem into two bilevel problems. Specifically, firstly fix the architecture (α) as the supernet and use the classification loss (LC val ) to search training loss, then use this searched loss (LQ val ) to search the architecture. As shown in Table 5, the second strategy achieved higher performance and gained a large advantage.

4.4 The Searched Results of Different Tasks

To better analyze our triple-level model for different tasks and explore the differences between tasks. We plotted the searched architectures and training loss of LLIE, SIHR and UIE tasks in Figure 7. For simplicity, we just selected one searched layer of the architectures from the HNU to show. It should be noted that the searched weights of TV and Laplacian loss are small to zero and don't show great differences among tasks. To focus on more useful information, we didn't draw them in the heatmap. And for easier observation, we softmax the searched loss weights of each task. As can be seen, the searched architectures were different in tasks. With the help of Hadamard principle, 1-C or 1-D was enough in this layer for LLIE and UIE task, while SC was searched in SIHR tasks. And the dominant loss also differed greatly in tasks, which were ℓ1, VGG and SSIM loss for LLIE, SIHR and UIE respectively. As shown in Figure 8, we showed the feature processed by the searched layer in second group of the VGG network. And present the 34-th, 42-th,51-th and 60-th channel feature to a colormap. As can be seen, the perceptual features required of SIHR and UIE were similar, while that of LLIE focused more on architectural characteristics.

5 Concluding Remarks

We proposed a new framework to combine domain knowledge with a new search strategy to perform architecture and loss search for Hadamard-based image restoration. We first developed a Hadamard-inspired neural unit to explore taskrelated knowledge. And by cascading multiple units into feature convert modules, we can obtain the whole holistic structure of our enhancement network. And we designed an architecture search space and an image quality-oriented loss search space. Then we designed a hierarchical bilevel learning strategy to solve the triple-level problems including the architecture, loss and network parameters. Our experiments were performed on three image restoration tasks and we obtained new state-of-the-art results. In the future, we will apply our built triple-level model technique to different vision tasks.

Acknowledgments

This work is partially supported by the National Natural Science Foundation of China (Nos. 61922019, 61733002 and 62027826), the major key project of Peng Cheng Laboratory (No. PCL2021A12) and the Fundamental Research Funds for the Central Universities.