Understanding `from_logits=True` vs `from_logits=False` in `tf.losses.CategoricalCrossentropy` for U-Net

When training convolutional neural networks like U-Net for tasks such as image segmentation, the choice of loss function is pivotal. TensorFlow's `tf.losses.CategoricalCrossentropy` is commonly used for multi-class segmentation tasks, where the goal is to classify each pixel into one of several classes.

A key parameter in `CategoricalCrossentropy` is `from_logits`, which can either be set to `True` or `False`. This parameter significantly affects the behavior and results of the training process. Below, I will provide a detailed explanation of how these configurations differ and their implications on U-Net model training.

The Role of `from_logits` Parameter

The `from_logits` parameter in `tf.losses.CategoricalCrossentropy` tells the function whether the predictions (outputs) of the model are logits or normalized probabilities.

• Logits: Direct model outputs without any transformation. These can take any real value.
• Probabilities: Outputs that result from applying a `softmax` function to logits, which are constrained to lie between 0 and 1.

When we set:

• `from_logits=True`: The loss function assumes the model's outputs are raw logits. It internally applies a `softmax` to convert logits into probabilities before computing the cross-entropy loss.
• `from_logits=False`: The model is expected to output probabilities directly, so the `softmax` is not applied again within the loss function.

Implications on Training

1. Stability in Computation:
   When dealing with logits, numerical stability is a crucial concern. Using `from_logits=True` allows TensorFlow to merge `softmax` and cross-entropy operations into a single function (`softmax_cross_entropy_with_logits`), which provides a more numerically stable computation of the loss.
2. Model Architecture Consistency:
   Especially in U-Net or any segmentation network, it is often beneficial to let the network's last layer be a linear layer without `softmax`. This can be more flexible and consistent for the optimization process where `from_logits=True`.
3. Common Pitfalls:
   If `from_logits=False` is set wrongly and the model outputs logits, it can lead to highly unstable or meaningless loss calculations because probabilities exceeding 1 or below 0 are mathematically impossible under correct assumptions.

Practical Examples

Example with `from_logits=True`

Assume you have designed a U-Net model without a `softmax` layer on the last layer’s output. Here, setting: