Weakly-Supervised Convolutional Neural Networks for Multimodal Image Registration

In plain words

Doctors often guide a procedure with one scan type and plan it with another. Here the planning scan is a detailed MRI of the prostate, and the guiding scan is an ultrasound taken during the procedure, which is much harder to read. To aim a needle, the two images have to be lined up, and that is hard because the two scan types show different structures. Standard methods that try to match brightness patterns between the two images fail on this pair.

The authors' idea is to teach a neural network to line the images up using only a few labelled points or regions that a person has marked as corresponding in both images. Those labels are easier to collect than a complete pixel-by-pixel match. The network never sees a full ground-truth alignment. During training it is scored on whether the marked regions land on top of each other after it has moved the MRI. When it is used, it needs only the two images and no labels, and it returns the warp in a fraction of a second.

The paper also argues that the labels need not exist for every training pair, and that the network learns to fill in the rest of the alignment from the labelled cases.

TL;DR

A weakly-supervised 3D CNN that predicts a dense displacement field from a moving MR and a fixed TRUS volume. Training uses only corresponding anatomical label pairs, scored with a multiscale soft Dice that is differentiable, plus a bending-energy penalty on displacements §2.1 §2.2. Inference is fully automatic, needs no labels or initialisation, and runs about four registrations per second on the authors' GPU §4.1. Cross-validated accuracy on 108 prostate image pairs is a median TRE of 3.6 mm and a median gland Dice of 0.87 §4.1 tab. 1. Use it when you need learned deformable MR-to-ultrasound registration from sparse labels; expect a single-trial, single-centre dataset and no independent test set §5.

Problem

Intensity-based similarity measures are unreliable for multimodal registration here. Different physics means the same anatomy can look different, and different structures can look alike across modalities §1. Feature-based and motion-model methods need anatomy that is not consistently visible in both images, or manual segmentation or initialisation §1 §4.4. Supervised learning of voxel correspondence needs dense ground truth, which is not available for patient data; anatomical labels are more practical to obtain but are sparse and may not exist on every pair §1 §3.1. In this paper's own tests, pairwise intensity-based non-rigid registration on this dataset gave a median TRE of 19.6 mm after centroid alignment alone, and 20.6 to 25.6 mm with FFD registration on top of centroid initialisation §4.4 tab. 2.

Method

The network predicts a dense displacement field (DDF) u from an image pair. The moving image's labels are warped by u and compared with the fixed image's labels. Training maximises the expected label similarity over the training pairs, minus a deformation penalty, and never uses voxel-level correspondence §2.1 fig. 2.

Objective §2.1 eq. 1 eq. 2

Over N training image pairs, each with M_n label pairs, the utility is the average label similarity. The moving label is warped by the predicted DDF, and it is compared with the fixed label §2.1:

J = \frac{1}{N} \sum_{n=1}^{N} \frac{1}{M_n} \sum_{m=1}^{M_n} J_{mn}\left(l^B_{mn}, y^A_{mn}\right) eq. 1

Here y^A_mn is the moving label l^A_mn warped by the displacements u_n, and the network has parameters θ §2.1 eq. 1. Training minimises the negative utility plus a weighted deformation regulariser Ω §2.1:

\hat{\theta} = \arg\min_{\theta} \left[ -J(x^A, x^B, l^A, l^B; \theta) + \alpha \cdot \Omega(u) \right] eq. 2

The regulariser Ω penalises non-smooth displacements, and α sets its weight §2.1 eq. 2. In the baseline Ω is the bending energy and α = 0.5 §3.2.1.

Minibatch gradient with two-stage sampling §2.1 eq. 3

Each minibatch should contain the same number K of image-label pairs. The paper samples K image pairs uniformly, then one label pair uniformly from each sampled image pair's labels §2.1. It shows that the expectation of this minibatch gradient equals the full gradient of J eq. 3.

# Two-stage sampling for one minibatch  [§2.1]
# Input:  K, image pairs, labels per image pair
# Output: K image pairs, each with one label pair
def sample_minibatch(image_pairs, labels_per_pair, K):
    # Stage 1: draw K image pairs uniformly  [§2.1]
    chosen_pairs = uniform_sample(image_pairs, K)
    label_pairs = []
    for pair in chosen_pairs:
        # Stage 2: draw one label pair uniformly from that pair  [§2.1]
        label_pairs.append(uniform_sample(labels_per_pair[pair], 1))
    return chosen_pairs, label_pairs

Label similarity: multiscale soft Dice §2.2 fig. 3

Binary overlap metrics ignore distance when objects do not overlap, and they saturate. The paper therefore blurs each label at several Gaussian scales and averages a soft Dice over those scales §2.2. For two label maps a and b with values in [0, 1]:

S_{Dice}(\mathbf{a}, \mathbf{b}) = \frac{2 \sum_{i=1}^{I} a_i \cdot b_i}{\sum_{i=1}^{I} a_i + \sum_{i=1}^{I} b_i} eq. 4

The index i runs over I voxels §2.2 eq. 4. The label similarity for one label pair is the average of this Dice over Z = 7 scales, where f_σ is a 3D isotropic Gaussian filter §2.2:

J_k = \frac{1}{Z} \sum_{\sigma} S_{Dice}\left(f_\sigma(l^B_k), f_\sigma(y^A_k)\right) eq. 5

The scales are σ in {0, 1, 2, 4, 8, 16, 32}. At σ = 0 the filter is a Dirac delta, so the unfiltered binary label is included §2.2 eq. 5. The filter is differentiable and can be applied after warping and augmentation §2.2.

The paper also gives a cross-entropy variant, adapted to the same multiscale scheme and with two classes c in {1, 2} §2.2 eq. 6:

S_{CE}(\mathbf{a}, \mathbf{b}) = \sum_{i=1}^{I} \sum_{c=1}^{2} p_c(a_i) \log p_c(b_i) eq. 6

The text calls this a negative cross-entropy and recommends clipping extreme probabilities §2.2 eq. 6. See Pitfalls for the sign.

Network §2.3 fig. 4

The network is a 3D CNN with four down-sampling blocks and four up-sampling blocks §2.3.

• Each block has two convolution, batch-normalisation and ReLU layers §2.3.
• Down-sampling is max-pooling with stride two; up-sampling is transposed convolution with stride two §2.3.
• All kernels are 3×3×3, except 7×7×7 in the first convolution, which gives a larger receptive field §2.3.
• The number of channels starts at n0, doubles at each down-sampling step and halves at each up-sampling step §2.3.
• Residual shortcuts: eight standard shortcuts that sum feature maps over two convolutions, plus four trilinear additive up-sampling layers §2.3.
• Displacement skip layers: four summation skips add a trilinear up-sampled displacement from resolution levels s1 to s4 to the displacement at s0. Each skip uses one convolution as a bias term, with no batch normalisation and no nonlinearity §2.3 fig. 4.
• The sum of these displacements at s0 gives the single output DDF §2.3 fig. 5.

The paper says trilinear sampling gives bounded nonlinearity between linear convolutions, and that summing these displacement skips is cheaper than composing displacements across levels, which needs resampling §2.3.

Output configurations §3.2.1 fig. 5

Three configurations are compared for the output summation. The baseline sums the outputs at all levels s0 to s4 fig. 5a. The δ0 variant uses only the output at the input resolution fig. 5b. The δ1-4 variant sums the outputs at s1 to s4 and drops the finest level fig. 5c §3.2.1.

Hyperparameters

Results

Author-reported. Setup: 108 T2-weighted MR and 3D TRUS image pairs from 76 patients in the SmartTarget trial, 834 landmark pairs, 12-fold patient-level cross-validation. The metrics are the RMS landmark error (TRE) and the Dice on prostate glands §3.1 §3.3. Baseline is the paper's own network; comparisons are to earlier networks and to pairwise intensity-based registration on the same data §3.2 §3.4.

Table 1 tab. 1: summary of the cross-validation results, reproduced from the paper. The medians in bold are significantly different from the baseline per the paper.

Significance (text, p-values as reported in the paper) §4.2 §4.3:

• msCE (cross-entropy instead of multiscale Dice): TRE worse, p < 0.001. The paper also reports a better single-scale DSC, p = 0.046.
• L2 (average L2 norm of displacement gradients instead of bending energy): TRE worse, p = 0.049; DSC worse, p < 0.001.
• preFilt (labels pre-filtered at all scales): no TRE effect (p = 0.458) and no DSC effect (p = 0.498). Training was about 25% faster in these experiments, depending on implementation.
• δ0: TRE worse (p = 0.034), DSC worse (p = 0.003).
• δ1-4: no significant difference (TRE p = 0.477; DSC p = 0.316).
• Without trilinear additive up-sampling: median TRE 6.4 mm, p < 0.001.
• Baseline vs Global-Net and Composite-Net: the baseline is better on TRE and DSC, both p < 0.001 §4.3.

Pairwise comparison §4.4 tab. 2: the paper reports pairwise intensity-based FFD registrations on this data, with and without rigid initialisation, for NMI, NCC and SSD. The paper says all nine medians are above 24 mm for TRE and below 0.77 for DSC. Its own Table 2 lists these numbers next to earlier MR-to-ultrasound results from other datasets, marked as not directly comparable. The table is reproduced in part below.

tab. 2: the paper's footnote says the first five rows come from simplified runs on its own data, and the rest are reported numbers from prior studies on different data.

Qualitative and Jacobian analysis: example slices fig. 9; boxplots of the cross-validation results fig. 8; Jacobian and displacement inspections fig. 10. In the Jacobian check, the baseline and δ0 gave no negative Jacobian determinants. An illustrative network with α = 0.01 gave negative determinants §4.3.

Paper vs code

No official code.

The paper says re-implementations of all the compared networks are available as part of NiftyNet §3.2, and gives niftynet.io as the project site. The NiftyNet repository (github.com/NifTK/NiftyNet, Apache-2.0) is owned by the NifTK organisation and is described as an unmaintained general platform for medical imaging, not as code for this paper. The owner is not verifiably one of the authors or an author's lab, so it is not cloned. Search done: the paper text, abstract and arXiv comments contain no repository link; no project page is linked; GitHub search by title words, by "label-driven weakly-supervised registration", by "weakly supervised registration prostate MR TRUS" returned zero repositories; GitHub search "NiftyNet" was checked for owner and licence §3.2.

Pitfalls

• Sign of the cross-entropy term. The text calls S_CE a negative cross-entropy §2.2. As printed, eq. 6 has no minus sign, so the implementer must choose the sign that makes the total loss decrease. The paper does not say which it used eq. 6.
• σ units are not stated. The Gaussian standard deviations in eq. 5 do not say whether they are in voxels or millimetres. Input images are resampled to 0.8 mm isotropic, so the two readings differ by a factor of 0.8 §2.2 §3.1.
• Initialisation conflict. §2.3 describes the displacement skip layers as randomly initialised with zero mean and small variance. §3.2.1 says the final displacement prediction layers were initialised to zero for all reported results §2.3 §3.2.1. The zero-init choice is what the reported numbers used.
• Learning-rate schedule. Adam starts at 1e-5 for the baseline and 1e-6 for Global-Net and Composite-Net, with no schedule or number of iterations stated §3.2.1 §3.2.2. Divergence is reported for the affine-based networks at higher rates §3.2.2.
• Regularisation weight is sensitive. With α = 0.5, the baseline produced no negative Jacobian determinants. With α = 0.01, the paper's illustrative network did §4.3. The weight was chosen without an exhaustive search and is not tuned against an unseen set §5.
• The unbiasedness argument covers one sampler. Eq. 3 shows the gradient is unbiased for the two-stage sampling that draws K image pairs and then one label pair per pair §2.1 eq. 3. The paper does not analyse other samplers, such as one drawing labels from the whole batch.
• Labels are sparse and uneven. Each image pair has its own label count M_n, and the label types differ by pair. Some labels overlap, such as gland and zonal regions §2.1 §3.1 §5.
• Memory and time. The baseline fits a 12 GB GPU. The Composite-Net needed a 24 GB card for the same minibatch §3.2 §3.2.2. Training ran for 48 hours §3.2.
• Comparisons across papers are weak. The paper says earlier MR-to-TRUS numbers come from different datasets, label sets and training strategies, so a direct numerical comparison may be unfair §4.4 §3.2.2.
• No separate held-out test set. The paper says a sizeable dataset unseen during method development was not available, so generalisation was not tested conclusively. Cross-validation on the trial data is the evidence, and generalisation to other centres is future work §5.
• Label quality. The landmark annotation took more than 200 hours, and TRUS gland segmentations were manually edited from automatic contours §3.1.
• Unstated details. Bending-energy formula, Gaussian kernel truncation, volume crop sizes, Adam β values and batch-norm ε are not given §2.1 §2.2 §3.2.
• Predicting without explicit topology. The DDF is not topology-preserving by construction. The authors say the relatively heavy regularisation avoided negative Jacobians in their tests, and suggest a Jacobian-based penalty for applications with larger deformations §5 §4.3.

Lineage

• Builds on: the authors' preliminary report of label-driven weakly-supervised learning for multimodal registration §1 (Hu et al., 2018). Changes in this paper: the multiscale Dice, the new network without the global affine sub-network, and a larger rigorous comparison §1.
• Compared against: the Global-Net and Composite-Net from Hu et al. (2018) §3.2.2 fig. 6 fig. 7; pairwise FFD registration with NMI, NCC and SSD §3.4; the prior MR-to-TRUS prostate studies in tab. 2.
• Superseded by: not stated in this card.