# Multi-Task Learning for Small Brain Tumor Segmentation from MRI
## 3 Materials and Method
### Training Strategy  
In this paper, we used two loss functions for the segmentation task and feature reconstruction task separately. The first loss function is the generalized dice loss (GDL) [56]. Most of the previous methods use dice loss to calculate the difference between the predicted mask and ground truth. However, the distribution of the number of voxels in the tumor regions varied greatly in the brain tumor structure. Therefore, we take advantage of the GDL to solve the imbalanced class problem in the labels of subtumors. The GDL is calculated by the following formula:  
\[L_{GDL}=1-2\frac{\sum_{c=1}^{2}w_{c}\sum_{n}y_{nc}p_{nc}}{\sum_{c=1}^{2}w_{c} \sum_{n}y_{nc}+p_{nc}} \tag{4}\]  
where the \(y_{nc}\) and \(p_{nc}\) are the ground truth tumor mask and the predicted tumor mask of class c, respectively. The class weight \(w_{c}\) was used to provide invariance to the different label set properties.  
We used mean squared error to measure the difference between the reconstructed feature map and the input feature map of each U-module. The loss for the feature reconstruction task was calculated by the following formula:  
\[L_{reconstruction}=L_{recom1}+L_{recom2}=\sum_{i=1}^{n}\left(I_{orig1}-I_{recom 1}\right)^{2}+\sum_{i=1}^{n}\left(I_{orig2}-I_{recom2}\right)^{2} \tag{5}\]  
Figure 1: Our proposed multi-task learning network with two tasks. One task is brain tumor segmentation. We add the additional feature reconstruction task to our model as an auxiliary task to help retain important features when downsampling. The model takes a sequence of MRI modalities as input and generates the brain tumor mask.  
where, the \(L_{recon1}\) and \(L_{recon2}\) are the reconstruction losses of the first and second U-module block. \(I_{orig}\) is the input feature of U-module and \(I_{recon}\) is the reconstructed feature map.  
There were two tasks in our model. The first one is the tumor segmentation task and the second one is feature reconstruction. Therefore, we had two loss functions for optimization. We need to combine two loss functions together. The combined loss function is the sum of the loss functions of the main task along with the auxiliary task as shown in the following:  
\[L_{total}=\lambda_{1}L_{segmentation}+\lambda_{2}L_{reconstruction} \tag{6}\]  
where the \(L_{segmentation}\) is the loss function for our main task (tumor segmentation) and the \(L_{reconstruction}\) is the loss function for our auxiliary task (feature reconstruction). The \(\lambda_{1}\) and \(\lambda_{2}\) are the weight parameters. Since each loss function has its optimization direction and the range value of each loss function is different, we need to set up a training strategy for getting the best optimization. At the first 20 epochs, we set the \(\lambda_{1}\) and \(\lambda_{2}\) equal 0.5 for better parameter initialization. After, we set the main task weight the \(\lambda_{1}\) = 1 and auxiliary task weight \(\lambda_{2}\) = 0.05 so the network can focus on the main task.  
The reconstruction task was used to retain important features and help the network converge better in the training phase. During the inference phase, the reconstruction task was removed. We only used the main task to generate segmented tumor masks.