Jaccard Reward in Segmentation

• Jaccard Reward is a differentiable surrogate for the discrete Jaccard index, aligning network training directly with segmentation evaluation metrics like IoU.
• It replaces non-differentiable indicator functions with probabilistic predictions, enabling effective SGD optimization and offering bounded error guarantees.
• Validated across medical segmentation tasks, its implementation in frameworks like TensorFlow and PyTorch demonstrates robust performance compared to cross-entropy losses.

The Jaccard reward, often referred to as the soft-Jaccard loss in the context of neural network training for image segmentation, is a differentiable surrogate for the discrete Jaccard index (Intersection-over-Union, IoU). Used extensively in medical image segmentation, this reward function directly aligns the training objective with the evaluation metric, thereby addressing discrepancies inherent in cross-entropy-based optimization. It is grounded in risk minimization theory and enables stochastic gradient descent (SGD)-based training while maintaining provable control over the true Jaccard evaluation target (Bertels et al., 2019).

1. Mathematical Definition and Soft Surrogate

For segmentation, the Jaccard index compares binary masks over $d$ pixels/voxels, with $y\in\{0,1\}^d$ denoting ground truth labels and $p\in[0,1]^d$ denoting probabilistic predictions output by a network. The discrete Jaccard index is defined as: $$J(y,\tilde{y}) = \frac{|y \cap \tilde{y}|}{|y \cup \tilde{y}|} = \frac{\sum_{i=1}^d 1_{[y_i=1]} 1_{[\tilde{y}_i=1]}}{\sum_{i=1}^d (1_{[y_i=1]} + 1_{[\tilde{y}_i=1]} - 1_{[y_i=1]}1_{[\tilde{y}_i=1]})}$$ To support gradient-based learning, the soft-Jaccard surrogate replaces indicator functions with probabilistic predictions: $$J_{\text{soft}}(p, y) = \frac{\sum_{i=1}^d p_i y_i}{\sum_{i=1}^d (p_i + y_i - p_i y_i)}$$ The Jaccard reward (or equivalently, the soft-Jaccard loss) is then $\ell_J(p, y) = 1 - J_{\text{soft}}(p, y)$.

2. Risk Minimization and Approximation Bounds

Minimizing the expected loss $\mathbb{E}_{(x, y)}[1 - J(y, f(x))]$ is generally infeasible due to the non-differentiability of the true Jaccard index. The soft-Jaccard loss acts as a smooth surrogate, and a minimal requirement for any surrogate is that it "controls" the target metric. The Dice-Jaccard relationship,

$$D(y, \tilde{y}) = \frac{2 J(y, \tilde{y})}{1 + J(y, \tilde{y})} \quad\Longleftrightarrow\quad J(y, \tilde{y}) = \frac{D(y, \tilde{y})}{2 - D(y, \tilde{y})}$$

shows that Dice and Jaccard are multiplicatively equivalent up to a factor of 2: $$J \leq D \leq 2J, \quad 1 - D \leq 1 - J \leq 2(1 - D)$$ Thus, minimizing the expectation of $1 - J_{\text{soft}}$ provably minimizes the expected true $y\in\{0,1\}^d$0 up to this multiplicative factor. An absolute error bound $y\in\{0,1\}^d$1 also holds, but the essential guarantee lies in the multiplicative control.

3. Impossibility of Optimal Cross-Entropy Weighting

Weighted cross-entropy (wCE) is a common workaround for class imbalance in segmentation. Formally, wCE is equivalent to optimizing a weighted Hamming similarity: $y\in\{0,1\}^d$2 However, for any fixed $y\in\{0,1\}^d$3, there exist segmentations (notably with small target objects) such that the relative error

$y\in\{0,1\}^d$4

is unbounded. No uniform finite $y\in\{0,1\}^d$5 exists such that $y\in\{0,1\}^d$6 for all $y\in\{0,1\}^d$7. This establishes that no constant cross-entropy weighting yields a uniform surrogate for the Jaccard or Dice metrics, unlike soft-Jaccard, which maintains bounded approximation (Bertels et al., 2019).

4. Differentiability, Gradient Formulation, and Implementation

The soft-Jaccard loss is differentiable and admits simple closed-form expressions for gradients. If

$y\in\{0,1\}^d$8

then

$y\in\{0,1\}^d$9

and

$p\in[0,1]^d$0

For $p\in[0,1]^d$1, the backward pass requires

$p\in[0,1]^d$2

This formulation supports direct implementation in autodiff frameworks such as TensorFlow or PyTorch. For explicit implementation, the recommended procedure is:

• Forward pass: compute $p\in[0,1]^d$3, $p\in[0,1]^d$4, $p\in[0,1]^d$5, and $p\in[0,1]^d$6
• Backward pass: compute gradients $p\in[0,1]^d$7 per pixel/voxel.

5. Empirical Performance Across Medical Segmentation Tasks

Extensive evaluations on five medical image segmentation datasets (BRATS 2018, ISLES 2017/18, MO17, PO18) compared five loss functions: cross-entropy (CE), weighted cross-entropy (wCE), soft-Dice, soft-Jaccard, and Lovász-sigmoid (a convex extension of IoU). Results demonstrate:

• Both CE and wCE yield significantly lower Dice and Jaccard test scores (p < 0.05) than the soft-Jaccard, soft-Dice, or Lovász surrogates.
• There is no statistically significant difference in test performance among {soft-Dice, soft-Jaccard, Lovász} across all evaluated tasks.
• Weighted CE may offer marginal improvement for certain object-size ranges but can underperform drastically for small lesions.
• Metric-sensitive surrogates consistently outperform CE/wCE over all object-size bins.

Practitioner recommendation: Use a metric-sensitive surrogate (soft-Jaccard, soft-Dice, Lovász-IoU) when optimizing for Dice or Jaccard evaluation, selected for convenience or minor stability advantages. No CE weighting scheme can uniformly match IoU-based performance (Bertels et al., 2019).

6. Summary and Implications for Metric-sensitive Optimization

The soft-Jaccard reward

$p\in[0,1]^d$8

is simple to compute, directly optimizes the associated evaluation metric, provides a mathematically grounded control over the discrete Jaccard, and yields easily computable gradients for SGD. In practice, it is as effective as other differentiable metric surrogates (soft-Dice, Lovász-IoU) and systematically outperforms (weighted) cross-entropy in segmentation scenarios where Dice or Jaccard is the primary evaluation metric. These results support the transition from standard cross-entropy-based optimization to metric-sensitive surrogates in modern convolutional neural network segmentation pipelines (Bertels et al., 2019).