Lovász-Softmax Loss: Optimizing IoU for Semantic Segmentation
· 3 min read · 436 words
Lovász-Softmax is a surrogate loss designed to optimize intersection over union (IoU), also called the Jaccard index, for multiclass semantic segmentation.
For a class with predicted set and ground-truth set ,
IoU is the evaluation metric we care about, but thresholding predictions into sets makes it non-differentiable. Pixel-wise cross-entropy is differentiable, but it optimizes independent classification decisions rather than the image-level overlap ratio.
Core idea
For each class, Lovász-Softmax computes pixel errors from the softmax probabilities, sorts those errors from largest to smallest, and takes a weighted sum. The weights are discrete gradients of the Jaccard loss as progressively more sorted pixels are included. This is the Lovász extension of the set loss.
Sorting is not cosmetic: it couples the pixels and gives more weight to errors that cause a larger change in IoU. The resulting piecewise-linear surrogate is differentiable almost everywhere and can be optimized with gradient methods.
Practical guidance
- Pass probabilities, not raw logits, to implementations that expect softmax scores.
- Exclude ignored labels before sorting.
- Decide how to handle classes absent from an image. Common implementations either average over all classes or only classes present in the target.
- Per-image and per-dataset optimization are different objectives. The original paper evaluates both.
- For stable early training, it is common to train with cross-entropy first or combine cross-entropy and Lovász-Softmax. That is a practical recipe, not a requirement of the loss.
The loss has no inference-time cost. Its main training overhead is sorting class-specific pixel errors.
When it helps
Lovász-Softmax is most attractive when leaderboard or product quality is judged by mean IoU and cross-entropy produces good pixel accuracy but mediocre region overlap. It does not replace careful class sampling, label cleaning, or boundary evaluation; it aligns the optimization objective more closely with one metric.
Primary source
- Maxim Berman, Amal Rannen Triki, and Matthew B. Blaschko, The Lovász-Softmax Loss: A Tractable Surrogate for the Optimization of the Intersection-Over-Union Measure in Neural Networks, CVPR 2018.
