PaperlayerView exact revision on arXiv

Source-linked AI summary

Constrained-CNN losses for weakly supervised segmentation

Hoel Kervadec, Jose Dolz, Meng Tang, Eric Granger, Yuri Boykov, Ismail Ben Ayed

arXiv:1805.04628v2cs.CV

TL;DR

Weakly supervised segmentation seeks to reduce the need for laborious pixel/voxel annotations, but directly optimizing global constraints in deep CNNs is computationally difficult. The paper introduces a differentiable penalty that enforces inequality constraints in the loss without proposal generation or dual iterates. Across three tasks, it achieves performance close to full supervision with few annotations and outperforms Lagrangian-based constrained CNNs, while its assessment does not cover severe abnormalities with sizes far outside training data.

  • Problem

    Directly enforcing flexible global inequality constraints in weakly supervised CNN segmentation is difficult because Lagrangian dual optimization is computationally intractable for large networks.

  • Method

    The paper adds a differentiable penalty directly to the segmentation loss to enforce inequality constraints, avoiding Lagrangian dual iterates and proposal generation.

  • Results

    Across three segmentation tasks, the method achieves results close to full supervision with only a small fraction of pixel annotations and substantially outperforms Lagrangian-based constrained CNNs.

  • Takeaways & Limitations

    Rough target-size estimates can yield highly competitive performance, while more precise bounds can approach full-supervision performance even with substantial size and shape variability.

  • Takeaways & Limitations

    Performance on images with severe abnormalities whose sizes differ substantially from the training set was not assessed.

Abstract

from arXiv · show

Weakly-supervised learning based on, e.g., partially labelled images or image-tags, is currently attracting significant attention in CNN segmentation as it can mitigate the need for full and laborious pixel/voxel annotations. Enforcing high-order (global) inequality constraints on the network output (for instance, to constrain the size of the target region) can leverage unlabeled data, guiding the training process with domain-specific knowledge. Inequality constraints are very flexible because they do not assume exact prior knowledge. However, constrained Lagrangian dual optimization has been largely avoided in deep networks, mainly for computational tractability reasons. To the best of our knowledge, the method of [Pathak et al., 2015] is the only prior work that addresses deep CNNs with linear constraints in weakly supervised segmentation. It uses the constraints to synthesize fully-labeled training masks (proposals) from weak labels, mimicking full supervision and facilitating dual optimization. We propose to introduce a differentiable penalty, which enforces inequality constraints directly in the loss function, avoiding expensive Lagrangian dual iterates and proposal generation. From constrained-optimization perspective, our simple penalty-based approach is not optimal as there is no guarantee that the constraints are satisfied. However, surprisingly, it yields substantially better results than the Lagrangian-based constrained CNNs in [Pathak et al., 2015], while reducing the computational demand for training. By annotating only a small fraction of the pixels, the proposed approach can reach a level of segmentation performance that is comparable to full supervision on three separate tasks. While our experiments focused on basic linear constraints such as the target-region size and image tags, our framework can be easily extended to other non-linear constraints.

1 Introduction

Weak supervision reduces reliance on laborious pixel/voxel annotations, while global inequality constraints can guide CNN segmentation using domain knowledge. The paper proposes a differentiable loss that imposes these constraints directly, avoiding proposal generation and expensive dual optimization.

  • Motivation: Weak supervision uses partial or uncertain labels to reduce the burden of pixel-level annotation and leverage unlabeled data.Examples include bounding boxes, points, scribbles, and image tags.
  • Constraint formulation: Global constraints can encode flexible prior knowledge such as lower and upper bounds on the target-region size.Image-tag presence and suppression constraints are special cases of inequality constraints on network outputs.
  • Optimization gap: Direct Lagrangian-dual optimization is computationally intractable for CNNs with millions of parameters because each dual iteration requires CNN optimization.Although the constraints may be linear and convex in the outputs, CNN non-convexity makes the overall problem difficult.
  • Prior work: Prior work synthesized fully labeled proposals from weak labels and trained the CNN to match a latent proposal distribution.This avoided direct dual optimization but mimicked full supervision through proposal generation.
  • Contribution: The proposed differentiable penalty enforces inequality constraints in the loss, outperforming the Lagrangian-based approach while reducing training demand.In cardiac segmentation, performance approached full supervision using 0.1% of full ground-truth labels.

2 Related work

Related work uses weak annotations and proposal generation to train segmentation networks, but proposal errors can reinforce themselves. Global constraints remain difficult to impose directly on deep CNN outputs, especially in medical imaging.

  • Weak supervision: Weakly supervised segmentation commonly uses image tags, bounding boxes, scribbles, or points instead of complete masks.Medical-image weak supervision remains comparatively scarce in the literature.
  • Limitations: Proposal errors can reinforce themselves in self-taught learning schemes and undermine convergence guarantees.Direct regularizers avoid extra proposal-generation steps but do not explicitly impose global constraints.
  • Proposal strategies: Many proposal-based methods alternate CNN learning with proposal generation, often using dense CRF refinement or related inference steps.The generated proposals are reused as pseudo-labels for subsequent CNN training.
  • Constrained CNNs: Direct Lagrangian optimization is generally impractical for modern deep CNNs because each dual step requires retraining a network with millions of parameters.Pathak et al. instead constrained a latent distribution representing a fake ground truth and minimized KL divergence to CNN outputs.

3 Proposed loss function

The proposed loss adds a differentiable penalty to partial cross-entropy so soft size constraints can be optimized with standard stochastic gradient descent. The penalty is inactive within the permitted interval and responds to violations outside it.

  • Loss definition: The proposed weakly supervised loss combines partial cross-entropy with a weighted differentiable constraint term.The positive constant λ controls the importance of the constraints.
  • Penalty shape: The penalty is quadratic below the lower bound, quadratic above the upper bound, and zero when the target-region size lies within bounds.This creates a soft penalty for violating the permitted size interval.
  • Optimization: The differentiable term is compatible with standard stochastic gradient descent during back-propagation.Its gradient contributes directly to the network-parameter update.
  • Gradient behavior: When the size constraint is satisfied, the penalty gradient is zero and does not affect the current network-parameter update.When the target region falls below its lower bound, the gradient increases the softmax outputs.
  • Evaluation context: Figure 1 illustrates the differentiable loss used to impose soft size constraints on the target region.The method is evaluated against fully supervised and Lagrangian-based constrained CNN settings.

4.1 Medical Image Data:

The evaluation uses three publicly available medical-imaging datasets covering cardiac, vertebral-body, and prostate segmentation. These datasets vary in anatomy, acquisition conditions, resolution, and available ground-truth annotations.

  • Datasets: The experiments cover cardiac, vertebral-body, and prostate segmentation across three publicly available datasets.Each dataset represents a different medical-imaging application.
  • Cardiac data: The cardiac dataset contains 100 cine MR exams from the 2017 ACDC Challenge, including normal subjects and several defined cardiac pathologies.The listed pathologies include dilated and hypertrophic cardiomyopathy, myocardial infarction, and abnormal right ventricle.
  • Vertebral-body data: The vertebral-body dataset contains 23 3D T2-weighted MR images from 23 patients, sampled to 39×305×305 voxels.Each image includes seven manually identified and segmented vertebral bodies.
  • Prostate data: The prostate dataset contains T2-weighted MR images from 50 patients acquired across centers, MRI vendors, and scanning protocols.Forty patients were used for training and ten for validation.

4.2 Weak annotations:

The study generates weak annotations from fully labeled images using erosion or random-point strategies, testing robustness to label-generation strategy and location.

  • Two weak-label strategies are evaluated: binary erosion for the left-ventricle dataset and random points for vertebral-body and prostate datasets.Figure 2 shows corresponding fully annotated images and weak labels.
  • Erosion: 0.1% of labeled pixels were annotated for the left-ventricle erosion strategy.A 10×10 erosion kernel was reduced to 7×7 or smaller when the label disappeared.
  • Random point: 0.02% of dataset pixels received ground-truth labels under the random-point strategy.Each annotation was a circle with maximum radius 4 pixels centered on a randomly selected foreground point.

4.3 Different levels of supervision:

The experiments compare supervision levels ranging from partial labels to full pixel annotations, while adding image-tag or target-size constraints to weak supervision.

  • Baselines: The fully supervised baseline uses foreground and background labels for every pixel, whereas partial cross-entropy uses only the labeled-pixel fraction.The partial cross-entropy model is the lower baseline, and the fully supervised model is the upper baseline.
  • Image tags bounds: Image-tag constraints encode target absence with predicted size 0 and target presence with a positive predicted size.These coarse bounds convey information about whether the target region is present.
  • Common bounds: Common size bounds are computed from minimum and maximum target sizes across slices of one patient, then expanded by factors of 0.9 and 1.1.For ACDC, the example bounds are 60 and 2000 when the target is present, and 0 otherwise.
  • Individual bounds: Individual bounds use the true target-region size for each slice to provide more precise supervision than common bounds.The experiment tests whether narrower slice-specific size ranges improve performance.
  • Mixed supervision: Mixed supervision combines full cross-entropy on n fully annotated images with constrained partial cross-entropy on the remaining m−n weakly labeled images.The weakly labeled images use common size bounds.

4.4 Constraining a 3D volume:

The formulation extends size constraints from individual 2D images to the 3D volume formed by a batch containing all slices of that volume.

  • The 3D formulation constrains the target-region volume across a batch containing all 2D slices of the volume.The batch is represented using labels and associated predictions for each slice.
  • 3D volume constraints provide less supervision than the 2D scenarios, where slices have independent supervision such as image tags.

4.5 Training and implementation details:

The experiments use ENet or a residual U-Net, train with Adam from scratch, evaluate using DSC, and compare the proposed loss with a reimplemented Lagrangian-proposal method.

  • ENet is used for left-ventricle and vertebral-body segmentation, while a fully residual U-Net is used for the more difficult prostate task.The architectures were selected for accuracy–inference-time trade-offs or task difficulty.
  • Networks are trained from scratch with Adam, starting at a learning rate of 5×10^-4 and halving it after 20 validation epochs without improvement.Volumes are sliced into 256×256 images and zero-padded when needed.
  • Segmentation performance is evaluated with the Dice similarity coefficient (DSC).The implementation uses PyTorch and a GTX 1080 Ti GPU with 11 GB of video memory.
  • Lagrangian-proposal comparison: The reimplemented Lagrangian-proposal method alternates projected-gradient-ascent proposal synthesis with cross-entropy optimization for fixed synthetic labels.The PGA learning rate is 5×10^-5, and early stopping is added when the dual optimization converges.

5 Results

Across left-ventricle, vertebral-body, and prostate experiments, direct constraint losses improved weakly supervised segmentation, especially with image-specific bounds, while reducing training overhead versus Lagrangian proposals.

  • Experimental design: The experiments compare direct constraint losses with Lagrangian proposals, weak supervision, hybrid supervision, and full supervision across three segmentation tasks.The evaluation includes size constraints, multiple supervision levels, and qualitative comparisons.
  • Left-ventricle segmentation: 0.8708 DSC was achieved on left-ventricle segmentation with individual bounds, only 2% below full supervision.The Lagrangian-proposal method remained similar to common loose bounds and did not exploit the more precise individual information.
  • Left-ventricle segmentation: 0.8580 mean DSC was obtained for 3D target-volume segmentation using less supervision than in the 2D setting.The result suggests that annotators may label only a fraction of slices rather than every 2D slice.
  • Hybrid training: Adding 5, 10, and 25 fully annotated images increased DSC by 4%, 5%, and 6%, respectively, with 25% of fully labeled images approaching full-supervision performance.Hybrid training outperformed training using weakly annotated images alone.
  • MR-T2 segmentation: 0.8604 DSC was reached on vertebral-body segmentation with individual bounds, only 3% below full supervision.Individual bounds also handled difficult cases better than common bounds, which achieved 0.7900 DSC.
  • MR-T2 segmentation: Individual bounds improved prostate segmentation to 0.8298 DSC, while common bounds improved results over tags by approximately 3%.The gap between weak supervision with individual bounds and full supervision was larger for prostate segmentation than for the other datasets.

6 Discussion

The proposed differentiable penalty directly enforces inequality constraints in the loss, avoiding dual iterates and proposal generation. Across tasks, it benefits from weakly annotated data, can approach full supervision with precise bounds, and remains untested on severe abnormalities unlike those in training.

  • The differentiable penalty enforces inequality constraints directly in the loss, avoiding expensive Lagrangian dual iterates and proposal generation.
  • Weakly annotated data with the direct size loss is especially beneficial when limited fully annotated data is available.
  • Highly competitive performance is obtained with rough target-size estimates, while precise bounds can approach full supervision despite large size and shape variability.
  • Segmentation of severe abnormalities whose sizes differ substantially from training examples was not assessed.
  • The framework can extend to nonlinear constraints involving invariant shape moments, centroid localization, and image-intensity statistics.
  • The penalty approach is faster and more stable than Lagrangian-dual optimization, but it provides no guarantee that constraints are satisfied.

7 Conclusion

The paper introduces a simple direct constrained-CNN loss for weakly supervised segmentation. It reports performance close to full supervision across three tasks with only a small fraction of pixels annotated and identifies extensions to nonlinear constraints.

  • The novel loss function performs significantly better than Lagrangian optimization for weakly supervised image segmentation.
  • Annotating only a small fraction of pixels yields results close to full supervision across three different tasks, with negligible computation overhead.
  • The direct constrained-CNN loss can extend beyond target-region size and image-tag constraints to invariant shape moments and other region statistics.
  • The approach has potential to close the gap between weakly and fully supervised learning in semantic medical image segmentation.
Loading 1805.04628v2…