Selecting the Best Optimizers for Deep Learning based Medical Image Segmentation

SGD and Mini-batch gradient descent were first optimizers used for training neural networks. The updating rule for these optimizers include only the value of last iteration as shown in Eq. 1. Choosing appropriate value for a LR is challenging in these optimizers since if LR is very small then convergence is very slow; and if LR is set high, the optimizer will oscillate around global minima instead of converging:

where $\theta$ is network parameters, $~{}\alpha$ is LR, J is cost function to be minimized (function of $\theta$, $X$(input), and $Y$(labels)). The equation 1 can be considered as an updating rule for SGD and mini-batch gradient descent by choosing $X$ and $Y$ as whole samples, a single sample, or a batch of samples in a dataset.

The Momentum optimizer was designed to accelerate the optimization process by taking into account the values from previous iterations, weighted by a factor known as ”momentum,” as mentioned in [5]. The updating for this optimizer is defined as:

where $~{}\beta$ denotes the momentum rate (MR). In Momentum optimizer, the past iterations don’t play any role in cost function and cost function is only calculated for the current iteration only. Also, similar to LR, choosing a proper value for MR is challenging and it has a correlation with LR too.

Nesterov accelerated gradient [9] (NAG) was then introduced to address the limitation of momentum optimizers as well as to accelerate the convergence by including information from previous iterations in calculating the gradient of the cost function as shown in the following equation:

Compared to optimizers with fixed LR/MR, the NAG optimizer generally shows improved performance in both convergence speed and generalizability.

A significant disadvantage of optimizers with a fixed LR/MR is that they cannot incorporate information from the gradients of past iterations in adjusting the learning and momentum rates. For instance, they cannot increase the learning rate for dimensions with a small slope to improve convergence, or reduce the learning rate for dimensions with a steep slope to avoid oscillation around the minimum point. Adagrad [3] is one of fist adaptive LR optimizers used in deep networks adapting the learning rate for each parameter in the network by dividing the gradient of each parameter by its sum of the squares of gradient, as follows:

where $G_{i}$ is a diagonal (square) matrix and each diagonal element equal to the sum of the square of gradient of its corresponding parameters:

where  I is the current iteration.

One of the drawbacks of AdaGrad is gradient vanishing due to accumulation of all past square gradient in denominator of Equation 6 during the training. This leads the gradients to converge to zero after several epochs in training. However,  AdaDelta,  RMSProp, and  ADAM optimizers solved this problem by considering a sum of the past samples within a pre-defined window. ADAM optimizer’s updating rule uses past squared gradient (as scale) and also like momentum, it keeps an exponentially decaying average of past gradients. Hence, these adaptive optimizers have advantages over the AdaGrad by adaptively changing both RL and ML as well as resolving the gradient vanishing issue:

Adaptive learning methods are costly because they are required to calculate and keep all the past gradients and their squares to update the next parameters. Also, the adaptive learning optimizer may converge into different minima in comparison with fixed learning rate optimizers [7, 6, 8].

Alternatively, Cyclic learning rate (CLR) was proposed to change the learning rate during training, which needed no additional computational cost. CLR is a method for training neural networks that involves periodically changing the learning rate during training. As mentioned earlier, the learning rate is typically adjusted according to a predetermined schedule, such as increasing the learning rate from a low value to a high value and then decreasing it back to the low value over a set number of training iterations. The learning rate is then reset and the process is repeated. This can help the optimization process by allowing the model to make larger updates at the beginning of training and smaller updates as training progresses, potentially leading to faster convergence and better model performance [7]. Later in Figure 3a, we show how we use CLR in our methodology.

Cardiovascular diseases (CVDs) are the leading cause of death worldwide according to the World Health Organization (WHO). CVDs lead to millions of deaths annually and are expected to cause over 23.6 million deaths in 2030 [10]. Cine-MR imaging can provide valuable information about cardiac diseases due to its excellent soft tissue contrast. For example, ejection fraction (EF), an important metric measuring how much blood the left ventricle pumps out with each contraction, can be measured with Cine-MRI. To this end, radiologists often manually measure the volume of the heart at the end of the systole (ES) and the end of the diastole (ED) to measure EF. This is a time-consuming process with known inter-, and intra-observer variations. Due to its significance in functional assessment of heart, there have been numerous machine learning based automated algorithms developed in the literature for measuring EF. In this study, we make our efforts in this application due to its importance in the clinic.

There is a considerable amount of research dedicated to the problem of cardiac segmentation from MR or CT images. Since Xu et al. found a correlation between motion characteristics and tissue properties, they developed a combined motion feature learning architecture for distinguishing myocardial infarction [11]. In our another attempt, CardiacNet in [12] proposed a multi-view CNN to segment the left atrium and proximal pulmonary veins from MR images following by an adaptive fusion. The shape prior information from deep networks were used to guide segmentation network to delineate cardiac substructures from MR images [13, 14]. As previously stated, the literature and methodologies for cardiac segmentation are extensive. Readers are invited to consult references [15] and [16] for more comprehensive information.

Within the DB, a concatenation operation is done for combining the feature maps (through direction (axis) of the channels) for the last three layers. So, if the input to $\it{l^{th}}$ layer is $\mathbf{X}_{\it{l}}$, then the output of $\it{l^{th}}$ layer is:

Since we are doing concatenation before each layer (except the first one), the output of each layer can be calculated only by considering the input and output of first layer as following:

where ^⌢ is concatenation operation. In addition, for initialization $F(\mathbf{X}_{\it{-1}})$ and $F(\mathbf{X}_{\it{0}})$ are considered as {} and $\mathbf{X}_{\it{1}}$ respectively which {} is an empty set and there are $\it{L}$ layers inside of the block.

Assuming the number of output features for each layer is $\mathbf{K_{out}}$ (channel out) and the number of input features for first layer is $\mathbf{K_{in_{1}}}$ (channel in). Then, the feature maps growth (channel out) for second, third, …, and $\it{L^{th}}$ layer are $\mathbf{K_{out}}+\mathbf{K_{in_{1}}}$, 2$\mathbf{K_{out}}+\mathbf{K_{in_{1}}}$, …, and $(\it{L}-1)\mathbf{K_{out}}+\mathbf{K_{in_{1}}}$ respectively. The growth rate for the DB is the same as fourth layer.

Smith et al [7] argued that a cycling learning may be a more effective alternative to adaptive optimizations especially from generalization perspective. Basically, cyclic learning includes a pre-defined cycle (such as triangle or Gaussian function) that learning rate is changing according to that cycle. Here, we hypothesize (and show later in the results section) that having a cyclic momentum in Nesterov optimizer (Eq. 2) can lead to a better accuracy in segmentation task in generalization phase. As a reminder, momentum in Eq. 2 was used to consider the past iterations by a coefficient called momentum. So, choosing the proper value for momentum is challenging. To this end, we propose changing the MR in the same way that we changed the LR, and we considered the cyclic triangle function for both MR and LR as illustrated in the Figure 3. $cycle_{lr}$ and $cycle_{mr}$ determine the period of triangle function for LR and MR are defined by:

where $C_{lr}$ and $C_{mr}$ are positive even integer numbers, It is number of iteration per each epoch.

In Figures 3a and 3b, the cyclic function for different values of $C_{lr}$ and $C_{mr}$ are illustrated. LR during whole training can be determined from equation 11:

where $max_{lr}$ and $min_{lr}$ are maximum and minimum values of LR function, respectively. $i$ is the iteration indicator during whole training process and $i\in\{1,2,\dots,It\times Ep\}$, which $Ep$ is total number of epochs in training and $N$ is a set of natural number. MR can also be determined as:

where $max_{mr}$ and $min_{mr}$ are maximum and minimum values of MR function, respectively.

Equations 11 and 12 are used to determine the values of LR and MR in each iteration during training. One of the challenges in using these cyclic LR and MR functions are determining the values of some variables in the equations including $max_{lr}$, $min_{lr}$, and $C_{lr}$ for LR; and also $max_{mr}$, $min_{mr}$, and $C_{mr}$ for MR. For finding the the $max_{lr}$ and $min_{lr}$ values, as it suggested in [7], one can run the networks with different LR values for a few epochs and then these values are chosen according to how network accuracy changes. Since, when both LR and MR change dynamically and the one value can affect the other one (considering the optimizer formula), it makes more challenging to find the CLMR optimum parameters by proposed solution. It means we need to train large number of networks in order to determine the optimum values of $max_{lr}$, $min_{lr}$, $max_{,r}$, and $min_{mr}$ which is not computationally feasible. Also, a heuristic method was suggested in [7] to find the best value of $C_{lr}$.

In this paper we propose an alternative way to find best cyclic functions with minimum computational cost. We set fixed values for $max_{lr}$, $min_{lr}$, $max_{,r}$, and $min_{mr}$ parameters and make sure that the selected values cover a good range of values for both LR and MR in practice (illustrated in Figure 3). Then, we did a computationally reasonable heuristic search for finding the appropriate amount of $C_{lr}$ and $C_{mr}$ from the values shown in Figure 3. Since, changing the values of $C_{lr}$ and $C_{mr}$ leads to change in the values of LR and MR in different iterations, there is no need to find the optimum values for minimum and maximum, we did search in 2D space of $C_{lr}$ and $C_{mr}$ to find their optimal values.

For investigating the performance of proposed method, a dataset from Automatic Cardiac Diagnosis Challenge (ACDC-MICCAI Workshop 2017) were used [27]. These data set includes 150 cine-MR images: 30 normal cases, 30 patients with myocardial infarction, 30 patients with dilated cardiomyopathy, 30 patients with hypertrophic cardiomyopathy, and the remaining 30 patients with abnormal RV. While 100 cine-MR images were used for training (80) and validation (20), the remaining 50 images were used for testing with online evaluation by the challenge organizers. For a fair validation in training procedures, four subjects from each category have been chosen. The binary masks for ground truths of three substructures were provided by the challenge organizers for training and validation while test set was evaluated online (unseen test set. Three substructures are right ventricle (RV), myocardium of left ventricle(Myo.), and left ventricle (LV) at two time points of end-systole (ES) and end-diastole (ED).

The MRIs were obtained using two MRI scanners of different magnetic strengths (1.5T and 3.0T). Cine-MR images were acquired with a SSFP sequence in short axis while on breath hold (and gating). In particular, a series of short axis slices cover the LV from the base to the apex, with a thickness of 5 mm (or sometimes 8 mm) and sometimes an inter-slice gap of 5 mm. The spatial resolution goes from 1.37 to 1.68 $mm^{2}/pixel$ and 28 to 40 volumes cover completely or partially the cardiac cycle.

The networks were trained for a fixed number of epochs (100) and it was confirmed that they are fully trained. All the images were resized to $200\times 200$ in short axis by using B-spline interpolation. Then, as a preprocessing step, we applied anisotropic filtering and histogram matching to the whole data set. The total number of 2D slices for training was about 1690 and batch size of 10 were chosen for training. Hence, the number of iteration per epoch is $\frac{1690}{10}=169$ and we have a total number of iteration $100\times 169=16,900$ in training. The Cross Entropy loss function was chosen for minimization. All the networks were implemented on Tensorflow with using NVIDIA TitanXP GPUs.

We calculated Dice Index (DI) and also Cross Entropy (CE) loss on validation set for investigating our proposed optimizer along with other optimizers. In Figures 5a and b, the CE and DI curves versus iterations for U-Net architecture for different optimizers are illustrated. As these curves show, the DI in U-Net with ADAM optimizer is increasing rapidly and sharply at the very beginning and then it is almost fixed afterwards. Although our proposed optimizer (CLMR(C_lr=20, C_mr=20)) is not learning as fast as ADAM optimizer at very beginning in U-Net, it gets better accuracy than ADAM finally. This phenomenon is clearer in CE curves. The quantitative results on test set in Table 2 support the same observation and conclusion. Further, the same pattern happens for DenseNet_2 architecture in Figure 6. This confirms our hypothesis that adaptive optimizers converges faster but to potentially to different local minimas in comparison with classical SGD optimizers.

Figure 5 shows the behavior of U-Net architecture with CLMR optimizer performing 2% increase in dice index (in all three substructures as well as average) than its CRL optimizer counterpart. This proves that having a cyclic momentum rate can yield to better efficiency and accuracy than having a simple cyclic learning rate. The results on the test set comparing CLR and CLMR optimizers in Table 2 support this conclusion too.

Moreover, the curves of DI and CE among different architectures, trained by ADAM and CLMR, are demonstrated in Figure 4 a and b. Although the DenseNet_2 has less parameters in comparison with other architectures, it gets better results than the other architectures. These curves reveals some other important points about using different architectures: first, for all different architectures, the proposed CLMR optimizer works better than ADAM optimizer, indicating the power of proposed cyclic optimizer. Second, DenseNet architectures are getting better results than U-Net and Enc_Dec architectures, which are highly over-parameterized architectures than DenseNet and their saturation can be linked to this information too. Third, comparison between curves of DenseNet_1 and DenseNet_2 shows that having a higher GR (growth rate) in dense connections is more important than having dense block with high number of parameters. Since, DenseNet_2, with GR=24 reached better results in comparison with DenseNet_1 with twice of number of parameters in end of each dense block in comparison to DenseNet_2 and GR=16. These results are supported by the dice metric obtained from test data and are mentioned in Table 2.

Finally, the DI on test data with online evaluation for different architectures with different optimizers are summarized in Table 2. In order to have a better comparison, the box plot of all methods are drawn in Figure 7. As the figure shows, the dice statistic obtained from CLMR is better than other optimizers most of the time in addition to its superior efficiency. In addition, qualitative results for different methods are shown in Figures 8 and 9: the contours for RV, Myo., and LV in ED for different methods and architectures and also ground-truth across four slices from Apex to Base. Usually, segmentation of RV near the Apex is harder than others because RV is almost vanishing at this point. As a result, some methods may not even detect the RV at slices near the Apex. Figure 9 shows the contours for RV, Myo., and LV in ES for different methods and architectures and also ground-truth across four slices from Apex to Base. Since at ES heart is at minimum volume; thus, it is more difficult to segment substructures. The contours generated with DenseNet_2 method is more similar to ground-truth in both ED and ES, which shows the generalizability of the proposed method with an efficient architecture choice.

No Conflict of Interest.

This project is supported by NIH funding: R01-CA246704, R01-CA240639, U01-DK127384-02S1, R03-EB032943-02, and R15-EB030356.

Data is available under MICCAI 2017 ACDC challenge.