Auto Seg-Loss: Searching Metric Surrogates for Semantic Segmentation

1 INTRODUCTION . Loss functions are of indispensable components in training deep networks , as they drive the feature learning process for various applications with specific evaluation metrics . However , most metrics , like the commonly used 0-1 classification error , are non-differentiable in their original forms and can not be directly optimized via gradient-based methods . Empirically , the cross-entropy loss serves well as an effective surrogate objective function for a variety of tasks concerning categorization . This phenomenon is especially prevailing in image semantic segmentation , where various evaluation metrics have been designed to address the diverse task focusing on different scenarios . Some metrics measure the accuracy on the whole image , while others focus more on the segmentation boundaries . Although cross-entropy and its variants work well for many metrics , the mis-alignment between network training and evaluation still exist and inevitably leads to performance degradation . Typically , there are two ways for designing metric-specific loss functions in semantic segmentation . The first is to modify the standard cross-entropy loss to meet the target metric ( Ronneberger et al. , 2015 ; Wu et al. , 2016 ) . The other is to design other clever surrogate losses for specific evaluation metrics ( Rahman & Wang , 2016 ; Milletari et al. , 2016 ) . Despite the improvements , these handcrafted losses need expertise and are non-trivial to extend to other evaluation metrics . In contrast to designing loss functions manually , an alternative approach is to find a framework that can design proper loss functions for different evaluation metrics in an automated manner , motivated by recent progress in AutoML ( Zoph & Le , 2017 ; Pham et al. , 2018 ; Liu et al. , 2018 ; Li et al. , 2019 ) . Although automating the design process for loss functions is attractive , it is non-trivial to apply an ∗Equal contribution . †This work is done when Hao Li and Chenxin Tao are interns at SenseTime Research . ‡Corresponding author . AutoML framework to loss functions . Typical AutoML algorithms require a proper search space , in which some search algorithms are conducted . Previous search spaces are either unsuitable for loss design , or too general to be searched efficiently . Recently Li et al . ( 2019 ) and Wang et al . ( 2020 ) proposed search spaces based on existing handcrafted loss functions . And the algorithm searches for the best combination . However , these search spaces are still limited to the variants of cross-entropy loss , and thus do not address the mis-alignment problem well . In this paper , we propose a general framework for searching surrogate losses for mainstream nondifferentiable segmentation metrics . The key idea is that we can build the search space according to the form of evaluation metrics . In this way , the training criteria and evaluation metrics are unified . Meanwhile , the search space is compact enough for efficient search . Specifically , the metrics are first relaxed to the continuous domain by substituting the one-hot prediction and logical operations , which are the non-differentiable parts in most metrics , with their differentiable approximations . Parameterized functions are introduced to approximate the logical operations , ensuring that the loss surfaces are smooth while effective for training . The loss parameterization functions can be of arbitrary families defined on [ 0 , 1 ] . Parameter search is further conducted on the chosen family so as to optimize the network performance on the validation set with the given evaluation metric . Two essential constraints are introduced to regularize the parameter search space . We find that the searched surrogate losses can effectively generalize to different networks and datasets . Extensive experiments on Pascal VOC ( Everingham et al. , 2015 ) and Cityscapes ( Cordts et al. , 2016 ) show our approach delivers accuracy superior than the existing losses specifically designed for individual segmentation metrics with a mild computational overhead . Our contributions can be summarized as follows : 1 ) Our approach is the first general framework of surrogate loss search for mainstream segmentation metrics . 2 ) We propose an effective parameter regularization and parameter search algorithm , which can find loss surrogates optimizing the target metric performance with mild computational overhead . 3 ) The surrogate losses obtained via the proposed searching framework promote our understandings on loss function design and by themselves are novel contributions , because they are different from existing loss functions specifically designed for individual metrics , and are transferable across different datasets and networks . 2 RELATED WORK . Loss function design is an active topic in deep network training ( Ma , 2020 ) . In the area of image semantic segmentation , cross-entropy loss is widely used ( Ronneberger et al. , 2015 ; Chen et al. , 2018 ) . But the cross-entropy loss is designed for optimizing the global accuracy measure ( Rahman & Wang , 2016 ; Patel et al. , 2020 ) , which is not aligned with many other metrics . Numerous studies are conducted to design proper loss functions for the prevalent evaluation metrics . For the mIoU metric , many works ( Ronneberger et al. , 2015 ; Wu et al. , 2016 ) incorporate class frequency to mitigate the class imbalance problem . For the boundary F1 score , the losses at boundary regions are up-weighted ( Caliva et al. , 2019 ; Qin et al. , 2019 ) , so as to deliver more accurate boundaries . These works carefully analyze the property of specific evaluation metrics , and design the loss functions in a fully handcrafted way , which needs expertise . By contrast , we propose a unified framework for deriving parameterized surrogate losses for various evaluation metrics . Wherein , the parameters are searched by reinforcement learning in an automatic way . The networks trained with the searched surrogate losses deliver accuracy on par or even superior than those with the best handcrafted losses . Direct loss optimization for non-differentiable evaluation metrics has long been studied for structural SVM models ( Joachims , 2005 ; Yue et al. , 2007 ; Ranjbar et al. , 2012 ) . However , the gradients w.r.t . features can not be derived from these approaches . Therefore , they can not drive the training of deep networks through back-propagation . Hazan et al . ( 2010 ) proposes to optimize structural SVM with gradient descent , where loss-augmented inference is applied to get the gradients of the expectation of evaluation metrics . Song et al . ( 2016 ) further extends this approach to non-linear models ( e.g. , deep neural networks ) . However , the computational complexity is very high during each step in gradient descent . Although Song et al . ( 2016 ) and Mohapatra et al . ( 2018 ) have designed efficient algorithms for the Average Precision ( AP ) metric , other metrics still need specially designed efficient algorithms . Our method , by contrast , is general for the mainstream segmentation metrics . Thanks to the good generalizability , our method only needs to perform the search process once for a specific metric , and the searched surrogate loss can be directly used henceforth . Applying the searched loss for training networks brings very little additional computational cost . Surrogate loss is introduced to derive loss gradients for the non-differentiable evaluation metrics . There are usually two ways for designing surrogate losses . The first is to handcraft an approximated differentiable metric function . For the IoU measure , Rahman & Wang ( 2016 ) propose to approximate the intersection and union seperately using the softmax probabilities in a differentiable form , and show its effectiveness on binary segmentation tasks . Berman et al . ( 2018 ) further deal with multi-class segmentation problems by extending mIoU from binary inputs to the continuous domain with the convex Lovàsz extension , and their method outperforms standard cross entropy loss in multi-class segmentation tasks . For the F1 measure , dice loss is proposed by Milletari et al . ( 2016 ) as a direct objective by substituting the binary prediction with the softmax probability . In spite of the success , they do not apply for other metrics . The second solution is to train a network to approximate the target metric . Nagendar et al . ( 2018 ) train a network to approximate mIoU . Patel et al . ( 2020 ) design a neural network to learn embeddings for predictions and ground truths for tasks other than segmentation . This line of research focuses on minimizing the approximation error w.r.t . the target metrics . But there is no guarantee that their approximations provide good loss signals for training . These approximated losses are just employed in a post-tuning setup , still relying on cross-entropy pre-trained models . Our method significantly differs in that we search surrogate losses to directly optimize the evaluation metrics in applications . AutoML is a long-pursued target of machine learning ( He et al. , 2019 ) . Recently a sub-field of AutoML , neural architecture search ( NAS ) , has attracted much attention due to its success in automating the process of neural network architecture design ( Zoph & Le , 2017 ; Pham et al. , 2018 ; Liu et al. , 2018 ) . As an essential element , loss function has also raised the interest of researchers to automate its design process . Li et al . ( 2019 ) and Wang et al . ( 2020 ) design search spaces based on existing human-designed loss functions and search for the best combination parameters . There are two issues : a ) the search process outputs whole network models rather than loss functions . For every new network or dataset , the expensive search procedure is conducted again , and b ) the search space are filled with variants of cross-entropy , which can not solve the mis-alignment between cross-entropy loss and many target metrics . By contrast , our method outputs the searched surrogate loss functions of close form with the target metrics , which are transferable between networks and datasets . 3 REVISITING EVALUATION METRICS FOR SEMANTIC SEGMENTATION . Various evaluation metrics are defined for semantic segmentation , to address the diverse task focusing on different scenarios . Most of them are of three typical classes : Acc-based , IoU-based , and F1-score-based . This section revisits the evaluation metrics , under a unified notation set . Table 1 summarizes the mainstream evaluation metrics . The notations are as follows : suppose the validation set is composed of N images , labeled with categories from C classes ( background included ) . Let In , n ∈ { 1 , . . . , N } be the n-th image , and Yn be the corresponding ground-truth segmentation mask . Here Yn = { yn , c , h , w } c , h , w is a one-hot vector , where yn , c , h , w ∈ { 0 , 1 } indicates whether the pixel at spatial location ( h , w ) belongs to the c-th category ( c ∈ { 1 , . . . , C } ) . In evaluation , the ground-truth segmentation mask Yn is compared to the network prediction Ŷn = { ŷn , c , h , w } c , h , w , where ŷn , c , h , w ∈ { 0 , 1 } . ŷn , c , h , w is quantized from the continuous scores produced by the network ( by argmax operation ) . Acc-based metrics . The global accuracy measure ( gAcc ) counts the number of pixels correctly classified . It can be written with logical operator AND as Eq . ( 1 ) . The gAcc metric counts each pixel equally , so the results of the long-tailed categories have little impact on the metric number . The mean accuracy ( mAcc ) metric mitigates this by normalizing within each category as in Eq . ( 2 ) . IoU-based metrics . The evaluation is on set similarity rather than pixel accuracy . The intersectionover-union ( IoU ) score is evaluated between the prediction and the ground-truth mask of each category . The mean IoU ( mIoU ) metric averages the IoU scores of all categories , as in Eq . ( 3 ) . In the variants , the frequency weighted IoU ( FWIoU ) metric weighs each category IoU score by the category pixel number , as in Eq . ( 4 ) . The boudary IoU ( BIoU ) ( Kohli et al. , 2009 ) metric only cares about the segmentation quality around the boundary , so it picks the boundary pixels out in evaluation and ignores the rest pixels . It can be calculated with Eq . ( 5 ) , in which BD ( yn ) denotes the boundary region in map yn . BD ( yn ) is derived by applying XOR operation on the min-pooled ground-truth mask . The stride of the Min-Pooling ( · ) is 1 . F1-score-based metrics . F1-score is a criterion that takes both precision and recall into consideration . A well-known metric of this type is boundary F1-score ( BF1-score ) ( Csurka et al. , 2013 ) , which is widely used for evaluating boundary segmentation accuracy . The computation of precision and recall in BF1-score is as in Eq . ( 6 ) , where BD ( ŷn ) and BD ( yn ) are derived from Eq . ( 5 ) . Max pooling with stride 1 , Max-Pooling ( · ) , is applied on the boundary regions to allow error tolerance .

Summary: This paper aims to directly optimize the metrics of semantic segmentation tasks, such as mIoU, which is different from the most existing methods which minimize the cross-entropy as a proxy. The metrics typically contain one-hot labels and logical operations. In order to directly optimize them, the authors first relax the one-hot label/prediction by Softmax. Then the logical operations applied on the one-hot label are extended by a continuous parameter function which is Monotonical and has the same output as the logical operation with 0/1 input. Finally, the authors describe a reinforcement learning framework to optimize the metrics parameterization (i.e., the outer objective), while the inner objective (i.e., the segmentation network) is trained by standard SGD. The experiments have been performed on Pascal VOC 2012 and Cityscapes datasets, showing the searched loss outperformed the traditional ones such as cross-entropy.