Analysis of Loss Functions for Colorectal Polyp Segmentation Under Class Imbalance ^†

Abstract

Class imbalance is a persistent limitation in polyp segmentation, commonly resulting in biased predictions and reduced accuracy in identifying clinically relevant structures. This study systematically evaluated 12 loss functions, including standard, weighted, and compound formulas, applied to colon polyp segmentation using the UNet-VGG16 fixed architecture on the Kvasir-SEG dataset. The encoder was frozen to isolate the effect of loss functions under the same training conditions. A fixed random seed was used in all experiments to ensure reproducibility and control variance during training. The results reveal that the combined loss functions, namely WBCE combined with Dice and Tversky combined with Focal, achieved the top Dice scores of 0.8916 and 0.8917, respectively. Tversky plus Focal also provided the highest sensitivity of 0.8885, and WBCE obtained the best average IoU of 0.8120. Tversky loss showed the lowest error rate of 4.99, indicating stable optimization. These results clarify the influence of loss function selection on segmentation performance in scenarios characterized by considerable class imbalance.

1. Introduction

Artificial intelligence (AI) has become increasingly prominent as an effective means for alleviating diagnostic burden through advanced segmentation techniques in medical imaging, especially with the growing burden on healthcare systems [1]. Medical image segmentation represents a crucial computer vision task aimed at automatically distinguishing anatomical and pathological patterns from images [2]. Through advances in computation vision and deep learning (DL) technologies, polyp segmentation has gained attention as a pivotal area of research [3]. However, due to a number of issues inherent in medical data, accurate segmentation remains a challenge despite the continued progress of DL models. The main issue is class imbalance, where the foreground region covers a tiny fraction of the overall image [4]. This challenge is particularly evident in the segmentation of colonic polyps, where the morphologic variability of polyps, including differences in shape, size, and texture, combined with limited visibility on endoscopic images, makes them difficult to identify [5]. The similar issue arises in the segmentation of pulmonary nodules when nodules constitute only a small fraction of the total CT volume, resulting in an imbalance in the distribution of positive and negative samples [6]. Segmentation of COVID-19-affected lung lesions faces similar challenges, as infected areas are often characterized by heterogeneous texture, indistinct boundaries, and considerable variability in shape and dimension [7]. Recently, deep learning techniques have started to play an important role in solving these segmentation problems by efficiently extracting features from medical images. Fully convolutional networks (FCNs) [8] UNet [9] and SegNet [10], commonly applied in semantic medical segmentation tasks, have shown a competitive performance [11]. However, a key challenge in training these models remains the large class imbalance [12]. In many medical imaging tasks, foreground structures of interest are often less than 3–10% of the whole image pixel count. To overcome this problem, loss functions serve as a fundamental component in training DL models for segmentation [13]. In particular, they can be modified to penalize the misclassification of underrepresented classes more strongly, thereby guiding the model to better train on rare but clinically important foreground regions. The performance of automatic segmentation under class imbalance depends on the loss function [14], which determines the optimization objective and gradient distribution. The authors of [15] proposed a lightweight DL model based on MobileNetV3 and DeepLabV3+ for the segmentation of colorectal polyps, their study aimed at optimizing the model architecture. In [16], a transformer-based model was proposed for segmenting colon polyps using a hybrid loss function combining the Tversky focal function, binary cross-entropy and Jaccard index, and the method obtained a Dice score of 0.9048. Standard loss functions often favor the majority of classes, reducing accuracy in minority regions. To address this challenge, weighting factors are used in modified functions [17]. However, comparative analyses of loss functions under class imbalance remain limited in current studies.

This research conducts a controlled comparison of 12 loss functions using a fixed model and sequential random sampling to evaluate their performance in segmentation problems with imbalance.

2. Materials and Methods

The study used the Kvasir-SEG dataset [18] consisting of 1000 endoscopic colorectal images with polyp masks standardized to 256 × 256 pixels. To expand the data set to 1500 images, augmentation method was applied [19]. A key challenge was the large variability in polyp size of around 3–11% of the image area, creating a class imbalance. A fixed random seed was applied to all stochastic components of the experimental process [20], including augmentation operations, data partitioning, and initialization of model parameters.

The model in this study, UNet-VGG16, employs a UNet architecture with a pre-trained VGG16 encoder [21] that is fixed during training for hierarchical feature representation. The encoder consists of successive convolutional layers with increasing depth [22].

Binary cross-entropy loss (BCE) determines the negative logarithmic probability of the true class over the predicted probability distribution [23].

$$l_{BCE}=-{\displaystyle \frac{1}{N}}{\displaystyle \sum _{i=1}^N[y_i\log p_i+(1-y_i)\log (1-p_i)]}$$

(1)

Weighted binary cross entropy (WBCE) modifies the standard BCE by introducing a class-dependent weighting factor [24]. It increases the contribution of the minority class to the loss function.

$$l_{WBCE}=-{\displaystyle \frac{1}{N}}{\displaystyle \sum _{j=1}^N[w\cdot y_i\log p_i+(1-w)·(1-y_i)\log (1-p_i)]}$$

(2)

In (2), w ∈ [0, 1] serves as a weighting factor amplifying the contribution of the positive class. This formulation penalizes errors on underrepresented regions more heavily [25].

$$l(p,y)=\left\{\begin{array}{l}-\log (p)ify=1\\ -\log (1-p)ify=0\end{array}\right\}$$

(3)

A generalized form using the probability associated with the correct class is introduced in (4):

$$p_t=\left\{\begin{array}{ll}p& ify=1\\ 1-p& ify=0\end{array}\right\}$$

(4)

In focal loss, a modulating factor is applied ${(1-p_t)}^{\gamma }$ to reduce the losses attributed to highly classified samples:

$$l_{focal}=-\alpha \cdot {(1-p_t)}^{\gamma }·\log (pt)$$

(5)

In (5), $\alpha \in \left[1,0\right]$ serves as a scaling factor to account for differences in class distribution, while $\gamma \ge 0$ regulates the intensity of loss suppression from correctly predicted instances.

Tversky loss is a generalization of Dice by incorporating asymmetric weights for false positives and false negatives [26], which allows for flexibility in managing the accuracy-recall trade-off in unbalanced segmentation tasks.

Surface loss [27] is a boundary loss function bounded by the interval [0, 1] that is computed using precomputed class weights to account for imbalance and does not change the optimization problem at the batch level. It is based on the distance to the object edge.

$$l_{Surface}=1-{\displaystyle \frac{{\displaystyle {\sum }_{c=1}^C{w}_c{\displaystyle {\sum }_{n=1}^N{d}_n^{(c)}{(1-t_n^{(c)}+p_n^{(c)})}^2}}}{{\displaystyle {\sum }_{c=1}^{C}w{\displaystyle {\sum }_{n=1}^N{(d_n^{(c)})}^2}}}},$$

(6)

where C refers to the number of classes, N is the total number of spatial elements, the coefficient $w_c$ is the weight assigned to class c, $d_n^{(c)}$ denotes the boundary distance [28] from pixel n to the class c, the term $t_n^{(c)}$ represent the binary ground truth label, and $p_n^{(c)}$ is the model prediction.

Mean squared error (MSE) [29] loss for segmentation quantifies variations across pixels by penalizing the mean square of the residuals between the estimated results and the corresponding reference annotations across all of the categories. In segmentation tasks with class imbalance, it enables equal contribution from all classes by updating predictions for both true and false classes.

$$l_{MSE}={\displaystyle \frac{1}{2N}}{\displaystyle \sum _{n=1}^N{\displaystyle \sum _{c=1}^C{(p_n^{(c)}-t_n^{(c)})}^2}},$$

(7)

where N denotes the cardinality of the pixel space, C equals the total number of semantic classes, and $p_n^{(c)}$ represents the model-assigned probability that a pixel n and belongs to the class c, $t_n^{(c)}$.

Mean absolute error (MAE) [30] loss quantifies the average magnitude of residuals between the predicted outputs and target values.

$$l_{MAE}={\displaystyle \frac{1}{N}}{\displaystyle \sum _{n=1}^N{\displaystyle \sum _{c=1}^C\left|p_n^{(c)}-t_n^{(c)}\right|}},$$

(8)

where N denotes the cardinality of the pixel space, C equals the total number of semantic classes, and $p_n^{(c)}$ represents the model-assigned probability that a pixel n and belongs to the class c, $t_n^{(c)}$.

Jaccard loss [31] estimates the discrepancy between the predicted results and reference masks through direct minimization of the discrepancy based on IoU.

$$l_{Jaccard}=1-{\displaystyle \frac{{\displaystyle {\sum }_{n-1}^N{\displaystyle {\sum }_{c=1}^C{p}_n^{(c)}\cdot t_n^{(c)}}}}{{\displaystyle {\sum }_{n=1}^N{\displaystyle {\sum }_{c=1}^C(p_n^{(c)}+t_n^{(c)}-p_n^{(c)}\cdot t_n^{(c)})}}}},$$

(9)

where N denotes the cardinality of the pixel space, C equals the total number of semantic classes, and $p_n^{(c)}$ represents the model-assigned probability that a pixel n and belongs to the class c, $t_n^{(c)}$.

Dice loss is a region-based loss function [32] that measures the overlap between predicted and true segmentation masks.

$$l_{Dice}=1-{\displaystyle \frac{2{\displaystyle {\sum }_{n=1}^N{p}_n+{\displaystyle {\sum }_{n=1}^N{t}_n+\epsilon }}}{{\displaystyle {\sum }_{n=1}^N{p}_n+{\displaystyle {\sum }_{n=1}^N{t}_n+\epsilon }}}},$$

(10)

where N denotes the cardinality of the pixel space, C equals the total number of semantic classes, and $p_n^{(c)}$ represents the model-assigned probability that a pixel n and belongs to the class c, $t_n^{(c)}$.

Dice + BCE loss is a combined loss function that combines region-based and pixel-based supervision.

$$l_{Dice-BCE}=\alpha \cdot l_{BCE}+(1-\alpha )\cdot l_{Dice}$$

(11)

WBCE + Dice loss addresses class imbalance by weighting pixel-wise binary cross-entropy and combining it with region-based Dice loss.

$$l_{WBCE+Dice}=\alpha \cdot l_{WBCE}+(1-\alpha )\cdot l_{Dice}$$

(12)

Tversky + Focal loss combines the Tversky index, which controls for false positives and false negatives, with Focal loss, which emphasizes hard samples.

$$l_{Tversky+Focal}=\alpha \cdot l_{Tversky}+(1-\alpha )\cdot l_{Focal}$$

(13)

Various measures, primarily classification accuracy and spatial correspondence were used to assess the effectiveness of the model. These include Dice score, mean IoU and accuracy. The Dice score metric estimates the spatial coverage overlap between predicted (P) and reference areas (G), giving an indication of segmentation consistency [33].

Mean IoU estimates the average overlap across all classes, offering a balanced metric.

$$MeanIOU={\displaystyle \frac{1}{N}}{\displaystyle \sum _{i=1}^N\frac{\left|P_i\cap G_i\right|}{P_i\cup G_i}}$$

(14)

Sensitivity, also known as recall, determines the model capacity to distinguish the relevant class [34]. Accuracy measures the proportion of correctly classified pixels relative to the total number of pixels [35,36].

The experiments were run on an A100 GPU (40 GB) with Python 3.10 and TensorFlow 2.13.

3. Results

The quantitative analysis presented in

Table 1. Comparative evaluation of loss functions based on validation metrics for image segmentation.

Loss	Accuracy	Dice	Mean IOU	Sensitivity
BCE	0.9553	0.8882	0.8104	0.8576
WBCE	0.9557	0.8831	0.8120	0.8593
Tversky	0.9546	0.8892	0.8059	0.8438
Surface	0.9525	0.8847	0.7982	0.8616
MSE	0.9543	0.8727	0.8066	0.8516
MAE	0.9518	0.8819	0.7953	0.8344
Jaccard	0.9529	0.8857	0.7996	0.8532
Focal	0.9552	0.7963	0.8102	0.8483
Dice	0.9539	0.8883	0.8052	0.8466
Dice + BCE	0.9538	0.8843	0.8014	0.8402
WBCE + Dice	0.9557	0.8916	0.8104	0.8498
Tversky + Focal	0.9540	0.8917	0.8096	0.8885

demonstrates the comparative effectiveness of 12 different loss functions evaluated on a validation subset of 150 images.

In terms of accuracy, WBCE and WBCE + Dice showed the highest value of 0.9557, which is an increase of 0.04 % compared to BCE at 0.9553. This minor improvement is due to the class weighting mechanism in WBCE. For the Dice coefficient, the combined Tversky + Focal and WBCE + Dice losses attained values of 0.8917 and 0.8916, respectively. These results are 0.39 and 0.38 % higher than the BCE and 0.34 and 0.33 % higher than the standard Dice losses. In terms of mean IOU, WBCE performed the highest at 0.8120, which is 0.16% higher than BCE at 0.8104 and 1.67% higher than MAE at 0.7953. MAE and MSE equally accounted for pixel-by-pixel deviations, without considering the class distribution. The highest sensitivity was observed for Tversky + Focal at 0.8885. This corresponds to an increase of 3.09% over Dice at 0.8466 and 1.18% over BCE at 0.8576. The improvement in sensitivity is explained by the parameters α and β of the Tversky index, which control the compromise between false-positive and false-negative results. The error comparisons for each loss function are summarized below in

Table 2. Assessment of stability of loss functions by difference and ratio of errors on validation.

Loss	Validation	Difference	Error Ratio
BCE	0.2020	0.1935	23.6540
WBCE	0.0992	0.0932	16.5767
Tversky	0.1093	0.0874	4.9933
Surface	0.0381	0.0329	7.3658
MSE	0.0312	0.0290	13.6952
MAE	0.0388	0.0320	5.6681
Jaccard	0.1798	0.1243	3.2367
Focal	0.0147	0.0141	24.77
Dice	0.1105	0.0940	6.6742
Dice + BCE	0.1385	0.1266	11.6415
WBCE + Dice	0.2020	0.1935	23.6540
Tversky + Focal	0.1281	0.1200	15.9704

, including validation loss values, absolute differences, and relative error rates.

As shown in Table 2, BCE produced the highest validation loss of 0.2020 and an error ratio of 23.65, indicating overfitting and limited generalization. WBCE reduced the validation loss by 50.9% and the error ratio by 7.08 points, demonstrating the advantage of class overfitting. Tversky achieved a lower validation loss of 0.1093 and the lowest error ratio among single losses, 4.99, reflecting stable learning under class imbalance. The surface yielded a validation loss of 0.0381 and an error ratio of 7.36. MSE and MAE resulted in low validation losses of 0.0312 and 0.0388, with error rates of 13.69 and 5.67, respectively. These values indicate good convergence, but limited structural representation. The validation loss by Jaccard’s criterion reached 0.1798, with an error rate of 3.23. Focal loss produced the lowest validation loss of 0.0147, but the highest error ratio of 24.77, suggesting sensitivity to complex samples and potential instability. Dice and Dice + BCE showed validation losses of 0.1105 and 0.1385, with error ratios of 6.67 and 11.64, respectively, suggesting a balance between regional and pixel-by-pixel learning. The combined Tversky + Focal loss provided 0.1281 validation loss and an error rate of 15.97, improving the stability performance over BCE-based losses.

Training time, average epoch duration, and frames per second (FPS) are provided in

Table 3. Assessment of computational efficiency of loss functions in terms of training time and processing speed.

Loss	Training	Epoch	FPS
BCE	527.60	10.55	153.52
WBCE	524.68	10.49	154.38
Tversky	523.97	10.48	154.59
Surface	559.32	11.19	144.82
MSE	524.75	10.50	154.36
MAE	526.22	10.52	153.93
Jaccard	554.64	11.09	146.04
Focal	524.57	10.49	154.41
Dice	528.67	10.57	153.22
Dice + BCE	547.41	10.95	147.97
WBCE + Dice	535.44	10.71	151.28
Tversky + Focal	542.83	10.86	149.22

.

Tversky showed the lowest training time of 523.97 s, followed by WBCE, MSE, and Focal at around 524 s. These features also showed high FPS values, with Tversky’s reaching 154.59, indicating the efficiency of executing a single frame. The longest training time was recorded for Surface at 559.32 s, the epoch duration was 11.19 s, and the lowest FPS was 144.82. The training time of Jaccard and Dice + BCE functions also exceeded 547 s, indicating a higher computational load, probably due to their structural operations over the prediction and truth regions. Combined loss functions such as WBCE + Dice and Tversky + Focal showed moderate training times of 535.44 and 542.83 s, with FPS values of 151.28 and 149.22, respectively.

The validation results are plotted in Figure 1, allowing for a comparison of the stability of convergence under identical training conditions.

The diagram in Figure 1 illustrates that most loss functions achieve near-optimal validation performance within the first 15–20 epochs. However, notable fluctuations are observed in certain functions. There is a distinct peak in the Jaccard loss curve around epoch 30, accompanied by a temporary deterioration in the validation results for the other metrics. Likewise, Dice + BCE exhibits increased variability between epochs 25 and 35, especially in the sensitivity and Dice coefficient curves. In contrast, loss functions such as WBCE, Tversky, and WBCE + Dice maintain a consistent trajectory without sharp deviations after epoch 20. The relationship between the validation loss and Dice coefficient for all estimated loss functions is depicted in Figure 2.

The chart in Figure 2 indicates a weak positive correlation between validation loss and Dice score with an estimated correlation coefficient of 0.32. Loss functions such as Tversky + Focal and WBCE + Dice achieve high Dice coefficient scores despite higher validation losses. In contrast, MAE, Surface, and MSE show low loss values, but moderate Dice results. Focal loss is considerably different, having the lowest loss and the lowest Dice coefficient score, confirming that loss mitigation does not directly lead to segmentation accuracy. The correlation between validation loss and IoU for the estimated loss functions is shown in Figure 3, providing a comparative analysis of segmentation consistency relative to error magnitude.

The diagram in Figure 3 demonstrates a weak positive correlation between validation loss and IoU, with an estimated correlation coefficient of 0.21. Loss functions such as WBCE, Focal, and WBCE + Dice achieve relatively high IoU values despite the differences in validation loss, indicating consistent segmentation at the region level. In contrast, MAE and Surface exhibit low validation losses, but the lowest IoU values.

The qualitative results in Figure 4 and Figure 5 illustrate the influence of different loss functions on the fixed model under conditions of class imbalance and structural similarity in polyp segmentation.

Most loss functions provide consistent localization, although there are slight differences in the accuracy of boundary detection and the number of false positives. Loss functions such as WBCE, Dice, WBCE + Dice and Tversky + Focal demonstrate more accurate localization. These features allow for better separation of the polyp from the mucosal background compared to the standard BCE or Focal.

Figure 5 provides a second sample of segmentation under severe class imbalance, illustrating how different loss functions cope with the complex morphology of polyps and their separation from surrounding tissues.

Figure 5 presents a second qualitative sample with a greater class imbalance compared to the first sample, with a smaller and less defined polyp area. Features such as Dice, Tversky, WBCE + Dice, and Tversky + Focal show relatively stable localization and shape preservation, although minor over-segmentation and imprecise boundaries are observed. In contrast, BCE, MSE, and Focal produce noticeable false positives and fragmentation. Compared to the first example in Figure 4, where most loss functions produced consistent predictions, this second case shows the limitations of some functions in scenarios with extreme foreground sparsity and morphological complexity.

4. Conclusions

Segmentation of polyp images commonly encounters the problem of class imbalance, where anatomically important foreground structures take up a minimal proportion of the image area compared to the background. This imbalance challenges the effectiveness of conventional loss functions that weight all pixel contributions uniformly and, hence, bias the model optimization towards majority classes. Addressing this, class-sensitive loss function formulas that incorporate weighting mechanisms have been developed to enhance the influence of underrepresented regions during training. In the present study, 12 loss functions, including standard, weighted, and composite variants, were rigorously evaluated under the UNet-VGG16 fixed architecture on the Kvasir-SEG dataset. Special attention has been paid to studying the interaction between pixel accuracy and region-level coherence under imbalanced conditions. Notably, WBCE introduces class-dependent modulation, and composite formulas such as WBCE plus Dice and Tversky plus Focal combine additional targets to refine the gradient behavior. Empirical results showed that Tversky plus Focal achieved the highest Dice coefficient of 0.8917 and sensitivity of 0.8885, while WBCE achieved the highest average IoU of 0.8120. Tversky losses showed excellent stability of the optimization with the lowest error rate of 4.99. However, the analysis was limited to the Kvasir-SEG dataset and the fixed UNet-VGG16 architecture with a frozen coder, which may limit generalizability. The impact of adaptive weighting and dynamic loss scheduling was not investigated. The study’s limitations include the use of a single, homogeneous set of endoscopic images and a fixed UNet VGG16 configuration with a frozen coder, which restricts the generalizability of the results to other imaging protocols and architectures. Additionally, the hyperparameters of the loss functions were not optimized through advanced search, and the evaluation was conducted without multi-round stratified testing or statistical testing of differences. Future work will involve using multicenter collections with different imaging conditions, implementing domain adaptation to remove bias in the distributions, and testing the same loss functions in hybrid convolutional models of varying depth.