Shape-prior regularization

Overview

A pixel-wise loss such as cross-entropy or a smooth ℓ1 loss does not see the shape of the whole output. A prediction can score well per pixel and still be an anatomically impossible shape 1705.08302 §1. A shape-prior regularizer adds a second term to the training loss. That term maps the prediction and the ground truth into a learned low-dimensional code and penalises their distance 1705.08302 §2.2. The code is learned from the data, so no explicit shape model has to be specified by hand.

Every variant shares three parts: a pixel-wise reconstruction loss, a code distance between prediction and target, and a small weight on that distance. The variants differ in what is coded. For label maps, an autoencoder learned on the masks provides the code 1705.08302 §2.2. For intensity images, a predictor trained to match the autoencoder's codes provides it 1705.08302 §2.4. The code is used only as a training signal, so the deployed network is the same as the baseline 1705.08302 §2.4.

Variants

Label-map shape term (ACNN-Seg)

A convolutional autoencoder is trained on ground-truth label maps. Its encoder maps the predicted class probabilities and the true labels to 64-value codes. The segmentation network is trained with cross-entropy plus the squared code distance, plus weight decay on the convolution filters 1705.08302 §2.2 1705.08302 eq. 1. The reported gain over the sub-pixel 3D baseline on cardiac MR is endocardial Dice .939 against .923 and myocardial Dice .811 against .773 1705.08302 tab. 1. Card: Anatomically Constrained Neural Networks (ACNN).

Intensity-image shape term (ACNN-SR)

Super-resolution outputs are intensity images, which the label-map autoencoder cannot encode. A predictor maps an intensity volume to the same 64-value code space. The predictor is trained in two stages: first the autoencoder on masks, then the predictor to match the autoencoder's codes. Then both are trained jointly with the super-resolution network 1705.08302 §2.4. The super-resolution network is trained with a smooth ℓ1 reconstruction loss plus the squared code distance 1705.08302 §2.3 1705.08302 eq. 2. The reported SSIM is .796 against .784 for 3D-UNet and .783 for SR-CNN on the same test set, and the reported inference time is 0.06 s against 0.29 s for SR-CNN 1705.08302 tab. 3. The same learned codes were also used for a classification test on cardiac MR: with both the ED and ES phases, T-L codes reached 91.6 % accuracy against 86.6 % for 64-value PCA codes of the label maps 1705.08302 §3.6.

When to use which

Pick the variant from the output type first. The paper's evidence is all from cardiac imaging, on three datasets, and the reported gains are small (a few Dice points in the segmentation table). In the ultrasound table the Dice gain has p ≈ 0.05, which the text calls not significant 1705.08302 tab. 2 1705.08302 §3.4.

Implementation sketch

Common core: the pixel-wise loss, plus λ1 times the squared distance between the codes of prediction and target, plus λ2 / 2 times the squared weights. The two variants differ in how the code is computed, as the two pseudocode blocks below show. Eq. 1 updates only the segmentation weights, which implies a frozen encoder. The paper does not state this in words 1705.08302 §2.2 1705.08302 eq. 1.

# Label-map variant (ACNN-Seg)  [arxiv:1705.08302 §2.2] [eq. 1]
def seg_loss(images, labels, seg_net, ae_encoder, lam1, lam2):
    # Per-class probability maps on the output grid  [§2.1]
    probs = seg_net(images)
    # Pixel-wise loss; mean or sum over voxels is not stated  [§2.1]
    ce = cross_entropy(probs, labels)
    # Codes of prediction and target from the frozen encoder  [§2.2]
    shape = squared_distance(ae_encoder(probs), ae_encoder(labels))
    # Weight decay on convolution filter weights  [§2.2]
    return ce + lam1 * shape + (lam2 / 2) * conv_weight_norm(seg_net)

# Intensity-image variant (ACNN-SR)  [arxiv:1705.08302 §2.3] [§2.4] [eq. 2]
def sr_loss(lr_images, hr_targets, sr_net, predictor, lam1, lam2):
    # Super-resolved volume from LR-grid features  [§2.3]
    sr = sr_net(lr_images)
    # Smooth L1 (Huber) reconstruction, threshold 1  [§2.3] [eq. 2]
    recon = smooth_l1(sr - hr_targets, threshold=1.0)
    # Codes of synthesised and true volumes from the predictor  [§2.4]
    # Whether predictor weights also update is not stated
    shape = squared_distance(predictor(sr), predictor(hr_targets))
    # Weight decay on convolution filter weights  [§2.4]
    return recon + lam1 * shape + (lam2 / 2) * conv_weight_norm(sr_net)

The reported values for λ1 and λ2 are 0.01 and 5 × 10⁻⁶, with a fixed learning rate of 0.001. The optimiser variant and schedule are not stated 1705.08302 supp. Implementation Details. The code dimension is 64 1705.08302 tab. 5. The label-map autoencoder's code layer has no non-linearity 1705.08302 supp. Network Structure.

Contradictions

None known. The paper reports that a joint update of the two stages converged better than alternating updates, and does not compare them further 1705.08302 §2.4 1705.08302 §3.2.

Papers

• Anatomically Constrained Neural Networks (ACNN): introduces the code-distance shape term for cardiac segmentation and MR super-resolution, with a training-only autoencoder and predictor.