Introduction

Dice loss is popular in segmentation because it focuses on overlap between prediction and target, which makes it useful for imbalanced masks. The tricky part is not the formula itself, but getting the TensorFlow or Keras implementation to match tensor shapes, class layout, and probability semantics. A lot of broken implementations fail because they threshold predictions, reduce across the wrong axes, or mix binary and multiclass cases incorrectly.

Binary Dice Loss: Use Probabilities, Not Hard Labels

For binary segmentation, the usual setup is:

• 'y_true contains binary masks,'
• 'y_pred contains probabilities from a sigmoid output,'
• the loss is computed on soft values.

1import tensorflow as tf
2
3
4def dice_loss(y_true, y_pred, smooth=1e-6):
5    y_true = tf.cast(y_true, tf.float32)
6    y_pred = tf.cast(y_pred, tf.float32)
7
8    axes = (1, 2, 3)
9    intersection = tf.reduce_sum(y_true * y_pred, axis=axes)
10    denominator = tf.reduce_sum(y_true + y_pred, axis=axes)
11
12    dice = (2.0 * intersection + smooth) / (denominator + smooth)
13    return 1.0 - tf.reduce_mean(dice)

The important detail is that y_pred should stay soft. Do not round or threshold during training, because that kills gradient quality.

A Minimal Keras Model Example

A typical binary segmentation setup looks like this:

1inputs = tf.keras.Input(shape=(128, 128, 1))
2x = tf.keras.layers.Conv2D(8, 3, padding="same", activation="relu")(inputs)
3outputs = tf.keras.layers.Conv2D(1, 1, activation="sigmoid")(x)
4
5model = tf.keras.Model(inputs, outputs)
6model.compile(optimizer="adam", loss=dice_loss)

If the final layer uses sigmoid, the loss should receive probabilities in the range [0, 1].

Why Thresholding Is Wrong During Training

This is a common mistake:

# wrong for training
binary_pred = tf.cast(y_pred > 0.5, tf.float32)

That turns the model output into a hard decision before the loss is computed. The gradient through that thresholding step is not useful for optimization. Dice loss is supposed to work with soft predictions during training and hard masks only during evaluation or visualization.

Multiclass Dice Loss Needs Axis Care

For multiclass segmentation, the model usually outputs shape like [batch, height, width, classes]. The ground truth is often one-hot encoded to the same shape.

1import tensorflow as tf
2
3
4def multiclass_dice_loss(y_true, y_pred, smooth=1e-6):
5    y_true = tf.cast(y_true, tf.float32)
6    y_pred = tf.cast(y_pred, tf.float32)
7
8    axes = (1, 2)
9    intersection = tf.reduce_sum(y_true * y_pred, axis=axes)
10    denominator = tf.reduce_sum(y_true + y_pred, axis=axes)
11
12    dice = (2.0 * intersection + smooth) / (denominator + smooth)
13    return 1.0 - tf.reduce_mean(dice)

Here the reduction preserves the class dimension so each class gets its own overlap score before averaging.

Logits Versus Probabilities

Another common error is feeding raw logits into a Dice implementation that assumes probabilities. If your final layer does not apply sigmoid or softmax, convert appropriately inside the loss or change the model so the loss receives normalized outputs.

For binary segmentation:

• sigmoid output pairs naturally with binary Dice,
• raw logits require explicit tf.nn.sigmoid before overlap is computed.

For multiclass segmentation:

• softmax output is the usual choice,
• raw logits require tf.nn.softmax before Dice overlap.

Dice Plus Cross-Entropy Is Common

Dice loss is often paired with binary cross-entropy or categorical cross-entropy because Dice emphasizes overlap but may not provide the best optimization behavior alone in every case.

def bce_dice_loss(y_true, y_pred):
    bce = tf.keras.losses.binary_crossentropy(y_true, y_pred)
    return tf.reduce_mean(bce) + dice_loss(y_true, y_pred)

This is a common practical choice for medical imaging and foreground-background segmentation tasks.

Shape And Type Consistency Matter

A correct formula can still fail if the tensors do not align.

Check these points:

• 'y_true and y_pred should have matching shapes,'
• 'y_true should be numeric and castable to float,'
• binary masks should use one channel if the model predicts one channel,
• multiclass targets should match the softmax class dimension.

Many “Dice loss is broken” issues are actually data pipeline mismatches.

Common Pitfalls

• Thresholding predictions before computing the loss.
• Using raw logits in a Dice formula that expects probabilities.
• Reducing over the wrong axes and accidentally mixing batch or class dimensions.
• Using integer masks and relying on implicit dtype behavior instead of explicit casting.
• Applying a binary Dice implementation to multiclass outputs without one-hot alignment.

Summary

• Dice loss should usually operate on soft probabilities, not thresholded predictions.
• Binary and multiclass segmentation need slightly different reduction logic.
• Match the loss to the model output activation, such as sigmoid or softmax.
• Many practical setups combine Dice loss with cross-entropy for more stable training.
• Most implementation bugs come from shape, axis, or probability-semantic mistakes rather than from the formula itself.