Unexpected behavior of Dice Loss on extremely imbalanced binary segmentation

username: ILoveSorasakiHina
cooked: <p>Hello everyone:</p>
<p>I’m currently working on a binary semantic segmentation task. The model takes as input grid-based rainfall data from the past three hours and predicts whether the rainfall at each grid point will exceed a certain threshold from 2 to 9 hours into the future (similar to afternoon thunderstorm forecasting).</p>
<p>This is a multi-channel output task where each pixel is a binary classification problem. The dataset is highly imbalanced, with a very low proportion of positive samples. I tested two datasets with different positive sample ratios:</p>
<p>Group A: 0.3% positive samples</p>
<p>Group B: 0.9% positive samples</p>
<p>The model I used is a U-Net. The loss function is Dice Loss, batch size is 4, and the optimizer is Adam. Below are my experimental results:</p>
<p>Group A:</p>
<p>Final CSI Score: 0.0018915347754955292</p>
<p>Final Recall Score: 0.4033149182796478</p>
<p>Final F1score Score: 0.0037759272381663322</p>
<p>Final Precision Score: 0.001896842964924872</p>
<p>Group B:</p>
<p>Final CSI Score: 0.002810424193739891</p>
<p>Final Recall Score: 0.016717325896024704</p>
<p>Final F1score Score: 0.005605095531791449</p>
<p>Final Precision Score: 0.003367003286257386</p>
<p>Below are my code:</p>
<pre data-code-wrap="python"><code class="lang-python">from typing import Optional

import torch
from torch import nn

from kornia.core import Tensor
from kornia.core.check import KORNIA_CHECK, KORNIA_CHECK_IS_TENSOR
from kornia.losses._utils import mask_ignore_pixels
from kornia.utils.one_hot import one_hot

# based on:
# https://github.com/kevinzakka/pytorch-goodies/blob/master/losses.py
# https://github.com/Lightning-AI/metrics/blob/v0.11.3/src/torchmetrics/functional/classification/dice.py#L66-L207


def dice_loss(
    pred: Tensor,
    target: Tensor,
    average: str = "macro",
    eps: float = 1e-8,
) -&gt; Tensor:
    
    KORNIA_CHECK_IS_TENSOR(pred)

    if not len(pred.shape) == 4:
        raise ValueError(f"Invalid pred shape, we expect BxNxHxW. Got: {pred.shape}")

    if not pred.shape[-2:] == target.shape[-2:]:
        raise ValueError(f"pred and target shapes must be the same. Got: {pred.shape} and {target.shape}")

    if not pred.device == target.device:
        raise ValueError(f"pred and target must be in the same device. Got: {pred.device} and {target.device}")
    num_of_classes = pred.shape[1]
    possible_average = {"micro", "macro"}
    KORNIA_CHECK(average in possible_average, f"The `average` has to be one of {possible_average}. Got: {average}")

    pred_soft: Tensor = pred.sigmoid()

    # set dimensions for the appropriate averaging
    dims: tuple[int, ...] = (2, 3)

    if average == "micro":
        dims = (1, *dims)

        pred_soft = pred_soft
        target = target

    intersection = torch.sum(pred_soft * target, dims)
    cardinality = torch.sum(pred_soft + target, dims)

    dice_score = 2.0 * intersection / (cardinality + eps)
    dice_loss = -dice_score + 1.0

    # reduce the loss across samples (and classes in case of `macro` averaging)
    if average == "macro":
        dice_loss = dice_loss.sum(-1)

    dice_loss = torch.mean(dice_loss)

    return dice_loss


class DiceLoss(nn.Module):
    r"""Criterion that computes Sørensen-Dice Coefficient loss.

    According to [1], we compute the Sørensen-Dice Coefficient as follows:

    .. math::

        \text{Dice}(x, class) = \frac{2 |X| \cap |Y|}{|X| + |Y|}

    Where:
       - :math:`X` expects to be the scores of each class.
       - :math:`Y` expects to be the one-hot tensor with the class labels.

    the loss, is finally computed as:

    .. math::

        \text{loss}(x, class) = 1 - \text{Dice}(x, class)

    Reference:
        [1] https://en.wikipedia.org/wiki/S%C3%B8rensen%E2%80%93Dice_coefficient

    Args:
        average:
            Reduction applied in multi-class scenario:
            - ``'micro'`` [default]: Calculate the loss across all classes.
            - ``'macro'``: Calculate the loss for each class separately and average the metrics across classes.
        eps: Scalar to enforce numerical stabiliy.
        weight: weights for classes with shape :math:`(num\_of\_classes,)`.
        ignore_index: labels with this value are ignored in the loss computation.

    Shape:
        - Pred: :math:`(N, C, H, W)` where C = number of classes.
        - Target: :math:`(N, H, W)` where each value is
          :math:`0 ≤ targets[i] ≤ C-1`.

    Example:
        &gt;&gt;&gt; N = 5  # num_classes
        &gt;&gt;&gt; criterion = DiceLoss()
        &gt;&gt;&gt; pred = torch.randn(1, N, 3, 5, requires_grad=True)
        &gt;&gt;&gt; target = torch.empty(1, 3, 5, dtype=torch.long).random_(N)
        &gt;&gt;&gt; output = criterion(pred, target)
        &gt;&gt;&gt; output.backward()

    """

    def __init__(
        self,
        average: str = "macro",
        eps: float = 1e-8,
    ) -&gt; None:
        super().__init__()
        self.average = average
        self.eps = eps

    def forward(self, pred: Tensor, target: Tensor) -&gt; Tensor:
        return dice_loss(pred, target, self.average, self.eps)

</code></pre>
<p>I observed a very strange phenomenon:</p>
<p>In the dataset with a lower positive sample ratio (0.3%), the model tends to predict more positive samples, resulting in a noticeably higher recall but extremely low precision.</p>
<p>In contrast, with a slightly higher positive sample ratio (0.9%), the model becomes more conservative — recall drops to an almost unusable level, while precision slightly improves.</p>
<p>The model architecture, parameters, and training logic are exactly the same for both datasets. The only difference is the positive sample ratio, yet this shift in prediction tendency is completely reversed.</p>
<p>I would like to ask if anyone has encountered a similar phenomenon when using Dice Loss. Can this behavior be explained by the mathematical properties of Dice Loss?</p>