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

TL;DR

This paper introduces Structured SIR, a memory-efficient inference method for high-dimensional image registration that captures complex, multi-modal uncertainty distributions more accurately than traditional variational approaches.

Contribution

The paper presents a novel high-dimensional covariance parameterisation and a Sampled Importance Resampling algorithm for efficient, expressive uncertainty quantification in dense 3D image registration.

Findings

Produces better calibrated uncertainty estimates than variational methods

Achieves equivalent or improved accuracy in brain MRI registration

Enables effective modeling of multi-modal posterior distributions

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.