# Next Generation Loss Function for Image Classification

Shakhnaz Akhmedova
Centre for Artificial Intelligence in
Public Health Research
Robert Koch Institute
Berlin, Germany
<EMAIL_ADDRESS>
&Nils Körber
Centre for Artificial Intelligence in
Public Health Research
Robert Koch Institute
Berlin, Germany
<EMAIL_ADDRESS>
0000-0003-2927-1974

###### Abstract

Neural networks are trained by minimizing a loss function that defines the
discrepancy between the predicted model output and the target value. The
selection of the loss function is crucial to achieve task-specific behaviour
and highly influences the capability of the model. A variety of loss functions
have been proposed for a wide range of tasks affecting training and model
performance. For classification tasks, the cross entropy is the de-facto
standard and usually the first choice. Here, we try to experimentally
challenge the well-known loss functions, including cross entropy (CE) loss, by
utilizing the genetic programming (GP) approach, a population-based
evolutionary algorithm. GP constructs loss functions from a set of operators
and leaf nodes and these functions are repeatedly recombined and mutated to
find an optimal structure. Experiments were carried out on different small-
sized datasets CIFAR-10, CIFAR-100 and Fashion-MNIST using an Inception model.
The 5 best functions found were evaluated for different model architectures on
a set of standard datasets ranging from 2 to 102 classes and very different
sizes. One function, denoted as Next Generation Loss (NGL), clearly stood out
showing same or better performance for all tested datasets compared to CE. To
evaluate the NGL function on a large-scale dataset, we tested its performance
on the Imagenet-1k dataset where it showed improved top-1 accuracy compared to
models trained with identical settings and other losses. Finally, the NGL was
trained on a segmentation downstream task for Pascal VOC 2012 and COCO-
Stuff164k datasets improving the underlying model performance.

_Keywords_ Loss function $\cdot$ Genetic Programming $\cdot$ Deep Learning
$\cdot$ Classification $\cdot$ Segmentation.

##  1 Introduction

Figure 1: NGL and CE functions (left) and their gradients (right) for
$y_{real}=1$.

Deep learning (DL) has received extensive research interest in developing new
image processing algorithms, and resulting models have been remarkably
successful in a variety of image analysis tasks such as classification
Krizhevsky et al. (2012), segmentation Ronneberger et al. (2015), object
detection Girshick (2015) and more. Deep learning methods represented by (or
partially based on) convolutional neural networks have become a well-known and
constantly developing research topic in the field of image analysis. Various
deep learning models, for example ResNet He et al. (2015), Inception Szegedy
et al. (2015), U-Net Ronneberger et al. (2015), DeepLab Chen et al. (2016),
have been proposed over the years for different image processing tasks and
domains. Vision transformers Dosovitskiy et al. (2020) constantly gain
popularity among researchers due to their high training efficiency and
scalability and are already outperforming convolutional neural networks on
multiple benchmarks for classification, segmentation and image generation
problems.

The loss function plays a crucial role in deep learning in general, as it
shows the direction for model training and affects the performance of the
model and training speed. There are well-known losses used for most image
processing tasks, such as cross entropy (CE) Cover and Thomas (2006), Focal
Loss Lin et al. (2017a), or mean square error Willmott and Matsuura (2005).
Recently, researchers started proposing loss functions developed for specific
applications or tasks. Usually, that is done by hand and requires a lot of
expert knowledge, time, parameter adjustments, and thus computational
resources, due to the complexity of modern models.

The loss function design can be represented as a symbolic regression problem
that can be solved by using Genetic Programming (GP) Augusto and Barbosa
(2000). The GP approach is frequently used to design previously unknown
heuristics in order to solve hard computational search problems K.Burke et al.
(2009). Genetic programming is a domain-independent method, which was used in
this study to genetically breed a population of loss functions for deep
learning models to solve image classification problems. GP was applied to find
a generally applicable loss function for DL classifiers independent of the
dataset and model architecture. Several experiments with different objective
functions using InceptionV3 Szegedy et al. (2015) trained on the CIFAR-10
Krizhevsky et al. (2009), Fashion-MNIST Xiao et al. (2017) and CIFAR-100
Krizhevsky et al. (2009) datasets were conducted to find an optimized loss
function.

The loss functions found were evaluated using ResNet50 and InceptionV3 models
trained on CIFAR-10 Krizhevsky et al. (2009), CIFAR-100 Krizhevsky et al.
(2009), Fashion-MNIST Xiao et al. (2017), Malaria Rajaraman et al. (2018),
PCam Veeling et al. (2018), Colorectal Histology Kather et al. (2016) and
Caltech 101 Fei-Fei et al. (2004) datasets. In the next step, the best
function, denoted as NGL, was used to train ResNet He et al. (2015) models and
Swin transformers Liu et al. (2021) to classify the ImageNet-1k dataset. In
the end, it was applied to train U-Net and DeepLab models on Pascal VOC 2012
Everingham et al.  and COCO-Stuff164k Caesar et al. (2018) datasets to show
the applicability for segmentation.

Among the evaluated functions, NGL demonstrated exceptionally good results
outperforming baseline losses, such as cross entropy loss, focal loss,
symmetric cross entropy and dice loss. Moreover, further investigation has
showed that the reason for these good results may be self-regulation, which is
inherently present in the NGL due to its mathematical definition.

Thus, the main contributions in this paper can be summarized as follows:

  * •

Experimental results show that GP is capable to design a good loss function
for deep neural networks without any prior knowledge neither about the model
nor about the problem at hand;

  * •

NGL, a new highly efficient loss function, applicable for any DL model to
solve classification and segmentation problems, was found;

  * •

Self-regularization naturally implemented in the NGL indicates a direction for
loss function design in general;

  * •

Models trained by the NGL achieved same or better performance on a variety of
datasets including ImageNet-1k and COCO-Stuff164k datasets.

The rest of this paper is organized as follows. Section 2 briefly reviews
related work of deep neural networks used for image analysis and loss function
design. The Genetic Programming algorithm and proposed search procedure are
described in Section 3, while the experimental settings and results are
presented in Section 4. Obtained results are discussed in Section 5. Finally,
the paper is concluded in Section 6.

##  2 Related Work

The progress in solving image processing and visual recognition tasks is
related to the rapid development of convolutional neural networks (CNNs) and
subsequently vision transformers, which have been already used for real-world
applications in different scientific fields. A plethora of convolutional
neural network architectures have been proposed, including the widely used
Inception network Szegedy et al. (2014) that explores the problem of multi-
scale fusion in convolution calculations to better characterize image
information, and ResNet He et al. (2015) with residual blocks to solve the
vanishing gradients problem.

Dosovitskiy et al. Dosovitskiy et al. (2020) proposed the Vision Transformer
(ViT) by formulating image classification as a sequence prediction task of the
image patch sequence similarly to Vaswani et al. (2017), thereby capturing
long-term dependencies within the input image. Transformer-based models for
image processing have been rapidly developed, and the following models are
considered as some of the most representative: Vision Transformers (ViT)
Dosovitskiy et al. (2020), Data-efficient image Transformers (DeiT) Touvron et
al. (2020) and Swin-Transformer Liu et al. (2021).

Most deep learning models initially designed to solve classification problems
can be applied also for segmentation tasks, as segmentation can be formulated
as a classification problem of pixels with semantic labels (semantic
segmentation) or partitioning of individual objects (instance segmentation).
In the recent decade there were a lot of works presenting new and more
powerful segmentation methods, U-Net Ronneberger et al. (2015), which was
initially developed for biomedical image segmentation, is one of the most
widely used architectures not only in the life sciences. Additionally,
DeepLabv1 Chen et al. (2016), DeepLabv2 Chen et al. (2016) and DeepLabv3+ Chen
et al. (2018) are among some of the most popular image segmentation
approaches, which use dilated convolution to address the decreasing resolution
in the network (caused by max-pooling and striding) and Atrous Spatial Pyramid
Pooling (ASPP), that probes an incoming convolutional feature layer with
filters at multiple sampling rates.

Aside from developing deeper networks with more complex structures and
features to get better performance, better loss functions have also been
proven to be effective on improving the model performance in most recent works
Tian et al. (2022), Jadon (2020). Many loss functions have been utilized for
neural networks based on softmax activation function, such as Mean Square
Error (MSE) loss function Willmott and Matsuura (2005), Cross Entropy (CE)
loss function Cover and Thomas (2006) (Binary Cross Entropy loss function when
images are divided into two classes, Categorical Cross Entropy for multiple
classes), etc. In most state-of-the-art supervised learning problems,
practitioners typically use large capacity deep neural networks together with
cross-entropy loss. The latter can be explained by the fact that the
traditional CE loss is supported by clear theory, easy training and good
performance. Nowadays, there are many variants of CE that have been proposed
in the past few years, e.g. symmetric CE (SCE) Han and Wu (2017), or Focal
Loss Lin et al. (2017a).

Loss functions used for segmentation can be divided into four groups described
in Jadon (2020): distribution-based, region-based, boundary-based, and
compounded losses. The most well-known and commonly used distribution-based
loss functions used to train segmentation models are binary cross entropy Yi-
de et al. (2004) and focal loss Lin et al. (2017b). Region-based loss
functions such as dice loss Sudre et al. (2017), and Tversky loss Salehi et
al. (2017) calculate the similarity between images. Hausdorff distance loss
Ribera et al. (2018) is defined as boundary-based loss functions in Jadon
(2020), while the exponential logarithmic loss Wong et al. (2018) is
considered as the compounded loss.

Most of the mentioned loss functions can be considered as heuristic methods
designed by experts by using theory and the domain knowledge for specific
tasks and/or models. Moreover, loss function design requires great effort from
experts to explore the large search space, which is usually sub-optimal in
practice. Nevertheless, only several commonly used losses are usually
implemented for model training, for example, categorical cross entropy for
classification or dice loss for segmentation. Recently, the automated search
of suitable loss functions without domain knowledge has received much
attention from the computer vision community. Reinforcement learning
algorithms were used to learn better loss functions with good generalization
ability on different image analysis tasks Xu et al. (2018), Oh et al. (2020).
However, loss functions found by using these approaches either have task-
specific requirements, such as environment interaction in reinforcement
learning, or remain fixed after training.

A lot of works on automatic loss search follow the search algorithms used in
AutoML He et al. (2019). In Li et al. (2021) the authors explore the
possibility of searching loss functions automatically from scratch for generic
tasks, e.g., semantic and instance segmentation, object detection, and pose
estimation. In studies such as Wang et al. (2020) and Li et al. (2019) the
main focus is on the search for particular hyper-parameters within the fixed
loss formula. AutoML-Zero Real et al. (2020) proposes a framework to construct
machine learning algorithms from simple mathematical operations, but the
search space and search strategies are specialized, which limits its potential
application. The authors of the study presented in Gonzalez and Miikkulainen
(2019) propose a framework for searching classification loss functions by
using evolutionary algorithms but the searched loss poorly generalizes to
large-scale datasets. These works motivated us to design loss functions by
using the Genetic Programming algorithm, which uses a large set of primitive
mathematical operations and terminal values, capable of training different
models regardless of their structure or the task at hand.

##  3 Method

Genetic Programming (GP) is a population-based evolutionary optimization
algorithm, which can be applied to design different heuristics, depending on
the problem at hand, and, additionally, it is considered as a machine learning
tool, as it can be used to discover a functional relationship between features
in data. GP was initially inspired by the biological evolution, including
natural processes such as mutation and selection. In GP, each solution is
represented by a tree structure, and each of these trees are evaluated
recursively to produce the resulting multivariate expression. There are two
types of nodes used for tree-based GP: the terminal node, also called leaf,
which is randomly chosen from the set of variables, and the tree node, which
can be chosen from a predefined set of operators. An example of the generated
trees is given in Appendix A (Figure A). The GP search process starts with
random initialization of a set of potential solutions, which is also called a
population of individuals, in the functional search space. The number of
individuals is the predefined parameter $n$, which does not change through the
whole search process. The overall GP search process is summarized in Algorithm
1. Thus, for this study a set of loss functions was randomly generated, and
each loss function was represented as a tree, where the terminals were
randomly chosen from the set $\\{y_{pred},y_{real},\mathbb{R}\\}$, while the
operators were chosen from the set
$\\{+,-,\times,\div,\times(-1),\sqrt{,}\log,\exp,\sin,\cos\\}$. To utilize
these functions and terminals for loss function generation, several
modifications were applied (mentioned modifications are described in Appendix
A).

Algorithm 1 An overview of the implemented GP approach

Randomly initialize $n$ trees (generation $0$), each representing one formula,
e.g., loss function

Evaluate GP fitness function $F$ for each individual in the population

Determine the best individual

Create an empty external archive $A$

while generation number is less than $T$ do

for each individual from population do

Generate a child individual by applying the crossover operator

if $rand_{1}<M_{ST}$ then

Apply subtree mutation to the generated child

end if

if $rand_{2}<M_{N}$ then

Apply one-point mutation to the generated child

end if

end for

Evaluate GP fitness function $F$ for each generated child individual

Update the best individual

Create new population of size $2\times n$ by concatenating children and parent
individuals

Select $n$ best individuals from the new population: create new generation

Update the external archive $A$

end while

After initialization, the main search loop starts, which consists of the
iteratively repeated steps crossover, mutation, fitness function evaluation
and selection (the number of GP steps or generations is denoted as $T$). The
fitness function $F$ determines the quality of the individual and is the
crucial part of the optimization process, as our goal was to experimentally
find a robust loss function that can be generalized to a variety of datasets.

The crossover operator is used to exchange the subtrees between two
individuals. An example of the implemented crossover operator is demonstrated
in Appendix A (Figure A). Mutation can be applied in various ways, but in this
study two variants were used:

  * •

a random subtree in the tree is chosen and replaced with another randomly
generated subtree (subtee mutation);

  * •

a random node in the tree is chosen and replaced with another randomly
generated node (one-point mutation).

Both mutation operators as well as crossover have their own parameters, such
as subtree and node mutation $M_{ST}$, $M_{N}$ and crossover $Cr$ rates, which
determine how often individuals will be changed during the main search loop.
Finally, all individuals and the ones generated after crossover and mutation
steps are combined into one population, and then the selection operator is
applied. In this study only $n$ most fit individuals from the combined
population are chosen for the next generation on each GP step.

Moreover, the success-history-based adaptation strategy was used to improve
the efficiency of the GP approach. To be more specific, at the beginning of
the main search loop an empty external archive $A$ was created. The maximum
size of this archive was set equal to the population size, $n_{A}=n$. All
individuals not passing selection could have been saved in the archive with a
given probability $p_{A}$. Individuals saved in the external set were used
during crossover step with some probability $Cr_{A}$.

##  4 Experiments

###  4.1 Loss function search

As mentioned in the previous section, loss function search was performed by
the GP algorithm and essentially can be described as an optimization process.
A set of individuals, where each individual is a mathematical formula
representing a loss function, is generated. These individuals change by
crossover and mutation operators and then a new set of loss functions is
selected. These actions are repeated a given number of times, called
generations, and in the end the best individual or loss function is
determined. In this implementation, the crucial part is selection, which in
turn depends on how the fitness of the loss function is defined, since only
the fittest individuals are transferred over to the next generation.

In this study, five experiments were conducted to search the loss function by
GP, which only differed in how the fitness function $F$ was defined.
Specifically, the following definitions of $F$ were used for each individual
loss for each experiment:

  * •

Train the InceptionV3 model from scratch one time on CIFAR-10 dataset, the
validation error was used as the fitness value;

  * •

Train the InceptionV3 model from scratch 3 times on Fashion-MNIST dataset, the
averaged validation error was used as the fitness value;

  * •

Fine-tune the pre-trained InceptionV3 model one time on CIFAR-10 dataset, the
validation error was used as the fitness value;

  * •

Fine-tune the pre-trained InceptionV3 model 3 times on Fashion-MNIST dataset,
the averaged validation error was used as the fitness value;

  * •

Train the InceptionV3 model from scratch on CIFAR-10, Fashion-MNIST and
CIFAR-100 datasets (once each). The validation error for each dataset was
compared to the respective validation error obtained by the same model trained
using CE loss, and individual fitness was represented by the pair of numbers
indicating the number of wins and the percentage of improvement compared to
CE.

It should be noted that a loss function was discarded if its values were not
in the range $[10^{-5},10^{5}]$ and a new one was generated instead of it. The
fitness function evaluation procedure is described in details in Appendix B.1
(Algorithm B.1).

The model and datasets were selected to limit the training time required per
generation while maintaining a reasonable evaluation of the selected loss
functions. Regardless of the experiment, a batch size of 128, data
augmentation including horizontal flip, width and height shift ($0.1$ for
both), zoom ($0.2$), Adam optimizer Kingma and Ba (2015) with reduction of the
learning rate on plateau, and 50 epochs for CIFAR-10 and CIFAR-100 and $35$
for Fashion-MNIST were used to train the network. The top layer of the
InceptionV3 model was replaced by $9$ new layers, consisting of Flatten,
BatchNormalization ($\times 3$) Ioffe and Szegedy (2015), Dense ($\times 3$),
Dropout ($\times 2$) Srivastava et al. (2014), and a final Softmax layer.

The following parameters were used for GP regardless of the experiment:
$n=16$, $T=100$, $M_{ST}=0.3$, $M_{N}=0.1$, $Cr=0.7$, $p_{A}=Cr_{A}=0.5$,
minimum tree height and maximum tree size were set to 2 and 100, respectively.

The best performing functions for each of the five GP experiments were
evaluated for different models and multiple datasets, all functions are listed
in Appendix B.1.

###  4.2 Evaluation

####  4.2.1 Small datasets for classification

The evaluation of the five found functions was performed by training ResNet50
and InceptionV3 on seven datasets, which differed by the number of images,
classes, by the type of images (grayscale and RGB), and their sizes. The brief
description of datasets used in this study is given in Table 1 ($N$ is the
number of classes). ResNet50 was added to the initial experiments, as all the
found functions were evaluated on InceptionV3 model during the GP search to
rule out model specific properties. The top layers of both InceptionV3 and
ResNet50 networks were not included, in both cases instead $9$ new layers such
as Flatten, BatchNormalization ($\times 3$), Dense ($\times 3$) and Dropout
($\times 2$) were added.

Table 1: Brief description of the small datasets used for classification.

Dataset | $N$ | Dataset size | Class size | Image size | Image type
---|---|---|---|---|---
Malaria | $2$ | $27558$ | $13779$ | edge lengths of $40-400$ pixels | RGB
PCam | $2$ | $327680$ | $163840$ | $96\times 96$ | RGB
Colorectal Histology | $8$ | $5\times 10^{3}$ | $625$ | $150\times 150$ | RGB
CIFAR-10 | $10$ | $60\times 10^{3}$ | $6\times 10^{3}$ | $32\times 32$ | RGB
Fashion-MNIST | $10$ | $70\times 10^{3}$ | $7\times 10^{3}$ | $28\times 28$ | Grayscale
CIFAR-100 | $100$ | $60\times 10^{3}$ | $6\times 10^{2}$ | $32\times 32$ | RGB
Caltech 101 | $102$ | $9144$ | $40-800$ | edge lengths of $200-300$ pixels | RGB

Parameter settings and data preprocessing varied for all datasets but were
identical for all tested loss functions, specifics are shown in Appendix B.2
(Table B.2.1). Classification accuracy on the test set after the last epoch
was used as the evaluation metric.

Both InceptionV3 and ResNet50 models were trained five times for each dataset.
Obtained results were averaged for both models and compared to the averaged CE
(with and without $L2$ regularization), focal and SCE losses for the same
models. In the end, the percentage of accuracy improvement or decrease
compared to CE was determined. A final value from the range $[-0.5,0.5]$ was
considered as same performing, a value $>+0.5\%$ better performing and
$<-0.5\%$ performing worse compared to CE.

Results obtained for all evaluated loss functions used to train models from
scratch ($f_{1},...,f_{5}$) are given in Table 2, results obtained for pre-
trained models are presented in Appendix B.2 (Table B.2.2). Loss function
$f_{3}$ demonstrates the worst results compared to others. It showed
particularly bad results on Malaria, Colorectal Histology, CIFAR-100, and
Caltech 101 datasets, it should be noted that the last three mentioned
datasets have less than $10^{3}$ examples of each present class.
Interestingly, no function outperformed CE loss on Malaria, CIFAR-10, and
Fashion-MNIST, while 6 out of 9 functions showed significantly better results
on PCam dataset.

Table 2: Average accuracy (%) for ResNet50 and InceptionV3 each trained 5
times for each dataset for cross entropy (CE) and percentual change of the
tested loss functions compared to CE, best result per dataset is shown in
bold.

Loss | Malaria | PCam | Colorectal Histology | CIFAR-10 | Fashion-MNIST | CIFAR-100 | Caltech 101 | Mean
---|---|---|---|---|---|---|---|---
CE | 94.0 | 69.4 | 88.9 | 92.8 | 94.0 | 68.2 | 72.5 | $\pm 0$
SCE | -0.06 | +0.62 | -0.45 | -3.66 | -2.42 | -0.80 | +2.71 | -0.58
Focal | +0.34 | +1.03 | \+ 2.89 | -0.64 | -0.02 | -2.14 | -0.78 | +0.10
CE + $L2$ | +0.11 | -0.64 | +0.88 | +0.32 | \+ 0.09 | +0.38 | +3.67 | +0.69
$f_{1}$ | -0.27 | +0.81 | +3.27 | -0.34 | +0.31 | -0.84 | -3.67 | -0.10
$f_{2}$ | -0.09 | +2.89 | -7.56 | -0.81 | +0.04 | +0.60 | -3.61 | -1.22
$f_{3}$ | -27.98 | +3.24 | -62.45 | -1.79 | -0.48 | -59.36 | -41.67 | -27.21
$f_{4}$ | -0.21 | -0.13 | -1.98 | -0.14 | 0.00 | -0.32 | -6.49 | -1.32
$f_{5}(NGL)$ | -0.27 | +7.07 | +1.77 | +0.12 | +0.07 | +1.01 | +5.00 | +2.11

Overall, $f_{5}$ is the only function showing on average significantly better
results than CE, outperforming it on four out of seven datasets with an
improvement ranging from 1 to 7 $\%$, while there was no significant
difference on the remaining three datasets.

Given the overall good performance of the function found by an evolutionary
method, we renamed function $f_{5}$ as Next Generation Loss (NGL) function,
with the following formula:

$f_{NGL}=\frac{1}{N}\sum_{i=1}^{N}{\left[e^{(\alpha-y^{(i)}_{pred}\cdot(1+y^{(i)}_{real}))}-\cos(\cos(\sin(y^{(i)}_{pred})))\right]},$
(1)

where $\alpha=2.4092$, $N$ is the number of classes.

To further test the generalizability of NGL, we tested its performance on
ImageNet dataset and larger models.

####  4.2.2 ImageNet-1k

Experiments on ImageNet-1k dataset were performed using the ResNet He et al.
(2015) convolutional neural networks and Swin Liu et al. (2021) transformer
models using NGL, CE, SCE and focal losses with the same settings.

For ResNet, we used the AdamW optimizer Loshchilov and Hutter (2019), cosine
annealing learning rate scheduler with a $30$-epoch linear warm-up, while the
learning rate was equal to $0.003$. The number of epochs was set to $180$, a
batch size of $256$, and a weight decay of $0.01$ were used. For Swin-
Transformers, the AdamW optimizer was employed for $400$ epochs using a cosine
decay learning rate scheduler and $60$ epochs of linear warm-up. A batch size
of $1024$, an initial learning rate of $0.001$, and a weight decay of $0.05$
were used. All of the augmentation and regularization strategies were the same
as in Liu et al. (2021). Results on ImageNet-1k are shown in Table 3. It
should be noted, that for both ResNet and Swin models CE was applied with
label smoothing and $L2$ regularization.

Table 3: Comparison of ImageNet results for models trained on CE, SCE NGL and
focal losses. †denote results reported as 10-crop testing, all other results
show single crop accuracy.

Model | Original top1-acc | Retrained (CE) top1-acc | SCE top1-acc | Focal top1-acc | NGL top1-acc
---|---|---|---|---|---
ResNet101 He et al. (2015) | 78.25 † | 76.55 | 77.21 | 75.68 | 78.38
ResNet152 He et al. (2015) | 78.57 † | 76.92 | 77.51 | 76.06 | 78.99
Swin-T Liu et al. (2021) | 81.3 | 81.19 | 78.79 | 77.68 | 81.25
Swin-S Liu et al. (2021) | 83.0 | 83.1 | 80.68 | 79.16 | 83.0

For ResNet models NGL shows a clear improvement over other losses, increasing
the top-1 accuracy by $1-3\%$ . For the Swin architecture NGL significantly
outperforms SCE and focal losses, while it shows similar performance as
regularized CE.

The training process for ResNet101 model is demonstrated in Figure 2. It can
be seen that NGL showed slower increase in accuracy compared to other losses,
but converges later with a higher accuracy. Even though the NGL was found
using only very small datasets, it shows astonishingly good performance on
ImageNet-1k. For different architectures and model sizes NGL shows on average
superior performance compared to mentioned loss functions.

Figure 2: Mean validation accuracy of ResNet101 model on each epoch during
ImageNet-1k training.

####  4.2.3 Segmentation

Additionally, NGL was evaluated on segmentation problems. Specifically, it was
used to train two segmentation models, DeepLabv2 with ResNet101 as backbone
Nakashima (2018) and U-Net with ResNet34 as backbone. The U-Net model was
trained and evaluated only on Pascal VOC 2012 Everingham et al.  dataset,
while the DeepLabv2 model was used for both Pascal VOC 2012 and COCO-Stuff164k
Caesar et al. (2018) datasets.

The first experiments were conducted utilizing the U-Net model. The same
parameter setting was used for CE, focal, dice and NGL losses: images were
resized to $224\times 224$, Adam optimizer was applied, the learning rate was
$0.0001$ and reduced on plateau with a factor of $0.2$, patience $5$ and
$lr_{min}=10^{-6}$ was used to change it according to the IoU metric value on
validation dataset. The model was trained for $100$ epochs with a batch size
of $32$. For DeepLabv2, Adam optimizer was chosen, the learning rate $lr$ was
set to $2.5\times 10^{-5}$ and the polynomial learning rate scheduler with
weight decay $5\times 10^{-4}$ and power $0.9$ was applied to decrease $lr$
during $2\times 10^{4}$ iterations. For both models the number of program runs
was equal to $5$ and mean IoU value was used to evaluate the performance of
the model.

Results for Pascal VOC are presented in Table 4. Using the U-Net shows
approximately $3-7\%$ improvement in terms of mean IoU when NGL is used for
training compared to other losses, while showing the same results as CE and
ouperforming dice and focal losses using DeepLabv2.

Table 4: Comparison of results (mIoU values) on Pascal VOC trained by CE, focal and dice and NGL loss functions. Model | CE | Focal | Dice | NGL
---|---|---|---|---
U-Net | 50.1 | 49.1 | 46.1 | 52.8
DeepLabv2 | 76.7 | 75.9 | 74.7 | 76.7

Finally, the DeepLabv2 model was trained on COCO-Stuff164k using mentioned
loss functions. The parameter setting for these experiments was the same as
for Pascal VOC 2012 dataset with the maximum number of iterations set to
$3\times 10^{5}$. A comparison of results is shown in Table 5, showing slight
improvement for NGL loss compared to CE, and significant improvement compared
to focal and dice losses.

Table 5: Comparison of results achieved by DeepLabv2 model on COCO-Stuff164k dataset. | Pixel Accuracy | Mean IoU | Mean Accuracy
---|---|---|---
CE | $66.9$ | $39.8$ | $51.8$
Focal | $67.7$ | $39.5$ | $51.5$
Dice | $39.7$ | $17.3$ | $26.9$
NGL | $67.8$ | $39.9$ | $51.6$

The obtained results demonstrate that NGL is suitable for downstream
segmentation tasks and showing improved performance compared to other commonly
used losses.

##  5 Discussion

For qualitative analysis, the NGL function can be simplified to a binary
classification problem, similar as in Gonzalez and Miikkulainen (2019). For a
binary classification problem, NGL is as follows:

$\begin{split}f_{NGL}=\frac{1}{2}\left[e^{(\alpha-y^{(0)}_{pred}\cdot(1+y_{real}))}+e^{(\alpha-(1-y^{(0)}_{pred})\cdot(2-y_{real}))}\right]\\\
-\frac{1}{2}\left[\cos{(\cos{(\sin{(y^{(0)}_{pred})})})}+\cos{(\cos{(\sin{(1-y^{(0)}_{pred})})})}\right],\end{split}$
(2)

where $\alpha=2.4092$, $y_{pred}=(y^{(0)}_{pred},y^{(1)}_{pred})$,
$y^{(1)}_{pred}=1-y^{(0)}_{pred}$ and $y_{real}\in\\{0,1\\}$.

Let us assume that the true labels $y_{real}$ are either $0$ or $1$. The case
where $y_{real}=1$ is plotted in Figure 1 for the CE loss and NGL functions.
The cross entropy shows a monotonic decrease with $y_{pred}$ converging to 1,
the NGL shows a decrease in loss, resulting in a minimum at $0.57$ and
slightly increases when $y_{pred}$ converges to 1. This increase in loss for
$y_{pred}$ approaching the true value seems to be counter-intuitive, but may
be a factor for the overall good performance of the loss function. The
mentioned increase of the loss value may prevent the model from becoming too
confident in its output predictions and may provide an important advantage, as
it lowers the probability of overfitting. Thus, this could provide an implicit
form of regularization, enabling better generalization. A more detailed
explanation is given in Appendix C.

While some of the functions found during GP showed better performance on
certain datasets, NGL was the only function that performed better than CE
(with and without regularization), SCE, focal and dice losses on average and
was therefore used for ImageNet-1k and COCO-Stuff164k training, where it was
able to show its generalizability to larger datasets and models. The NGL is
independent of additional parameters, it is differentiable and has an implicit
regularization. For larger datasets it could be observed that it converges
slower than other mentioned losses, but at a higher maximum accuracy. The NGL
was discovered experimentally and is not supported by theory, which in turn
provides less information about the confidence of the model. Nonetheless, for
a large proportion of applications the raw performance is the main goal to
maximize. The search process did not cover the entire search space and in the
future better performing functions may be found, thereby the NGL may point to
a promising direction when searching for general suitable functions for
classification tasks.

##  6 Code

Code can be found on the project page https://github.com/ZKI-PH-
ImageAnalysis/Next-Generation-Loss/tree/main.

##  7 Conclusions

This study proposes to use Genetic Programming to search a generally
applicable loss function for image classification tasks. During that process,
a new function was found and shown to outperform other losses, commonly used
for classification and segmentation tasks, on average on a variety of datasets
demonstrating its general applicability. Moreover, it was shown that proposed
loss function can be applied to train a variety of model architectures.
Further analysis suggested that improvements provided by the new loss result
from implicit regularization that reduces overfitting to the data.

## References

  * Krizhevsky et al. [2012] Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton.  Imagenet classification with deep convolutional neural networks.  _Advances in neural information processing systems_ , 25, 2012.
  * Ronneberger et al. [2015] Olaf Ronneberger, Philipp Fischer, and Thomas Brox.  U-net: Convolutional networks for biomedical image segmentation.  _ArXiv_ , abs/1505.04597, 2015.  URL https://api.semanticscholar.org/CorpusID:3719281.
  * Girshick [2015] Ross Girshick.  Fast r-cnn.  In _Proceedings of the IEEE international conference on computer vision_ , pages 1440–1448, 2015.
  * He et al. [2015] Kaiming He, X. Zhang, Shaoqing Ren, and Jian Sun.  Deep residual learning for image recognition.  _2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_ , pages 770–778, 2015.  URL https://api.semanticscholar.org/CorpusID:206594692.
  * Szegedy et al. [2015] Christian Szegedy, Vincent Vanhoucke, Sergey Ioffe, Jonathon Shlens, and Zbigniew Wojna.  Rethinking the inception architecture for computer vision.  _2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_ , pages 2818–2826, 2015.  URL https://api.semanticscholar.org/CorpusID:206593880.
  * Chen et al. [2016] Liang-Chieh Chen, George Papandreou, Iasonas Kokkinos, Kevin P. Murphy, and Alan Loddon Yuille.  Deeplab: Semantic image segmentation with deep convolutional nets, atrous convolution, and fully connected crfs.  _IEEE Transactions on Pattern Analysis and Machine Intelligence_ , 40:834–848, 2016.  URL https://api.semanticscholar.org/CorpusID:3429309.
  * Dosovitskiy et al. [2020] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby.  An image is worth 16x16 words: Transformers for image recognition at scale.  _ArXiv_ , abs/2010.11929, 2020.  URL https://api.semanticscholar.org/CorpusID:225039882.
  * Cover and Thomas [2006] Thomas M. Cover and Joy A. Thomas.  Elements of information theory (2. ed.).  2006\.  URL https://api.semanticscholar.org/CorpusID:702542.
  * Lin et al. [2017a] Tsung-Yi Lin, Priya Goyal, Ross Girshick, Kaiming He, and Piotr Dollár.  Focal loss for dense object detection.  In _Proceedings of the IEEE international conference on computer vision_ , pages 2980–2988, 2017a.
  * Willmott and Matsuura [2005] Cort J. Willmott and Kenji Matsuura.  Advantages of the mean absolute error (mae) over the root mean square error (rmse) in assessing average model performance.  _Climate Research_ , 30:79–82, 2005.  URL https://api.semanticscholar.org/CorpusID:120556606.
  * Augusto and Barbosa [2000] D.A. Augusto and H.J.C. Barbosa.  Symbolic regression via genetic programming.  In _Proceedings. Vol.1. Sixth Brazilian Symposium on Neural Networks_ , pages 173–178, 2000.
  * K.Burke et al. [2009] Edmund K.Burke, Mathew R. Hyde, Graham Kendall, Gabriela Ochoa, Ender Ozcan, and John R. Woodward.  _Exploring Hyper-heuristic Methodologies with Genetic Programming_ , pages 177–201.  Springer Berlin Heidelberg, Berlin, Heidelberg, 2009.
  * Krizhevsky et al. [2009] Alex Krizhevsky, Geoffrey Hinton, et al.  Learning multiple layers of features from tiny images.  2009\.
  * Xiao et al. [2017] Han Xiao, Kashif Rasul, and Roland Vollgraf.  Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms, 2017.
  * Rajaraman et al. [2018] Sivaramakrishnan Rajaraman, Sameer K Antani, Mahdieh Poostchi, Kamolrat Silamut, Md A Hossain, Richard J Maude, Stefan Jaeger, and George R Thoma.  Pre-trained convolutional neural networks as feature extractors toward improved malaria parasite detection in thin blood smear images.  _PeerJ_ , 6:e4568, 2018.
  * Veeling et al. [2018] Bastiaan S Veeling, Jasper Linmans, Jim Winkens, Taco Cohen, and Max Welling.  Rotation equivariant CNNs for digital pathology.  June 2018.
  * Kather et al. [2016] Jakob Nikolas Kather, Cleo-Aron Weis, Francesco Bianconi, Susanne M Melchers, Lothar R Schad, Timo Gaiser, Alexander Marx, and Frank Gerrit Z"ollner.  Multi-class texture analysis in colorectal cancer histology.  _Scientific reports_ , 6:27988, 2016.
  * Fei-Fei et al. [2004] Li Fei-Fei, Rob Fergus, and Pietro Perona.  Learning generative visual models from few training examples: An incremental bayesian approach tested on 101 object categories.  _Computer Vision and Pattern Recognition Workshop_ , 2004.
  * Liu et al. [2021] Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo.  Swin transformer: Hierarchical vision transformer using shifted windows.  In _Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV)_ , 2021.
  * [20] M. Everingham, L. Van Gool, C. K. I. Williams, J. Winn, and A. Zisserman.  The PASCAL Visual Object Classes Challenge 2012 (VOC2012) Results.  http://www.pascal-network.org/challenges/VOC/voc2012/workshop/index.html.
  * Caesar et al. [2018] Holger Caesar, Jasper Uijlings, and Vittorio Ferrari.  Coco-stuff: Thing and stuff classes in context.  In _Computer vision and pattern recognition (CVPR), 2018 IEEE conference on_. IEEE, 2018.
  * Szegedy et al. [2014] Christian Szegedy, Wei Liu, Yangqing Jia, Pierre Sermanet, Scott E. Reed, Dragomir Anguelov, D. Erhan, Vincent Vanhoucke, and Andrew Rabinovich.  Going deeper with convolutions.  _2015 IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_ , pages 1–9, 2014.  URL https://api.semanticscholar.org/CorpusID:206592484.
  * Vaswani et al. [2017] Ashish Vaswani, Noam M. Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  In _NIPS_ , 2017.  URL https://api.semanticscholar.org/CorpusID:13756489.
  * Touvron et al. [2020] Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Herv’e J’egou.  Training data-efficient image transformers & distillation through attention.  In _International Conference on Machine Learning_ , 2020.  URL https://api.semanticscholar.org/CorpusID:229363322.
  * Chen et al. [2018] Liang-Chieh Chen, Yukun Zhu, George Papandreou, Florian Schroff, and Hartwig Adam.  Encoder-decoder with atrous separable convolution for semantic image segmentation.  In _ECCV_ , 2018.
  * Tian et al. [2022] Yingjie Tian, Duo Su, Stanislao Lauria, and Xiaohui Liu.  Recent advances on loss functions in deep learning for computer vision.  _Neurocomputing_ , 497:129–158, 2022.  URL https://api.semanticscholar.org/CorpusID:248642142.
  * Jadon [2020] Shruti Jadon.  A survey of loss functions for semantic segmentation.  _2020 IEEE Conference on Computational Intelligence in Bioinformatics and Computational Biology (CIBCB)_ , pages 1–7, 2020.  URL https://api.semanticscholar.org/CorpusID:220128180.
  * Han and Wu [2017] Bin Han and Yiquan Wu.  A novel active contour model based on modified symmetric cross entropy for remote sensing river image segmentation.  _Pattern Recognit._ , 67:396–409, 2017.  URL https://api.semanticscholar.org/CorpusID:2500127.
  * Yi-de et al. [2004] Ma Yi-de, Liu Qing, and Qian Zhi-bai.  Automated image segmentation using improved pcnn model based on cross-entropy.  _Proceedings of 2004 International Symposium on Intelligent Multimedia, Video and Speech Processing, 2004._ , pages 743–746, 2004.  URL https://api.semanticscholar.org/CorpusID:11039022.
  * Lin et al. [2017b] Tsung-Yi Lin, Priya Goyal, Ross B. Girshick, Kaiming He, and Piotr Dollár.  Focal loss for dense object detection.  _2017 IEEE International Conference on Computer Vision (ICCV)_ , pages 2999–3007, 2017b.  URL https://api.semanticscholar.org/CorpusID:47252984.
  * Sudre et al. [2017] Carole Helene Sudre, Wenqi Li, Tom Kamiel Magda Vercauteren, Sébastien Ourselin, and M. Jorge Cardoso.  Generalised dice overlap as a deep learning loss function for highly unbalanced segmentations.  _Deep learning in medical image analysis and multimodal learning for clinical decision support : Third International Workshop, DLMIA 2017, and 7th International Workshop, ML-CDS 2017, held in conjunction with MICCAI 2017 Quebec City, QC,…_ , 2017:240–248, 2017.  URL https://api.semanticscholar.org/CorpusID:21957663.
  * Salehi et al. [2017] Seyed Sadegh Mohseni Salehi, Deniz Erdoğmuş, and Ali Gholipour.  Tversky loss function for image segmentation using 3d fully convolutional deep networks.  In _MLMI@MICCAI_ , 2017.  URL https://api.semanticscholar.org/CorpusID:732793.
  * Ribera et al. [2018] Javier Ribera, David Güera, Yuhao Chen, and Edward J. Delp.  Weighted hausdorff distance: A loss function for object localization.  _ArXiv_ , abs/1806.07564, 2018.  URL https://api.semanticscholar.org/CorpusID:49322181.
  * Wong et al. [2018] Ken C. L. Wong, Mehdi Moradi, Hui Tang, and Tanveer F. Syeda-Mahmood.  3d segmentation with exponential logarithmic loss for highly unbalanced object sizes.  _ArXiv_ , abs/1809.00076, 2018.  URL https://api.semanticscholar.org/CorpusID:52157209.
  * Xu et al. [2018] Zhongwen Xu, H. V. Hasselt, and David Silver.  Meta-gradient reinforcement learning.  In _Neural Information Processing Systems_ , 2018.  URL https://api.semanticscholar.org/CorpusID:43966764.
  * Oh et al. [2020] Junhyuk Oh, Matteo Hessel, Wojciech M. Czarnecki, Zhongwen Xu, H. V. Hasselt, Satinder Singh, and David Silver.  Discovering reinforcement learning algorithms.  _ArXiv_ , abs/2007.08794, 2020.  URL https://api.semanticscholar.org/CorpusID:220633409.
  * He et al. [2019] Xin He, Kaiyong Zhao, and Xiaowen Chu.  Automl: A survey of the state-of-the-art.  _ArXiv_ , abs/1908.00709, 2019.  URL https://api.semanticscholar.org/CorpusID:199405568.
  * Li et al. [2021] Hao Li, Tianwen Fu, Jifeng Dai, Hongsheng Li, Gao Huang, and Xizhou Zhu.  Autoloss-zero: Searching loss functions from scratch for generic tasks.  _2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , pages 999–1008, 2021.  URL https://api.semanticscholar.org/CorpusID:232352674.
  * Wang et al. [2020] Xiaobo Wang, Shuo Wang, Cheng Chi, Shifeng Zhang, and Tao Mei.  Loss function search for face recognition.  _ArXiv_ , abs/2007.06542, 2020.  URL https://api.semanticscholar.org/CorpusID:220496317.
  * Li et al. [2019] Chuming Li, Chen Lin, Minghao Guo, Wei Wu, Wanli Ouyang, and Junjie Yan.  Am-lfs: Automl for loss function search.  _2019 IEEE/CVF International Conference on Computer Vision (ICCV)_ , pages 8409–8418, 2019.  URL https://api.semanticscholar.org/CorpusID:158046703.
  * Real et al. [2020] Esteban Real, Chen Liang, David R. So, and Quoc V. Le.  Automl-zero: Evolving machine learning algorithms from scratch.  In _International Conference on Machine Learning_ , 2020.  URL https://api.semanticscholar.org/CorpusID:212634211.
  * Gonzalez and Miikkulainen [2019] Santiago Gonzalez and Risto Miikkulainen.  Improved training speed, accuracy, and data utilization through loss function optimization.  _2020 IEEE Congress on Evolutionary Computation (CEC)_ , pages 1–8, 2019.  URL https://api.semanticscholar.org/CorpusID:167217832.
  * Kingma and Ba [2015] Diederik P. Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  In Yoshua Bengio and Yann LeCun, editors, _3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings_ , 2015.  URL http://arxiv.org/abs/1412.6980.
  * Ioffe and Szegedy [2015] Sergey Ioffe and Christian Szegedy.  Batch normalization: Accelerating deep network training by reducing internal covariate shift.  In _International Conference on Machine Learning_ , 2015.  URL https://api.semanticscholar.org/CorpusID:5808102.
  * Srivastava et al. [2014] Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov.  Dropout: A simple way to prevent neural networks from overfitting.  _Journal of Machine Learning Research_ , 15(56):1929–1958, 2014.
  * Loshchilov and Hutter [2019] Ilya Loshchilov and Frank Hutter.  Decoupled weight decay regularization, 2019.
  * Nakashima [2018] Kazuto Nakashima.  deeplab-pytorch.  https://github.com/kazuto1011/deeplab-pytorch, 2018.

##  Appendix A Method

Figure A: Left: an example of the solution representation for tree-based GP:
$\frac{0.5\times x}{y}-1.05\times\cos(x)$. Right: an example of the subtrees
exchange during crossover.

The GP search process starts with random initialization of a set of potential
solutions, which is also called a population of individuals, in the functional
search space. For this study a set of loss functions was randomly generated,
and each loss function was represented as a tree, where the terminals were
randomly chosen from the set $\\{y_{pred},y_{real},\mathbb{R}\\}$, while the
operators were chosen from the set
$\\{+,-,\times,\div,\times(-1),\sqrt{,}\log,\exp,\sin,\cos\\}$. To utilize
these functions and terminals for loss function generation, the following
modifications were applied:

  * •

$x\div y$ was changed to $x\div(y+\varepsilon)$;

  * •

$\sqrt{x}$ was changed to $\sqrt{|x|+\varepsilon}$;

  * •

$\log(x)$ was changed to $\log(|x|+\varepsilon)$, with $\varepsilon=10^{-8}$;

  * •

it was implemented as a rule that both $y_{pred}$ and $y_{real}$ should be
present in the formula, and $y_{pred}\in[0,1]$ (Softmax is always the last
layer of the model).

##  Appendix B Experiments

###  B.1 Loss function search

In this study, five experiments were conducted to search the loss function by
GP, which only differed in how the fitness function $F$ was defined.
Specifically, the following definitions of $F$ were used for the respective
experiments:

  * •

each individual loss function was used to train the InceptionV3 model from
scratch one time on CIFAR-10 dataset, the validation error was used as the
fitness value;

  * •

each individual loss function was used to train the InceptionV3 model from
scratch 3 times on Fashion-MNIST dataset, the averaged validation error was
used as the fitness value;

  * •

each individual loss function was used to fine-tune the pre-trained
InceptionV3 model one time on CIFAR-10 dataset, the validation error was used
as the fitness value;

  * •

each individual loss function was used to fine-tune the pre-trained
InceptionV3 model 3 times on Fashion-MNIST dataset, the averaged validation
error was used as the fitness value;

  * •

each individual loss function was used to train the InceptionV3 model from
scratch on CIFAR-10, Fashion-MNIST and CIFAR-100 datasets (once each). The
validation error for each dataset was compared to the respective validation
error obtained by the same model trained using CE loss, and individual fitness
was represented by the pair of numbers indicating the number of wins and the
percentage of improvement compared to CE.

For the first four experiments, the classification error was used to evaluate
individual’s fitness, and, therefore, the individual with the smallest
classification error was considered as the best one. However, for the last
experiment, individuals were compared in the following way:

  * •

if individual $f_{i}$ outperformed CE at least once and individual $f_{j}$
showed worse results on all 3 datasets, then $f_{i}$ was considered as the
better loss function;

  * •

if both individuals $f_{i}$ and $f_{j}$ outperformed CE the same number of
times, then the mean percentage of improvement for all winning datasets was
calculated and the loss with higher mean value was considered as the better
one;

  * •

if both individuals $f_{i}$ and $f_{j}$ showed worse results on all datasets
compared to the CE loss, then the one with lower mean second value was
considered as the better loss.

The fitness function $F$ evaluation procedure for all 5 experiments is
described in Algorithm B.1. It should be noted that in Algorithm B.1 the
number of datasets was set to 1 for the first 4 experiments and 3 for the last
one. Besides, the number of runs was equal to 1 for the first, third and last
experiments, while it was set to 3 for the remaining experiments.

Algorithm B.1 Fitness function $F$ evaluation

Test error values of the InceptionV3 model trained by using CE on all given
datasets

for each individual from population do

for each dataset used for a given experiment do

for each run do

Randomly generate train, val and test sets

Create an InceptionV3 model

Train the model for a given number of epochs

Evaluate model’s accuracy on the test set

if only one dataset is used for an experiment then

Save model’s test error

else

Determine whether individual outperformed CE or not

Save the percentage of improvement compared to CE

end if

end for

if only one dataset is used for an experiment then

Save individual’s test error averaged over all runs

else

Calculate individual’s number of wins compared to CE

Calculate the average percentage of improvement compared to CE

end if

end for

if only one dataset is used for an experiment then

Save individual’s averaged test error

else

Save the number of wins and average percentage of improvement as a pair of
values

end if

end for

The best performing functions for each of the five GP experiments are listed
below:

$f_{1}=\frac{1}{N}\sum_{i=1}^{N}{\frac{e^{(\sin^{2}(y^{(i)}_{pred}))}}{2\times(|(\beta-y^{(i)}_{real})\times
y^{(i)}_{pred}|+\varepsilon)}},$ (B.1.1)
$f_{2}=\frac{1}{N}\sum_{i=1}^{N}{\left[\frac{\cos(-y^{(i)}_{real})}{\sqrt{\left|y^{(i)}_{pred}\right|+\varepsilon}+\varepsilon}\times(y^{(i)}_{pred}-\gamma)+(y^{(i)}_{pred})^{4}+\sqrt{\left|y^{(i)}_{real}+y^{(i)}_{pred}\right|+\varepsilon}\right]},$
(B.1.2)
$f_{3}=\frac{1}{N}\sum_{i=1}^{N}{\sqrt{\left|\frac{y^{(i)}_{real}}{\sin(y^{(i)}_{real})+\varepsilon}+\frac{y^{(i)}_{pred}-e^{\delta}}{\zeta}\right|+\varepsilon}},$
(B.1.3)
$f_{4}=\frac{1}{N}\sum_{i=1}^{N}{\sin(\sin(y^{(i)}_{pred}-\eta)+(y^{(i)}_{real})^{3})},$
(B.1.4)
$f_{5}=NGL=\frac{1}{N}\sum_{i=1}^{N}{\left[e^{(\alpha-y^{(i)}_{pred}-y^{(i)}_{pred}\times
y^{(i)}_{real})}-\cos(\cos(\sin(y^{(i)}_{pred})))\right]},$ (B.1.5)

where $\alpha=2.4092$, $\beta=1.5494$, $\gamma=3.8235$, $\delta=3.1868$,
$\zeta=2.4428$, $\eta=2.6085$, $\varepsilon=10^{-8}$ and $N$ is the number of
classes. Here in all formulas $y_{pred}=(y^{(1)}_{pred},...,y^{(N)}_{pred})$
is obtained after applying softmax function to the generated prediction, thus,
$y^{(i)}_{pred}\in[0,1],i\in[1,N]$. Later in the study $f_{5}$ was denoted as
$NGL$.

###  B.2 Small datasets

Parameter settings and data preprocessing varied for all datasets but were
identical for all tested loss functions, specifics are shown in Table B.2.1.
Classification accuracy on the test set after the last epoch was used as the
evaluation metric.

Table B.2.1: Parameter settings for used small datasets.

| CIFAR-10 | Fashion-MNIST | CIFAR-100 | Malaria | PCam | Caltech 101 | Colorectal Histology
---|---|---|---|---|---|---|---
Epochs | $200$ | $100$ | $200$ | $400$ | $50$ | $400$ | $600$
Optimizer | Adam: $\beta_{1}=0.9$, $\beta_{2}=0.999$, $\epsilon=10^{-7}$ | SGD: momemtum $0.9$
Learning rate $lr$ | $0.001$ | $0.01$
Scheduler | Reduce on Plateau: factor $0.2$, patience $5$, $lr_{min}=10^{-4}$ | -
Batch size | $128$ | $32$ | $128$ | $32$
Image size | $32\times 32$ | $224\times 224$ | $75\times 75$ | $96\times 96$ | $224\times 224$ | $150\times 150$
Augmentation | Zoom $0.2$, width and height shift ($0.1$ for both), horizontral flip

The evaluation of the five found functions was performed by training ResNet50
and InceptionV3 on seven datasets, which differed by the number of images,
classes, by the type of images (grayscale and RGB), and their sizes.
Classification accuracy on the test set after the last epoch was used as the
evaluation metric. Both InceptionV3 and ResNet50 models were trained five
times for each dataset. Obtained results were averaged for both models and
compared to the averaged CE loss for the same models. In the end, the
percentage of accuracy improvement or decrease compared to CE was determined.
A final value from the range $[-0.5,0.5]$ was considered as same performing, a
value $>+0.5\%$ better performing and $<-0.5\%$ performing worse compared to
CE. Here results obtained for all five evaluated loss functions
($f_{1},...,f_{5}$) used to fine-tune pre-trained ResNet50 and InceptionV3
models are given in Table B.2.2.

Table B.2.2: Average accuracy (%) for pre-trained ResNet50 and InceptionV3
each fine-tuned 5 times for each dataset for cross entropy (CE) and percentual
change of the tested loss functions compared to CE, best result per dataset is
shown in bold.

Loss | Malaria | PCam | Colorectal Histology | CIFAR-10 | Fashion-MNIST | CIFAR-100 | Caltech 101 | Mean
---|---|---|---|---|---|---|---|---
CE | 96.7 | 72.4 | 90.5 | 94.9 | 94.6 | 74.8 | 84.5 | $\pm 0$
SCE | +0.11 | -0.83 | -2.13 | -2.12 | -1.82 | -1.25 | -0.02 | -1.13
Focal | +0.15 | +3.86 | +2.44 | +0.35 | +0.08 | -3.18 | +0.52 | +0.58
CE + $L2$ | +0.34 | +0.73 | -2.77 | +0.50 | -0.08 | -0.23 | -1.19 | -0.39
$f_{1}$ | +0.20 | +7.90 | +5.33 | -1.33 | -0.17 | +2.89 | -0.33 | +2.06
$f_{2}$ | -0.09 | +8.62 | +4.54 | +0.19 | +0.09 | +4.79 | +0.10 | +2.61
$f_{3}$ | -2.13 | +0.51 | -50.96 | +1.74 | -0.04 | +4.98 | +5.59 | -5.76
$f_{4}$ | -0.04 | +8.10 | +3.03 | +0.02 | -0.13 | +6.78 | +5.47 | +3.33
$f_{5}(NGL)$ | -0.26 | +7.07 | +5.01 | -0.37 | -0.30 | +5.29 | +3.45 | +2.84

##  Appendix C Discussion

All functions found by GP and listed in Section 1.1 were analysed the same way
as was done in Gonzalez and Miikkulainen (2019), thus, they were analysed in
the context of binary classification to show why and how they work. For
example, let us consider loss function $NGL$ for classification problem with
$N=2$ classes:

$\begin{split}NGL=\frac{1}{2}\left[e^{(\alpha-
y_{pred}^{(0)}\cdot(1+y_{real}^{(0)}))}+e^{(\alpha-
y_{pred}^{(1)}\cdot(1+y_{real}^{(1)}))}\right]-\\\
-\frac{1}{2}\left[\cos{(\cos{(\sin{(y_{pred}^{(0)})})})}+\cos{(\cos{(\sin{(y_{pred}^{(1)})})})}\right],\end{split}$
(B.2.1)

where $\alpha=2.4092$.

Since vectors $y_{pred}$ and $y_{real}$ sum to $1$, by consequence of being
passed through a softmax function, for binary classification
$y_{pred}=(y_{pred}^{(0)},1-y_{pred}^{(0)})$ and
$y_{real}=(y_{real}^{(0)},1-y_{real}^{(0)})$. This constraint simplifies the
binary $NGL$ loss to the following function of two variables
$(y_{pred}^{(0)},y_{real}^{(0)})$:

$\begin{split}NGL=\frac{1}{2}\left[e^{(\alpha-
y_{pred}^{(0)}\cdot(1+y_{real}^{(0)}))}+e^{(\alpha-(1-y_{pred}^{(0)})\cdot(2-y_{real}^{(0)}))}\right]-\\\
-\frac{1}{2}\left[\cos{(\cos{(\sin{(y_{pred}^{(0)})})})}+\cos{(\cos{(\sin{(1-y_{pred}^{(0)})})})}\right],\end{split}$
(B.2.2)

where $\alpha=2.4092$. This same methodology can be applied to the cross-
entropy loss and other functions found by GP and listed in Section 2.1.

Let us assume that true labels $y_{real}$ are either $0$ or $1$. The case
where $y_{real}=1$ is plotted in Fig. B.2 for the CE loss and loss functions
found by GP. The cross-entropy and $f_{3}$ losses are shown to be
monotonically decreasing, $f_{2}$ loss has a small region where it increases
and after that, similar to CE and $f_{3}$, it starts to decrease, while the
remaining three functions, $f_{1}$, $f_{4}$ and $NGL$, show an increase in the
loss value as the predicted label, $y_{pred}$, approaches the true label
$y_{true}$. This increase of the loss value allows the loss functions $f_{1}$,
$f_{4}$ and $NGL$ to prevent the model from becoming too confident in its
output predictions. The latter may provide an important advantage as it lowers
the probability of network’s overfitting. Thus, these loss functions provide
an implicit form of regularization, enabling better generalization.

Figure B.2: Functions found by GP and listed in Section 2.1

Results, achieved on small datasets, show that the only found function, which
monotonically decreases, $f_{3}$, demonstrates worst results as for models
trained from scratch and for pretrained models. Loss $f_{2}$ shows almost the
same results as CE on most of small datasets, but is significantly worse on
datasets such as Colorectal Histology and Caltech101. Despite the fact that
generally results for $f_{1}$ are good, it should be noted that function
$f_{1}$ (similarly to $f_{2}$) depends on the $\varepsilon$ value, which has
significant impact on the value of $f_{1}$ itself and makes the latter
differentiable. Therefore, only functions $f_{4}$ and $NGL$ have the mentioned
form of regularization and do not depend on additional parameters such as
introduced $\varepsilon$ value. It can be seen from Fig. B.2, that the minimum
for the $f_{4}$ loss, where $y_{real}=1$, lies near $0.61$, and the minimum
for the $NGL$ loss is near $0.57$ when $y_{real}=1$. All losses were used to
train ResNet101 model on ImageNet-1k dataset, however, experiments showed that
only $NGL$ showed reasonable performance.

Thus, loss $NGL$ was chosen for the experiments on large datasets, such as
ImageNet-1k and COCO-Stuff164k, as it doesn’t depend on any additional
parameters such as $\varepsilon$ to make it diffrentiable, has implicit
regularization and showed the best results on small datasets (pre-trained and
trained from scratch). Experiments on ImageNet-1k and COCO-Stuff164k showed
that found loss function $NGL$ is "slower" compared to CE, for example, as the
latter is capable to reach generally good accuracy very quickly. Nevertheless,
given the same number of epochs, $NGL$ loss manages to catch up and surpass CE
later during the training. It can be caused by the fact that at some point
during the training cross-entropy loss stops or barely improves deep learning
model’s accuracy, while $NGL$ function is more "cautious" during the training
due to its implicit regularization, which helps model to prevent overfitting.


11institutetext: Department of Statistics, LMU Munich
22institutetext: Department of Radiology, University Hospital, LMU Munich
33institutetext: Munich Center for Machine Learning (MCML)
33email<EMAIL_ADDRESS>
# Constrained Probabilistic Mask Learning for Task-specific Undersampled MRI
Reconstruction

Tobias Weber 112233 0000-0002-5430-2595 Michael Ingrisch 2233
0000-0003-0268-9078 Bernd Bischl 1133 0000-0001-6002-6980 David Rügamer 1133
0000-0002-8772-9202

###### Abstract

Undersampling is a common method in Magnetic Resonance Imaging (MRI) to
subsample the number of data points in k-space and thereby reduce acquisition
times at the cost of decreased image quality. In this work, we directly learn
the undersampling masks to derive task- and domain-specific patterns. To solve
this discrete optimization challenge, we propose a general optimization
routine called ProM: A fully probabilistic, differentiable, versatile, and
model-free framework for mask optimization that enforces acceleration factors
through a convex constraint. Analyzing knee, brain, and cardiac MRI datasets
with our method, we discover that different anatomic regions reveal distinct
optimal undersampling masks. Furthermore, ProM can create undersampling masks
that maximize performance in downstream tasks like segmentation with networks
trained on fully-sampled MRIs. Even with extreme acceleration factors, ProM
yields reasonable performance while being more versatile than existing
methods, paving the way for data-driven all-purpose mask generation.

###### Keywords:

MRI undersampling discrete optimization bernoulli.

##  1 Introduction

Undersampling is an important tool to speed up the acquisition time in
magnetic resonance imaging (MRI) by selectively sampling data points in
k-space. This can, however, result in decreased image quality and image
artifacts. To address this issue, various techniques have been developed to
enhance reconstructions and produce high-quality images from undersampled
data. The majority of research focuses on improving the reconstruction from an
undersampled MRI, e.g., via compressed sensing [16], or, more recently, deep
learning. Examples include transformer-based approaches, e.g., for the
reconstruction of radial trajectories [11] and sparse-view CTs [26], or
diffusion models [7, 8, 19]. Further approaches include variational models to
jointly synthesize and reconstruct MRI images [6], sharpening networks [10] to
counter the absence of high-frequency features in undersampled MRIs, or
learnable Fourier interpolation [9]. To account for meta information such as
the manufacturer, [15] condition the reconstruction network on side
information. Using adversarial methods, [4] is able to significantly enhance
images with extreme acceleration factors.

Rather than focusing on enhancing the image quality with a predetermined
undersampling pattern, the challenge addressed in this paper lies in
identifying the optimal mask in terms of reconstruction quality for a given
undersampling ratio. In the literature, this is mainly addressed by combined
approaches that simultaneously learn a reconstruction network and an
undersampling mask [1, 25, 27, 23, 24]. As in this work, the combined approach
from [28] considers probabilistic masks but obtains sparsity via pruning. In
contrast, and most closely related to our work, [20, 21] propose a direct mask
optimization scheme based on iterative gradient sampling (IGS), which
repeatedly determines k-space elements that contribute the most to a loss
criterion. They argue that undersampling masks could be optimized directly but
“the undersampled pattern is binary, which cannot be trained by the gradient
descent”.

In this work, we propose ProM: A fully differentiable probabilistic framework
for mask optimization. By framing the search for an optimal mask as a
probabilistic optimization problem for a pre-specified acceleration factor,
ProM is able to find the optimal undersampling distribution using ideas from
relaxed categorical optimization in deep learning research. In particular,
this allows tailored results to the given downstream task and anatomic region
(Figure 1), making it a versatile and data-driven all-purpose mask generator.

ProM undersampling mask optimization Figure 1:  Visualization of the ProM
optimization routine for the ACDC dataset with an acceleration factor of x8
for a 2D mask. The bottom row shows our Bernoulli mask distribution
$p_{\bm{\theta}}$, where a lighter color implies a higher probability for
sampling the respective entry in the cartesian k-space grid. Starting from a
randomly initialized distribution, ProM gradually optimizes $p_{\bm{\theta}}$
to maximize reconstruction quality while simultaneously increasing the
sparsity of the masks. The resulting distribution converges to a domain-
specific mask with desired acceleration factor that preserves most of the
image’s quality (top row). The original image is displayed on the right.

##  2 Methods

Following [31], we reformulate the idea of learning sparse neural network
weights and introduce a gradient descent routine to learn fully probabilistic
undersampling masks. In the following, we describe our routine for a single
image and later extend this idea to jointly optimize masks across a whole
dataset.

###  2.1 Probabilistic Undersampling Masks

In the following, we define $D$ as the number of elements on the 2-dimensional
or 3-dimensional k-space grid and use a vectorized notation for all objects
for simplicity. Thus,
$\bm{x}_{k}=(x_{k}^{(1)},\ldots,x_{k}^{(D)})^{\top}\in\mathbb{C}^{D}$ depicts
an image residing in k-space. Partial acquisition is augmented by applying a
binary mask $\bm{m}\in\\{0,1\\}^{D}$ to the fully-sampled $\bm{x}_{k}$
element-wise: $\bm{x}_{k}\odot\bm{m}$.

Instead of following a distinct sampling pattern for $\bm{m}$, we assume that
every element $m^{(i)},i=1,...,D,$ in $\bm{m}$ is the result of an independent
Bernoulli experiment of a random variable $M^{(i)}$ and distribution defined
by

$p_{\bm{\theta}}:=\mathbb{P}(\bm{M}=\bm{m})=\textstyle\prod_{i=1}^{D}\theta_{i}^{m^{(i)}}(1-\theta_{i})^{1-m^{(i)}},$
(1)

$\bm{M}=(M^{(1)},\ldots,M^{(D)})^{\top}$, and $\theta_{i}\in(0,1)$ the
sampling probability of the $i$th element. Eq. 1 can be thought of as a
Bayesian prior for the image mask. Once a mask is sampled, the translation of
the image from the under-sampled k-space to the complex image domain
$\mathcal{X}\subseteq\mathbb{C}^{D}$ is done by
$\bm{x}_{c}=\mathcal{F}^{-1}(\bm{x}_{k}\odot\bm{m})$, where $\mathcal{F}^{-1}$
is the inverse Fourier transform matrix. This procedure amounts to a simple
zero-filling reconstruction strategy.

We strive to allow the optimization of the (posterior) sampling distribution
(i.e., after accounting for the specific reconstruction task and data) for
arbitrary differentiable loss functions $\mathcal{L}$ used in general vision
deep learning. As these loss functions are typically designed to only work in
real-valued space, we transform $\bm{x}_{c}$ into a real-valued representation
using its magnitude image $\hat{\bm{x}}=|\bm{x}_{c}|\in\mathbb{R}^{D}$. The
quality of the reconstruction can then be assessed by
$\mathcal{L}(\hat{\bm{x}},\bm{x})$, where $\bm{x}$ is the fully-sampled
original image. Given a data point $\bm{x}$, our approach directly maximizes
the posterior (or equally, minimizes the posterior empirical risk) by finding
$\bm{\theta}$ as

$\operatorname*{\arg\min}_{\bm{\theta}}\int\mathcal{L}(\hat{\bm{x}},\bm{x})\,dp_{\bm{\theta}}=\operatorname*{\arg\min}_{\bm{\theta}}\mathbb{E}_{\bm{m}\sim
p_{\bm{\theta}}}\mathcal{L}(\hat{\bm{x}},\bm{x})\approx\operatorname*{\arg\min}_{\bm{\theta}}\frac{1}{L}\sum_{l=1}^{L}\mathcal{L}(\hat{\bm{x}}_{l},\bm{x})\,\text{,}$
(2)

where we approximate the expectation using $L$ Monte-Carlo samples
$\hat{\bm{x}}_{l}$.

Without further constraints, the optimal solution of Eq. 2 is
$\bm{\theta}=\bm{1}$, i.e., the fully-sampled image. We, therefore, introduce
an undersampling constraint. This can be done similar to [31] by limiting the
sum of all probabilities in $p_{\bm{\theta}}$ to a pre-specified value $S$,
i.e., $\sum_{i=1}^{D}\theta_{i}\leq S$. Practically speaking, $S$ will result
in the number of sampled k-space elements as $\sum_{i=1}^{D}\theta_{i}$ is the
expected value of $||\bm{m}||_{0}$. Given an acceleration factor $\alpha$, we
can define $S=\lfloor\frac{D}{\alpha}\rfloor$ and our final objective:

$\displaystyle\operatorname*{\arg\min}_{\bm{\theta}}\mathbb{E}_{\bm{m}\sim
p_{\bm{\theta}}}\mathcal{L}(\hat{\bm{x}},\bm{x})\>\>\>\>\>\>s.t.\>\textstyle\sum_{i=1}^{D}\theta_{i}\leq
S\>\text{and}\>S\in\\{0,...,D\\}.$ (3)

The constraint in Eq. 3, which can equally be expressed as an $\ell_{1}$-norm
penalty for $\bm{\theta}$, has an intrinsic affinity to be sparse (c.f. Figure
2(a)).

###  2.2 Differentiability over Reparameterization

So far, the proposed approach is fully differentiable including
$\mathcal{F}^{-1}$ and $|.|$, except for the stochastic sampling of $\bm{m}$.
In order to use modern autograd frameworks for stochastic masks, we apply the
Gumbel-Softmax trick [12] tailored to the Bernoulli distribution [31]. Let
$\bm{\rho}:=\log\left(\frac{\bm{\theta}}{1-\bm{\theta}}\right)$ be the log
odds-ratio for $\bm{\theta}$. Then, a “soft mask” $\bm{m}_{\text{\tiny soft}}$
allowing for differentiability can be obtained by sampling

$\bm{m}_{\text{\tiny
soft}}\sim\mathds{1}\left(\bm{\rho}+\bm{g}_{1}-\bm{g}_{0}\geq\bm{0}\right)\approx\sigma\left((\bm{\rho}+\bm{g}_{1}-\bm{g}_{0})\tau^{-1}\right)\,\text{,}$
(4)

where $\bm{g}_{1},\bm{g}_{0}$ are independent and identically distributed
samples from a Gumbel(0,1) distribution and the indicator function
$\mathds{1}$ is relaxed using a sigmoid function $\sigma$. A temperature
parameter $\tau$ controls the softness of the discrete approximation and is
annealed during optimization. Here, stochasticity is rerouted over the Gumbel
samples and thus a computational graph is able to propagate gradients to
$\bm{\theta}$. As the mask $\bm{m}$ needs to be strictly binary, which is not
the case for $\bm{m}_{\text{\tiny soft}}$, we adopt the straight-through
Gumbel estimator trick [12], yielding

$\bm{m}=\mathds{1}\left(\bm{m}_{\text{\tiny soft}}\geq
0.5\right)-sg[\bm{m}_{\text{\tiny soft}}]+\bm{m}_{\text{\tiny
soft}}\,\text{,}$ (5)

where $\mathds{1}$ is applied element-wise and returns our final binary mask
sample. $sg$ denotes the stop gradient operation, which blocks gradients from
backpropagation. In other words, Eq. (5) yields discrete values while we
obtain gradients for its soft approximation.

###  2.3 Constrained Optimization via Projected Gradient Descent

The optimization problem in Eq. 3 cannot be solved effectively with standard
gradient descent. Instead, we follow [31] and use a projected gradient
approach by first updating the unconstrained parameter vector
$\tilde{\bm{\theta}}=[\tilde{\theta}_{1},...,\tilde{\theta}_{D}]$ using
$\tilde{\bm{\theta}}=\bm{\theta}-\eta\nabla_{\bm{\theta}}\mathbb{E}_{\bm{m}\sim
p_{\bm{\theta}}}\mathcal{L}(\hat{\bm{x}},\bm{x})$ with $\eta$ being the
learning rate, and then project $\tilde{\bm{\theta}}$ into the space of valid
elements by solving

$\textstyle\sum_{i=1}^{D}\min(1,\max(0,\tilde{\theta_{i}}-\lambda))=S$ (6)

for $\lambda\in\mathbb{R}$, yielding

$\bm{\theta}=\min(\bm{1},\max(\bm{0},\tilde{\bm{\theta}}-\max(0,\lambda)\bm{1}))\,\text{.}$
(7)

A solution to Eq. (6) can be obtained using a convex solver or root-finding
method such as bisection search. To foster exploration at the beginning of the
training and allow for exploitation at later stages, we anneal $S$ during
optimization. First, exploration iterations with $S=D$ allow for unrestricted
optimization. Then, an annealing phase following the schedule of [32]
decreases $S$ to meet the desired acceleration factor. Finally, the
exploitation phase optimizes $\bm{\theta}$ under the nominal constrained $S$
(c.f. Figure 2(b)).

Our procedure is summarized in Algorithm 1 for a single image $\bm{x}_{k}$.
The same procedure can be applied to a whole dataset by iteratively taking
random batches for the optimization of $\bm{\bm{\theta}}$. Note that the only
trainable parameters in ProM are the $D$ parameters $\bm{\bm{\theta}}$.
Optimization can thus be done in a matter of seconds (single image) or a few
minutes (dataset).

(a) Histograms of probabilities in $\bm{\theta}$ during different optimization
phases. Starting from a random initialization, the sparsification constraint
leads to a distribution with probabilities that tend to be close to zero or
one.

ConstrainingExplorationExploitation (b) Sparsity versus reconstruction
quality. The amount of active masked elements in $\bm{m}$ is bounded by the
continuously annealed dense rate $\frac{S}{D}$. During optimization, the
average SSIM metric (measuring reconstruction) plummets when $\bm{\theta}$ is
limited by the constraint (at around 20% progress) but roughly stays constant
or even slightly improves while the number of active elements is further
restricted.

Figure 2: Optimization phases of ProM for a fastMRI Knee sample with acc.
factor x16.

##  3 Experiments

We investigate the performance of ProM using slices from ACDC [5], BraTS [17,
2, 3] and the fastMRI Knee [29] dataset:

ACDC are cardiac MRIs with 100 train and 50 test subjects. We extract the end-
diastolic frame in 256px resolution including segmentation labels of the left
and right ventricular cavity as well as the left ventricular myocardium,
yielding $548$ train and $338$ test slices. k-space data is emulated via
Fourier transform.

BraTS contains brain MRIs with T2-FLAIR, T1-, T1Gd- and T2-weighted
modalities. The goal of the segmentation is to determine the classes of whole,
core, and enhancing tumors. The dataset is split into 387 train and 97 test
subjects. We extract $19,350$ train and $4,850$ test slices using slice
indices between $60$ and $110$ with 256px resolution. k-space data is emulated
as in ACDC.

fastMRI Knee includes raw k-space data of single coil knee MRI. To focus on
pathologies, we extract the annotated subset of fastMRI+ [30] amounting to
$8,057$ train and $1524$ test slices with a center-crop to 320px resolution.

For ProM we use 2500 iterations with a learning rate of $0.01$ in the Adam
[13] optimizer. This configuration provided an ideal balance between runtime
and convergence. We use batches of size 32 and draw $L=4$ Bernoulli samples
for each sample, yielding a total batch size of 128. We found that similar to
[14], a low number of Monte Carlo samples is sufficient if the batch size is
large. The temperature $\tau$ follows a linearly decreasing schedule from $1$
to $0.03$ in the last step, which is in line with [31]. For reconstruction, we
use the mean squared error for $\mathcal{L}$. The final masks are obtained by
applying Bayesian model averaging on 10 different ProM solutions. Optimization
is done in PyTorch v.1.13 [18] on an NVIDIA A100 GPU. We compare our approach
against an equispaced mask with fully-sampled central region of 4% [29] and
random offset, a 2D variable density Gaussian mask, and the most recent mask
optimization approach Iterative Gradients Sampling (IGS; [21]). Reported
metrics represent the mean across 10 randomly initialized runs per method.

###  3.1 Domain-specific Masks

Each anatomic region (as a composition of different elemental shapes) yields a
distinct k-space representation. Our first experiment not only proves that
this is the case in practice but also shows that a data-driven optimization
routine for masks such as ProM is indeed necessary to facilitate optimal
reconstruction (c.f. Figure 3 for the results of ProM trained on the three
different datasets). For example, the cardiac ACDC consists dominantly of
elliptic primitives, which results in a completely different optimal mask than
fastMRI Knee with a lot of vertical lines and some horizontal elements in
image space.

| x4 | x8 | x16 | x32 | Example
---|---|---|---|---|---
ACDC |  |  |  |  |
BraTS |  |  |  |  |
fastMRI Knee |  |  |  |  |

Figure 3: Optimized undersampling 2D mask distributions for ACDC, BraTS and
fastMRI Knee (rows) with varying acceleration factors (columns). Different
anatomic regions (right) have a distinct unique optimal distribution.

###  3.2 Reconstruction Quality

To assess the reconstruction quality of ProM we evaluate the peak signal-to-
noise ratio (PSNR), structural similarity index measure (SSIM), and normalized
mean squared error (NMSE) for acceleration factor $\alpha$ x4 to x32 on the
test set of fastMRI Knee. Results (Table 1) show that the IGS method works
well, the additional flexibility of ProM to operate in 2D, however, allows to
obtain superior results. This advantage increases for higher $\alpha$. As an
instance, in the case of factor x4, the SSIM / NMSE of ProM is 5.00% / 19.23%
better compared to IGS, with the improvement increasing to $15.64$% /
$134.78\%$ for factor x32. Moreover, when compared to more conventional
masking approaches, the gain in quality of ProM is even larger. In addition,
Figure 4 shows that the ProM mask does not introduce noise nor artifacts,
unlike the other evaluated methods. However, even with a custom-tailored
undersampling strategy, high-frequency details are omitted to a certain
degree.

| Equispaced | 2D Gaussian | IGS | ProM | Full
---|---|---|---|---|---
x16 |  |  |  |  |
x32 |  |  |  |  |

Figure 4: Reconstructions of a fastMRI Knee sample. Equidistant spacing and
IGS display strong infolding artifacts, which are less pronounced with
Gaussian sampling. The reconstruction with the data-driven mask of ProM is
much closer to the fully-sampled image and reduces image noise as well as
artifacts very clearly. Table 1: Quality of reconstruction for the fastMRI
Knee dataset.

| x4 | x8 | x16 | x32
---|---|---|---|---
|  PSNR $\uparrow$ |  SSIM $\uparrow$ |  NMSE $\downarrow$ |  PSNR $\uparrow$ |  SSIM $\uparrow$ |  NMSE $\downarrow$ |  PSNR $\uparrow$ |  SSIM $\uparrow$ |  NMSE $\downarrow$ |  PSNR $\uparrow$ |  SSIM $\uparrow$ |  NMSE $\downarrow$
Equi. | 24.481 | 0.601 | 0.066 | 23.803 | 0.524 | 0.077 | 23.365 | 0.474 | 0.085 | 23.158 | 0.448 | 0.090
Gauss. | 29.570 | 0.664 | 0.027 | 28.118 | 0.560 | 0.035 | 23.342 | 0.440 | 0.112 | 17.373 | 0.299 | 0.396
IGS | 28.553 | 0.640 | 0.031 | 26.428 | 0.532 | 0.047 | 24.376 | 0.458 | 0.070 | 22.313 | 0.409 | 0.108
ProM | 29.787 | 0.672 | 0.026 | 28.472 | 0.570 | 0.034 | 27.575 | 0.511 | 0.040 | 26.739 | 0.473 | 0.046

###  3.3 Zero-shot Undersampled Segmentation

Visual reconstruction quality does not necessarily correlate with performance
in downstream tasks such as segmentation [20]. We now investigate the quality
of ProM in a transfer learning scenario, where we apply a frozen segmentation
network trained on fully-sampled MRIs and learn a mask to maximize its
performance with undersampled MRIs. For this, we choose a standard U-net [22]
implementation and training routine with channel multipliers of (16, 32, 64,
128, 256). We substitute $\mathcal{L}$ and use the trained segmentation
network $net$ as a proxy paired with a combination of Dice and cross-entropy
loss $\mathcal{L}_{\text{seg}}$, i.e. $\mathcal{L}(\hat{\bm{x}},\bm{x})$
becomes $\mathcal{L}_{\text{seg}}(net(\hat{\bm{x}}),\bm{x}_{\text{seg}})$
where $\bm{x}_{\text{seg}}$ is the segmentation label of $\bm{x}$.
Additionally, for the task of segmentation, we introduce a 1D variant of ProM
to better understand its gain in performance compared to other methods.

The segmentation performance is evaluated in Table 2 via Dice score and
Intersection-over-Union (IuO). Examples are visualized in Figure 5. The fully-
sampled MRI achieve a macro-averaged Dice score of 0.855 and an IoU of 0.763
for ACDC as well as 0.772 and 0.710 for BraTS, respectively. While 1D ProM
does not surpass the performance of IGS, our contention is that 1D ProM would
benefit from optimized parameters specifically tailored for 1D masking. 2D
ProM achieves competitive segmentation performance for $\alpha$ x8 and is
overall notably better for extreme $\alpha$ such as x32 with a Dice score /
IoU of $0.727$ / $0.606$ in ACDC and $0.706$ / $0.608$ in BraTS. The Gaussian
mask appears to be an efficient and simple method to achieve competitive
performance for smaller $\alpha$. Note that the ProM mask determined by
$\mathcal{L}_{\text{seg}}$ differs substantially from the one obtained to
maximize reconstruction quality (Figure 3), emphasizing the fact that an
optimal mask does not only vary with the dataset but also with the task.

| Equispaced | 2D Gaussian | IGS | ProM (1D) | ProM (2D) | Full
---|---|---|---|---|---|---
ACDC |  |  |  |  |  |
BraTS |  |  |  |  |  |

Figure 5: Segmentation with a pre-trained U-net using
$\mathcal{L}_{\text{seg}}$ in x16 accelerated ProM. The contours in ACDC
correspond to the left (o) and right (o) ventricular cavity as well as the
left ventricular myocardium (o). BraTS markings in the T1Gd sample imply whole
(o), core (o) and enhancing (o) tumor.  Table 2: Segmentation metrics using a
pre-trained U-net on ACDC and BraTS, showing that ProM (2D) excels especially
for extreme acceleration factors.

|  |  | ACDC | BraTS
---|---|---|---|---
|  |  |  Equi. |  Gauss. |  IGS |  ProM (1D) |  ProM (2D) |  Equi. |  Gauss. |  IGS |  ProM (1D) |  ProM (2D)
x8 | Dice | $\uparrow\,\,$ |  0.671 |  0.847 |  0.828 |  0.762 |  0.839 |  0.596 |  0.733 |  0.716 |  0.646 |  0.739
IoU | $\uparrow\,\,$ |  0.546 |  0.752 |  0.726 |  0.650 |  0.742 |  0.489 |  0.638 |  0.619 |  0.542 |  0.646
x16 | Dice | $\uparrow\,\,$ |  0.645 |  0.644 |  0.745 |  0.717 |  0.789 |  0.589 |  0.597 |  0.651 |  0.537 |  0.735
IoU | $\uparrow\,\,$ |  0.517 |  0.534 |  0.627 |  0.599 |  0.679 |  0.481 |  0.494 |  0.546 |  0.426 |  0.640
x32 | Dice | $\uparrow\,\,$ |  0.644 |  0.399 |  0.592 |  0.587 |  0.727 |  0.580 |  0.315 |  0.537 |  0.483 |  0.706
IoU | $\uparrow\,\,$ |  0.517 |  0.301 |  0.466 |  0.460 |  0.606 |  0.472 |  0.226 |  0.428 |  0.374 |  0.608

##  4 Conclusion

This paper proposes ProMas a general framework for data-driven and
probabilistic undersampling mask learning. We evaluated ProM on cardiac,
brain, and knee MRIs. Our approach shows promising performance, even when
reaching acceleration factors of x16 and x32. In clinical practice, ProM would
allow for significantly reduced acquisition times, e.g., for high-speed
interventions or pathology localization. Our method further allows deriving
these optimal masks in a matter of seconds from a single data sample without
requiring a large computing infrastructure. As there are no shape requirements
for ProM, future research directions could investigate patterns for radial
masks or the reconstruction of multi-coil images. Additionally, custom masks
for certain scanner parameters, acquisition protocols, or manufacturers could
be derived.

## Acknowledgments

This work has been partially funded by the Deutsche Forschungsgemeinschaft
(DFG, German Research Foundation) as part of BERD@NFDI - grant number
460037581. The authors gratefully acknowledge LMU Klinikum for providing
computing resources on their Clinical Open Research Engine (CORE).

## References

  * [1] Bahadir, C.D., Dalca, A.V., Sabuncu, M.R.: Learning-based optimization of the under-sampling pattern in MRI. In: Information Processing in Medical Imaging: 26th International Conference, IPMI. pp. 780–792 (2019)
  * [2] Bakas, S., Akbari, H., Sotiras, A., Bilello, M., Rozycki, M., et al.: Advancing the cancer genome atlas glioma mri collections with expert segmentation labels and radiomic features. Scientific data 4(1), 1–13 (2017)
  * [3] Bakas, S., Reyes, M., Jakab, A., Bauer, S., Rempfler, M., et al.: Identifying the best machine learning algorithms for brain tumor segmentation, progression assessment, and overall survival prediction in the brats challenge. arXiv:1811.02629 [cs, stat] (2018)
  * [4] Belov, A., Stadelmann, J., Kastryulin, S., Dylov, D.V.: Towards ultrafast mri via extreme k-space undersampling and superresolution. In: Medical Image Computing and Computer Assisted Intervention, MICCAI. pp. 254–264 (2021)
  * [5] Bernard, O., Lalande, A., Zotti, C., Cervenansky, F., Yang, X., Heng, P.A., Cetin, I., Lekadir, K., Camara, O., Ballester, M.A.G., et al.: Deep learning techniques for automatic mri cardiac multi-structures segmentation and diagnosis: is the problem solved? IEEE Transactions on Medical Imaging 37(11), 2514–2525 (2018)
  * [6] Bian, W., Zhang, Q., Ye, X., Chen, Y.: A learnable variational model for joint multimodal MRI reconstruction and synthesis. In: Medical Image Computing and Computer Assisted Intervention, MICCAI. pp. 354–364 (2022)
  * [7] Chung, H., Ye, J.C.: Score-based diffusion models for accelerated mri. Medical Image Analysis 80, 102479 (2022)
  * [8] Dar, S.U., Öztürk, Ş., Korkmaz, Y., Elmas, G., Özbey, M., et al.: Adaptive diffusion priors for accelerated mri reconstruction. arXiv:2207.05876 [cs, eess] (2022)
  * [9] Ding, Q., Zhang, X.: Mri reconstruction by completing under-sampled k-space data with learnable fourier interpolation. In: Medical Image Computing and Computer Assisted Intervention, MICCAI. pp. 676–685 (2022)
  * [10] Dong, S., Chen, E.Z., Zhao, L., Chen, X., Liu, Y., Chen, T., Sun, S.: Invertible sharpening network for mri reconstruction enhancement. In: Medical Image Computing and Computer Assisted Intervention, MICCAI. pp. 582–592 (2022)
  * [11] Gao, C., Shih, S.F., Finn, J.P., Zhong, X.: A projection-based k-space transformer network for undersampled radial mri reconstruction with limited training subjects. In: Medical Image Computing and Computer Assisted Intervention, MICCAI. pp. 726–736 (2022)
  * [12] Jang, E., Gu, S., Poole, B.: Categorical reparameterization with gumbel-softmax. In: 5th International Conference on Learning Representations, ICLR (2017)
  * [13] Kingma, D.P., Ba, J.: Adam: A method for stochastic optimization. In: 3rd International Conference on Learning Representations, ICLR (2015)
  * [14] Kingma, D.P., Welling, M.: Auto-encoding variational bayes. In: 2nd International Conference on Learning Representations, ICLR (2014)
  * [15] Liu, X., Wang, J., Peng, C., Chandra, S.S., Liu, F., Zhou, S.K.: Undersampled mri reconstruction with side information-guided normalisation. In: Medical Image Computing and Computer Assisted Intervention, MICCAI. pp. 323–333 (2022)
  * [16] Lustig, M., Donoho, D., Pauly, J.M.: Sparse MRI: The application of compressed sensing for rapid mr imaging. Magnetic Resonance in Medicine 58(6), 1182–1195 (2007)
  * [17] Menze, B.H., Jakab, A., Bauer, S., Kalpathy-Cramer, J., Farahani, K., Kirby, J., Burren, Y., Porz, N., Slotboom, J., Wiest, R., et al.: The multimodal brain tumor image segmentation benchmark (BRATS). IEEE Transactions on Medical Imaging 34(10), 1993–2024 (2014)
  * [18] Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., et al.: Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems 32 (2019)
  * [19] Peng, C., Guo, P., Zhou, S.K., Patel, V.M., Chellappa, R.: Towards performant and reliable undersampled mr reconstruction via diffusion model sampling. In: Medical Image Computing and Computer Assisted Intervention, MICCAI. pp. 623–633 (2022)
  * [20] Razumov, A., Rogov, O.Y., Dylov, D.V.: Optimal MRI undersampling patterns for ultimate benefit of medical vision tasks. arXiv:2108.04914 [cs, eess] (2021)
  * [21] Razumov, A., Rogov, O.Y., Dylov, D.V.: Optimal mri undersampling patterns for pathology localization. In: Medical Image Computing and Computer Assisted Intervention, MICCAI. pp. 768–779 (2022)
  * [22] Ronneberger, O., Fischer, P., Brox, T.: U-net: Convolutional networks for biomedical image segmentation. In: Medical Image Computing and Computer-Assisted Intervention, MICCAI. pp. 234–241 (2015)
  * [23] Wang, G., Luo, T., Nielsen, J.F., Noll, D.C., Fessler, J.A.: B-spline parameterized joint optimization of reconstruction and k-space trajectories (bjork) for accelerated 2d mri. IEEE Transactions on Medical Imaging 41(9), 2318–2330 (2022)
  * [24] Weiss, T., Senouf, O., Vedula, S., Michailovich, O., Zibulevsky, M., Bronstein, A.: Pilot: Physics-informed learned optimized trajectories for accelerated mri. arXiv:1909.05773 (2019)
  * [25] Weiss, T., Vedula, S., Senouf, O., Michailovich, O., Zibulevsky, M., Bronstein, A.: Joint learning of cartesian under sampling andre construction for accelerated mri. In: IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP). pp. 8653–8657. IEEE (2020)
  * [26] Xia, W., Yang, Z., Zhou, Q., Lu, Z., Wang, Z., Zhang, Y.: A transformer-based iterative reconstruction model for sparse-view ct reconstruction. In: Medical Image Computing and Computer Assisted Intervention, MICCAI. pp. 790–800 (2022)
  * [27] Xie, J., Zhang, J., Zhang, Y., Ji, X.: Puert: Probabilistic under-sampling and explicable reconstruction network for cs-mri. IEEE Journal of Selected Topics in Signal Processing 16(4), 737–749 (2022)
  * [28] Xue, S., Cheng, Z., Han, G., Sun, C., Fang, K., et al.: 2d probabilistic undersampling pattern optimization for mr image reconstruction. Medical Image Analysis 77, 102346 (2022)
  * [29] Zbontar, J., Knoll, F., Sriram, A., Murrell, T., Huang, Z., Muckley, M.J., Defazio, A., Stern, R., Johnson, P., Bruno, M., et al.: fastmri: An open dataset and benchmarks for accelerated mri. arXiv:1811.08839 [physics, stat] (2019)
  * [30] Zhao, R., Yaman, B., Zhang, Y., Stewart, R., Dixon, A., Knoll, F., Huang, Z., Lui, Y.W., Hansen, M.S., Lungren, M.P.: fastmri+: Clinical pathology annotations for knee and brain fully sampled multi-coil mri data. arXiv:2109.03812 [physics] (2021)
  * [31] Zhou, X., Zhang, W., Xu, H., Zhang, T.: Effective sparsification of neural networks with global sparsity constraint. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. pp. 3599–3608 (2021)
  * [32] Zhu, M., Gupta, S.: To prune, or not to prune: Exploring the efficacy of pruning for model compression. In: 6th International Conference on Learning Representations, ICLR, Workshop Track Proceedings (2018)

##  5 Supplementary Material

###  5.1 ProMAlgorithm

Algorithm 1 Optimization routine of ProM

k-space image $\bm{x}_{k}$, mask distribution $p_{\bm{\theta}}$, number of
samples $L$, acc. factor $\alpha$, iterations $i$, learning rate $\eta$,
criterion $\mathcal{L}$

for $1,...,i$ do

$\mathcal{M}\leftarrow\\{\bm{m}_{1},...,\bm{m}_{L}\\}$ with
$\bm{m}_{\cdot}\sim p_{\theta}$ $\triangleright$ Eq. (4) and (5)

$\bm{X}_{k}\leftarrow\text{expand}(\bm{x}_{k})$ $\triangleright$ Expand to
match $L$

$\hat{\bm{X}}\leftarrow\lvert\mathcal{F}^{-1}(\bm{X}_{k}\odot\mathcal{M})\rvert$
$\triangleright$ Undersampled magnitude images

$\bm{X}\leftarrow\text{expand}(\lvert\mathcal{F}^{-1}(\bm{x}_{k})\rvert)$
$\triangleright$ Fully-sampled image and expand

$\tilde{\bm{\theta}}\leftarrow\bm{\theta}-\eta\nabla_{\bm{\theta}}\mathcal{L}(\hat{\bm{X}},\bm{X})$
$\triangleright$ Compute batched loss and apply gradients

$S\leftarrow\text{anneal}(\alpha)$ $\triangleright$ Anneal to acc. factor over
iterations

$\bm{\theta}\leftarrow\text{project}(\tilde{\bm{\theta}},S)$ $\triangleright$
Eq. (6) and (7)

end for

###  5.2 $S$ annealing schedule

Acceleration factor $\alpha$, annealing start iteration $i_{min}$, annealing
end iteration $i_{max}$, current iteration $i_{cur}$,
$d_{target}=\frac{1}{\alpha}$.

$\displaystyle d_{cur}$
$\displaystyle=d_{target}+(1-d_{target})\left(1-\frac{i_{cur}-i_{min}}{i_{max}-i_{min}}\right)$
$\displaystyle S_{cur}$ $\displaystyle=d_{cur}\cdot D$

###  5.3 Proof of Constraint Projection

Proof for Eq. 6 and Eq. 7 is taken and adapted from [31]. Transforming updated
parameters $\tilde{\bm{\theta}}\in\mathbb{R}^{D}$ into $\bm{\theta}$, which
fulfills the sparsification constraint can be described as a least-squares
convex problem:

$\operatorname*{\arg\min}_{\bm{\theta}}\frac{1}{2}\lvert\lvert\tilde{\bm{\theta}}-\bm{\theta}\rvert\rvert^{2}\>\>\>\>\>\>s.t.\>\sum_{i=1}^{D}\theta_{i}=\bm{1}^{\top}\bm{\theta}\leq
S\>\text{and}\>0\leq\theta_{i}\leq 1\,\text{.}$ (8)

This can be solved by the Lagrangian multiplier method:

$\displaystyle\mathfrak{L}(\bm{\theta},\lambda)$
$\displaystyle=\frac{1}{2}\lvert\lvert\bm{\theta}-\tilde{\bm{\theta}}\rvert\rvert^{2}+\lambda(\bm{1}^{\top}\bm{\theta}-S)$
(9)
$\displaystyle=\frac{1}{2}\lvert\lvert\bm{\theta}-(\tilde{\bm{\theta}}-\lambda\bm{1})\rvert\rvert^{2}+\lambda(\bm{1}^{\top}\tilde{\bm{\theta}}-S)-\frac{n}{2}\lambda^{2}\,\text{,}$
(10)

where $\lambda\geq 0$ and $0\leq\theta_{i}\leq 1$. Minimizing w.r.t.
$\bm{\theta}$ results in

$\bar{}\bm{\theta}=\bm{1}_{\tilde{s}-\lambda\bm{1}\geq
1}+(\tilde{s}-\lambda\bm{1})_{1>\tilde{s}-\lambda\bm{1}>0}\,\text{.}$ (11)

Thus, for $\lambda\geq 0$

$\displaystyle g(\lambda)$
$\displaystyle=\mathfrak{L}(\bar{}\bm{\theta},\lambda)$ (12)
$\displaystyle=\frac{1}{2}\lvert\lvert[\tilde{s}-\lambda\bm{1}]_{-}+[\tilde{\bm{\theta}}-(\lambda+1)\bm{1}]_{+}\rvert\rvert^{2}+\lambda(\bm{1}^{\top}\tilde{\bm{\theta}}-\bm{\theta})-\frac{D}{2}\lambda^{2}$
(13)
$\displaystyle=\frac{1}{2}\lvert\lvert[\tilde{s}-\lambda\bm{1}]_{-}\rvert\rvert+\frac{1}{2}\lvert\lvert[\tilde{\bm{\theta}}-(\lambda+1)\bm{1}]_{+}\rvert\rvert^{2}+\lambda(\bm{1}^{\top}\tilde{\bm{\theta}}-\bm{\theta})-\frac{D}{2}\lambda^{2}$
(14)

and

$\displaystyle g^{\prime}(\lambda)$
$\displaystyle=\bm{1}^{\top}[\lambda\bm{1}-\tilde{\bm{\theta}}]_{+}+\bm{1}^{\top}[(\lambda+1)\bm{1}-\bm{\theta}]_{-}+(\bm{1}^{\top}\tilde{\bm{\theta}}-\bm{\theta})-D\lambda$
(15)
$\displaystyle=\bm{1}^{\top}\min(\bm{1},\max(\bm{0},\tilde{\bm{\theta}}-\lambda\bm{1}))-S$
(16)
$\displaystyle=[\textstyle\sum_{i=1}^{D}\min(1,\max(0,\tilde{\theta_{i}}-\lambda))]-S\,\text{.}$
(17)

With $g^{\prime}(\lambda)$ being a monotone function, $\lambda_{1}^{*}$ a
solution for $g^{\prime}(\lambda)=0$ can be obtained by e.g. a convex solver
or a bisection method. The maximum of $g(\lambda)$ with $\lambda\geq 0$ is at
$\lambda_{2}^{*}=\max(0,\lambda_{1}^{*})$. Eventually,

$\displaystyle\bm{\theta}^{*}$
$\displaystyle=\bm{1}_{\tilde{s}-\lambda_{2}^{*}\bm{1}\geq
1}+(\tilde{s}-\lambda_{2}^{*}\bm{1})_{1>\tilde{s}-\lambda_{2}^{*}\bm{1}>0}$
(18)
$\displaystyle=\min(\bm{1},\max(\bm{0},\tilde{\bm{\theta}}-\lambda_{2}^{*}\bm{1})$
(19)
$\displaystyle=\min(\bm{1},\max(\bm{0},\tilde{\bm{\theta}}-\max(0,\lambda_{1}^{*})\bm{1})\,\text{.}$
(20)