Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Structured SIR: Efficient and Expressive Importance-Weighted Inference for High-Dimensional Image Registration

Ivor J. A. Simpson, Neill D. F. Campbell

Abstract

Image registration is an ill-posed dense vision task, where multiple solutions achieve similar loss values, motivating probabilistic inference. Variational inference has previously been employed to capture these distributions, however restrictive assumptions about the posterior form can lead to poor characterisation, overconfidence and low-quality samples. More flexible posteriors are typically bottlenecked by the complexity of high-dimensional covariance matrices required for dense 3D image registration. In this work, we present a memory and computationally efficient inference method, Structured SIR, that enables expressive, multi-modal, characterisation of uncertainty with high quality samples. We propose the use of a Sampled Importance Resampling (SIR) algorithm with a novel memory-efficient high-dimensional covariance parameterisation as the sum of a low-rank covariance and a sparse, spatially structured Cholesky precision factor. This structure enables capturing complex spatial correlations while remaining computationally tractable. We evaluate the efficacy of this approach in 3D dense image registration of brain MRI data, which is a very high-dimensional problem. We demonstrate that our proposed methods produces uncertainty estimates that are significantly better calibrated than those produced by variational methods, achieving equivalent or better accuracy. Crucially, we show that the model yields highly structured multi-modal posterior distributions, enable effective and efficient uncertainty quantification.

Structured SIR: Efficient and Expressive Importance-Weighted Inference for High-Dimensional Image Registration

Abstract

Image registration is an ill-posed dense vision task, where multiple solutions achieve similar loss values, motivating probabilistic inference. Variational inference has previously been employed to capture these distributions, however restrictive assumptions about the posterior form can lead to poor characterisation, overconfidence and low-quality samples. More flexible posteriors are typically bottlenecked by the complexity of high-dimensional covariance matrices required for dense 3D image registration. In this work, we present a memory and computationally efficient inference method, Structured SIR, that enables expressive, multi-modal, characterisation of uncertainty with high quality samples. We propose the use of a Sampled Importance Resampling (SIR) algorithm with a novel memory-efficient high-dimensional covariance parameterisation as the sum of a low-rank covariance and a sparse, spatially structured Cholesky precision factor. This structure enables capturing complex spatial correlations while remaining computationally tractable. We evaluate the efficacy of this approach in 3D dense image registration of brain MRI data, which is a very high-dimensional problem. We demonstrate that our proposed methods produces uncertainty estimates that are significantly better calibrated than those produced by variational methods, achieving equivalent or better accuracy. Crucially, we show that the model yields highly structured multi-modal posterior distributions, enable effective and efficient uncertainty quantification.
Paper Structure (38 sections, 16 equations, 9 figures, 3 tables)

This paper contains 38 sections, 16 equations, 9 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Background
  3. Registration Notation
  4. Probabilistic Registration
  5. Sampling and Approximate Inference
  6. Amortised Variational Inference
  7. Uncertainty Calibration
  8. Importance Sampling
  9. Flexible Covariances
  10. Methods
  11. Sampled Importance Resampling (SIR)
  12. Likelihood Function
  13. Prior Function
  14. Covariance Parameterisation
  15. Network Architecture
  16. ...and 23 more sections

Figures (9)

  • Figure 1: Illustration of the distribution modelling problem, where the underlying probability distribution is illustrated in (\ref{['fig:overview_true_dist']}). (\ref{['fig:overview_spherical_sampling']}) shows an attempt to model this variationally using a spherical Gaussian, which assigns probability mass to a lot of area that is unlikely, but simultaneously does not cover the whole distribution; if used as a proposal distribution then many poor samples are drawn and this is exacerbated as dimensionality increases. (\ref{['fig:overview_supn_sampling']}) illustrates how our SIR sampling framework allows selection of high probability samples at both learning/inference time, which avoids penalising a model for being more exploratory. Moreover, using a structured Gaussian allows assigning probability mass to higher density areas improving the efficiency of sampling.
  • Figure 2: Left, an illustration of the proposal mean ($\boldsymbol{\mu}$ in SIR), variational $\boldsymbol{\mu}$, SIR sample mean and Oracle (best-sample), and standard deviation of the segmentation accuracy for different model variants. Right, a boxplot exemplifying the quality of the Oracle sample drawn from the approximate posterior distribution. Larger Dice Similarity Coefficient (DSC) values indicated greater segmentation accuracy.
  • Figure 3: Accuracy vs calibration for the different model variants for a set of structures, standard deviations plotted in gray. The best performance is in the bottom-right of the graph with a low ECE (better calibration) and a high DSC (better accuracy). The full version of the model (S+LC) offers improved calibration with equal or improved accuracy consistently over all structures.
  • Figure 4: Illustration of the structured nature of the samples, where we find the performance of the Oracle sample can lead to substantial structured gains over the average SIR deformation. On the left, we show an example where a lobe of the cerebral cortex is correctly aligned by the oracle sample, with a DSC improvement of 0.04. On the right, we observe substrantial structured changes in the lateral ventricle, with a DSC improvements of 0.03.
  • Figure 5: Illustration of the multi-modality of the posterior distribution, where disagreements between modes are marked in blue or orange. On the left, we find structured changes in the Thalamus where mode 1 has a DSC of 0.834 and mode 2 has a Dice of 0.824. On the right, we observe substantial differences in the left lateral ventricle, either on the interior or exterior border, with a DSC of 0.864 for mode 1 and 0.862 for mode 2.
  • ...and 4 more figures