Dealing with class imbalanced image datasets using the Focal Tversky Loss | by Robin Vinod | Towards Data Science

Class imbalanced datasets is a frequent problem experienced when trying to train segmentation networks. The first time I trained an image segmentation model, I was getting over 96% accuracy straight off the bat. Either I unknowingly made the biggest breakthrough since the inception of the first CNN, or there was something wrong with my approach.
I quickly realised that the pixels to be segmented accounted for a very small percentage of the total pixels in the image. All the model had to do was predict a fully black image (a.k.a no segmentation) and it was rewarded with over 90% accuracy. This is a common problem encountered in most image segmentation tasks, where the background class is much larger than the other classes.
In this story, I go through the techniques used to deal with class imbalanced problems and why the Focal Tverky Loss might be the best option for you.
Losses that deal with class imbalanceWhat is the Tversky Index?Focal Tversky LossConclusion
Losses that deal with class imbalance
What is the Tversky Index?
Focal Tversky Loss
Conclusion
a. Weighted Binary Cross Entropy
Once of the losses typically used to deal with class imbalance is the weighted binary cross entropy. The crux of the normal binary cross entropy is that it considers all pixels equally when calculating the loss. In a mask where 90% of the pixels are 0s and only 10% are 1, the network receives receives a low loss even if it misses all the 1s, which means the network is not learning anything.
Weighted binary cross entropy (WBCE) attempts to solve this by weighting the positive class.
For every class 1 predicted by the network, BCE adds log(p) to the loss while WBCE adds β log(p) to the loss.
Hence, if β > 1, class 1 is weighted higher, meaning the network is less likely to ignore it (lesser false negatives). Conversely, if β < 1, class 0 is weighted higher, meaning there will be lesser false positives.
By controlling, the value of β, you can reduce the problem of class imbalance by weighting the smaller class higher. However, the optimal value of β is hard to ascertain and requires many rounds of trial and error.
Pros: Simple smooth loss surface that is fast in training
Cons: Difficult to optimise and find the sweet spot
b. Dice Coefficient
The Dice Coefficient is well-known for being the go-to evaluation metric for image segmentation, but it can also serve as a loss function. Although not as widely used as other loss functions like binary cross entropy, the dice coefficient does wonders when it comes to class imbalance.
Unlike BCE, dice coefficient only considers the segmentation class and not the background class. The pixels are classified as True Positive (TP), False Negative (FN) and False Positive (FP).
The dice coefficient is a measure of overlap of the predicted mask and the groundtruth. Since it does not account for the background class, it cannot dominate over the smaller segmentation class. The dice coefficient outputs a score in the range [0,1] where 1 is a perfect overlap. Thus, (1-DSC) can be used as a loss function.
Considering the maximisation of the dice coefficient is the goal of the network, using it directly as a loss function can yield good results, since it works well with class imbalanced data by design.
However, one of the pitfalls of the dice coefficient is the potential for exploding gradients. Early in the training, the dice coefficient is close to 0, which can cause instability in training due to exploding gradients as a result of the network making large changes to weights. However, this is easily manageable using batch normalisation and ReLUs.
Pros: Easy solution to class imbalance without having to manually optimise any parameters
Cons: Potential for gradient explosion, and generally slower to train than BCE
The Tversky Index (TI) is a asymmetric similarity measure that is a generalisation of the dice coefficient and the Jaccard index.
The tversky index adds two parameters, α and β, where α + β = 1. In the case where α = β = 0.5, it simplifies into the dice coefficient. It simplifies to the Jaccard index if α = β = 1.
By setting the value of α > β, you can penalise false negatives more. This becomes useful in highly imbalanced datasets where the additional level of control over the loss function yields better small scale segmentations than the normal dice coefficient.
Although the tversky index is only a simple improvement over the dice coefficient, it can prove useful in edge cases where you need a finer level of control.
In the paper, Tversky loss function for image segmentation using 3D fully convolutional deep networks by Salehi et al., the tversky loss was used to obtain the most desirable performance for multiple sclerosis lesion segmentation, a typically class imbalanced problem.
The Focal Tversky Loss (FTL) is a generalisation of the tversky loss. The non-linear nature of the loss gives you control over how the loss behaves at different values of the tversky index obtained.
γ is a parameter that controls the non-linearity of the loss. As γ tends to positive ∞, the gradient of the loss tends to ∞ as the Tversky Index (TI) tends to 1. As γ tends to 0, the gradient of the loss tends to 0 as TI tends to 1.
Essentially, with a value of γ < 1, the gradient of the loss is higher for examples where TI > 0.5, forcing the model to focus on such examples. This behaviour can be useful towards the end of training as the model is still incentivised to learn even though TI is nearing convergence. However, at the same time, it will weight easier examples higher during the early stages of training, which can lead to poor learning.
In the case of class imbalance, the FTL becomes useful when γ > 1. This results in a higher loss gradient for examples where TI < 0.5. This forces the model to focus on harder examples, especially small scale segmentations which usually receive low TI scores.
In the image above, the blue line is the standard tversky loss. The purple line shows the higher gradient and higher loss when TI > 0.5 while the green line shows higher loss when TI < 0.5.
The table above is extracted from the paper, A Novel Focal Tversky loss function with improved Attention U-Net for lesion segmentation.
In the BUS dataset, the FTL showed significant improvements in all categories. In the ISIC dataset, although the overall dice coefficient was the highest, the model did not have the highest precision or recall.
Although models trained with FTL may not have the outright highest precision or recall, it is important to note that the balance between precision and recall was best when trained with FTL, which shows that the objective of tackling class imbalance was satisfied. The end result is the best overall dice coefficient.
The Focal Tversky Loss is an easy drop in solution to deal with class imbalance. Although the best parameter values will take some trial and error to determine, you should see good results with the following:
α = 0.7, β = 0.3, γ = 3/4
Keras Implementation
To get the full implementation of a U-Net with attention, recurrence and inception layers pre-implemented, please check https://github.com/robinvvinod/unet/.