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

Image registration is an ill-posed dense vision task, wheremultiple solutions achieve similar loss values, motivating probabilisticinference. Variational inference has previously been employed to capturethese distributions, however restrictive assumptions about the posteriorform can lead to poor characterisation, overconfidence and low-qualitysamples. More flexible posteriors are typically bottlenecked by the com-plexity of high-dimensional covariance matrices required for dense 3Dimage registration.In this work, we present a memory and computationally e!cient in-ference method, Structured SIR, that enables expressive, multi-modal,characterisation of uncertainty with high quality samples. We proposethe use of a Sampled Importance Resampling (SIR) algorithm witha novel memory-e!cient high-dimensional covariance parameterisationas the sum of a low-rank covariance and a sparse, spatially structuredCholesky precision factor. This structure enables capturing complex spa-tial correlations while remaining computationally tractable.We evaluate the e!cacy of this approach in 3D dense image registrationof brain MRI data, which is a very high-dimensional problem. We demon-strate that our proposed method produces uncertainty estimates that aresignificantly better calibrated than those produced by variational meth-ods, achieving equivalent or better accuracy. Crucially, we show thatthe model yields highly structured multi-modal posterior distributions,enable e"ective and e!cient uncertainty quantification.