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

Conference: ECCV 2026
Paper: ECCV 2026 Poster
Area: Medical Imaging
Keywords: Image Registration, Sampled Importance Resampling, Uncertainty Quantification, Structured Covariance, High-Dimensional Probabilistic Inference

TL;DR

To overcome the limitations of variational inference in capturing multi-modal posteriors and high-dimensional spatial correlations in dense 3D image registration, this paper introduces Structured SIR—a framework combining Sampled Importance Resampling with a compact "low-rank plus locally sparse Cholesky precision factor" covariance parameterization that delivers expressive, multi-modal, and well-calibrated probabilistic inference across 1.6 million displacement dimensions.

Background & Motivation

Dense 3D medical image registration establishes voxel-wise spatial correspondences between images. Because underlying anatomical deformations exhibit intrinsic ambiguities and multiple plausible transformations can yield virtually identical similarity losses, the task is fundamentally ill-posed. Characterizing the full posterior distribution via uncertainty quantification is therefore critical for reliable downstream clinical decision-making. However, classical Markov Chain Monte Carlo (MCMC) and stochastic gradient sampling schemes incur prohibitive computational latency for routine practice, while amortized variational inference models enforce overly simplistic, factorized Gaussian posterior assumptions that lead to acute posterior variance underestimation, overconfidence, and an inherent inability to represent multi-modal anatomical hypotheses.

The core tension stems from the mismatch between the high dimensionality of 3D spatial deformations and the expressive capacity of tractable covariance models. Registering high-resolution 3D brain MRI volumes involves estimating displacement fields spanning over 1.6 million dimensions. Under such dimensional scale, directly parameterizing full-rank covariances triggers memory explosion, whereas sequential multi-scale sampling techniques like PULPo require iterative forward passes that tightly bind sampling runtime to neural network depth.

This paper tackles the challenge by decoupling amortized single-pass proposal parameterization from test-time evidence-weighted hypothesis selection. Core idea: propose the Structured SIR framework, which predicts a structured proposal distribution whose covariance decomposes into a low-rank component and a 3D spatially sparse Cholesky precision factor, then performs two-stage Sampled Importance Resampling (SIR) to filter high-quality hypotheses and backpropagate informative gradients without memory bottlenecks.

Method

Overall Architecture

The pipeline consists of an amortized proposal network and an importance-weighted resampling mechanism. Given a fixed volume \(I_f\) and moving volume \(I_m\) (discretized over \(96 \times 72 \times 80\) voxels, yielding \(N_v = 552,960 \times 3 \approx 1.6 \times 10^6\) displacement dimensions), a 4-level convolutional U-Net predicts the proposal distribution parameters \(q(Z) = \mathcal{N}(\mu, \Sigma)\) in a single forward pass: the mean field \(\mu\), a low-rank covariance factor \(R\), and a spatially sparse Cholesky precision factor \(L\). A large pool of candidate samples is drawn without gradients and evaluated under both normalized cross-correlation likelihood and diffusion regularizers. Multinomial resampling then extracts a representative subset of high-probability hypotheses for gradient backpropagation and final posterior characterization.

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["Input Image Pair<br/>Fixed Image If and Moving Image Im"] --> B["Structured Covariance Parameterization<br/>Predict mean μ, low-rank factor R, and sparse Cholesky factor L"]
    B --> C["No-Gradient Large-Scale Candidate Sampling<br/>Draw Ns Gaussian vectors and reparameterize via sparse solver"]
    C --> D["Dual-Evidence Importance Weighting<br/>Evaluate NCC likelihood and diffusion prior to form weights"]
    D --> E["Dynamic Temperature Multinomial Resampling<br/>Resample Nk high-probability hypotheses to backpropagate gradients"]
    E --> F["Calibrated Uncertainty & Multi-Modal Deformations"]

Key Designs

1. Structured Covariance Parameterization: Combining Global Coordination with Local Spatial Correlations

To prevent memory exhaustion across 1.6 million dimensions while maintaining rich covariance expressivity, the proposal covariance matrix is parameterized as the sum of a low-rank matrix and the inverse of a sparse Gaussian Markov Random Field (GMRF) precision factor: $\(\Sigma = R R^\top + (L L^\top)^{-1}\)$ Here, \(R \in \mathbb{R}^{3N_v \times r}\) (with rank \(r=10\) or \(25\)) captures large-scale, long-range anatomical correlations across brain regions, whereas \(L \in \mathbb{R}^{3N_v \times 3N_v}\) is a lower-triangular sparse matrix with positive diagonal entries, covering a \(3 \times 3 \times 3\) local spatial neighborhood with cross-channel displacement coupling. Leveraging the Woodbury Matrix Identity and the matrix determinant lemma, evaluating the quadratic Mahalanobis distance and log-determinant under this Gaussian reduces to an eigenvalue decomposition of an \(r \times r\) core matrix \(M = I_r + S^\top S\) (where \(S = L^\top R\)): $\((vec(Z) - \mu)^\top \Sigma^{-1} (vec(Z) - \mu) = k^\top k - (S^\top k)^\top M^{-1} (S^\top k)\)$ with \(k = L^\top (Z - \mu)\). This compact formulation scales as \(\mathcal{O}(\text{nnz}(L))\) for sparse triangular solves plus \(\mathcal{O}(r^3)\) for low-rank inversion, allowing 1,200 samples to be generated, weighted, and evaluated in under 5 seconds on a single GPU.

2. Sampled Importance Resampling (SIR): Decoupling Candidate Exploration from Gradient Propagation

Standard variational inference penalizes exploration by averaging over all proposal regions via the ELBO, causing mode collapse into a blurry mean. Conversely, backpropagating through all candidates in importance sampling exhausts GPU memory on low-probability samples. The proposed SIR resolves this by adopting a dual-pass strategy: first, \(N_s\) candidate fields are drawn without gradients via \(Z^i = \mu + R \varepsilon_R^i + L^{-\top} \varepsilon_L^i\) (where \(\varepsilon_R, \varepsilon_L \sim \mathcal{N}(0, I)\)), computing unnormalized importance weights against the true target posterior: $\(w_i \propto \frac{p(I_f | I_m \circ Z^i) p(Z^i)}{q(Z^i)}\)$ Second, multinomial resampling draws \(N_k\) indices according to the normalized weights. The network then reconstructs these selected hypotheses using their stored noise components and backpropagates gradients only through these high-quality solutions: $\(\mathcal{L}_{\text{SIR}} = -\frac{1}{N_k} \sum_{k=1}^{N_k} \left[ \log p(I_f | I_m \circ Z^k) + \log p(Z^k) \right]\)$ Analogous to the wake update in Reweighted Wake-Sleep, this scheme prevents entropy collapse, decouples candidate exploration from memory overhead, and guides the proposal distribution toward the true multi-modal posterior.

3. Dynamic Temperature Scaling and Asymmetric Prior Regularization

In high-dimensional vision tasks, slight variations in likelihood or prior densities easily cause importance weights to spike into a one-hot distribution or degenerate into uniform noise. To stabilize sampling, an exponential moving average (EMA) standard deviation of weights is incorporated into a dynamic temperature scaling formula: $\(\tilde{w}_i = w_i^{\frac{\bar{\sigma}_w}{T}}\)$ with momentum \(\phi=0.9\) and base temperature \(T=3.0\), balancing sample selectivity against distribution breadth. Furthermore, an asymmetric diffusion prior is imposed: $\(p(Z) \propto \exp\left( -\lambda_z \sum_{j=1}^{N_v} \|\nabla z_j\|^2 - \lambda_\mu \sum_{j=1}^{N_v} \|\nabla \mu_j\|^2 \right)\)$ Setting the regularization weight on the proposal mean field \(\lambda_\mu = 2.5\) higher than that on sampled perturbations \(\lambda_z = 1.0\) guarantees that \(\mu\) anchors a smooth, structurally sound global baseline, while sampled perturbations are granted sufficient variance to discover distinct anatomical modes (such as bilateral ventricles and cortical folds).

Loss & Training

The likelihood is modeled as a Boltzmann distribution over Normalized Cross-Correlation (NCC): $\(p(I_f | I_m, Z) \propto \exp\left( \frac{\text{NCC}(I_f, I_m \circ Z)}{\sigma^2} \right)\)$ with likelihood scaling parameter set to \(\sigma = 0.5\). Models are trained for 50,000 steps with batch size 1 on an NVIDIA RTX A6000 (48 GB VRAM), peaking at 32 GB. The SOAP optimizer (Heavyball library) is used with learning rate \(2 \times 10^{-3}\), momentum parameters \(\beta_1 = 0.9, \beta_2 = 0.95\), and a cosine scheduler decaying to \(2 \times 10^{-4}\) after 30,000 steps.

Key Experimental Results

Main Results

Evaluated on the OASIS 3D brain MRI benchmark across 2,450 subject-to-subject test pairs with 35 anatomical structures, measuring Dice Similarity Coefficients (DSC) for the proposal mean \(\mu\), resampled mean \(\bar{Z}\), and best-sample Oracle, alongside the non-positive Jacobian determinant rate (Fold %). Notation: V denotes Variational Inference, I denotes naive Importance Sampling, S denotes Structured SIR; D is diagonal Cholesky, C is sparse Cholesky, and L is the low-rank component.

Model	\(\mu\) (DSC)	\(\bar{Z}\) (DSC)	Oracle (DSC)	\(\mu\) Fold %	\(\bar{Z}\) Fold %	Oracle Fold %
V+D	0.800 (0.04)	0.800 (0.04)	0.805 (0.04)	0.046 (0.04)	0.059 (0.05)	0.059 (0.05)
V+LD	0.799 (0.04)	0.798 (0.04)	0.805 (0.04)	0.046 (0.04)	0.059 (0.05)	0.059 (0.05)
V+C	0.802 (0.04)	0.802 (0.04)	0.807 (0.04)	0.045 (0.05)	0.057 (0.06)	0.057 (0.06)
V+LC	0.800 (0.04)	0.800 (0.04)	0.806 (0.04)	0.044 (0.03)	0.056 (0.04)	0.056 (0.04)
I+LC	0.800 (0.04)	0.800 (0.04)	0.805 (0.04)	0.050 (0.05)	0.063 (0.06)	0.063 (0.06)
S+D	0.801 (0.04)	0.801 (0.04)	0.805 (0.04)	0.048 (0.05)	0.062 (0.06)	0.062 (0.06)
S+LD	0.798 (0.04)	0.800 (0.04)	0.809 (0.04)	0.032 (0.03)	0.085 (0.08)	0.096 (0.10)
S+C	0.801 (0.04)	0.801 (0.04)	0.805 (0.04)	0.051 (0.03)	0.065 (0.04)	0.065 (0.04)
S+LC	0.794 (0.04)	0.801 (0.04)	0.812 (0.03)	0.011 (0.01)	0.115 (0.07)	0.163 (0.13)
S+LC (10L)	0.788 (0.04)	0.802 (0.04)	0.814 (0.03)	0.004 (0.01)	0.240 (0.14)	0.292 (0.25)
Vxm5 (5-model ensemble)	0.798 (0.03)	0.798 (0.03)	0.804 (0.03)	0.000 (0.00)	0.000 (0.00)	0.000 (0.00)

Ablation Study

Probabilistic calibration metrics evaluated within a 3-voxel dilated boundary around anatomical labels: Area Under the Sparsification Error (AUSE, lower is better), Expected Calibration Error (ECE, lower is better), mean label entropy, and Spearman rank correlation between label entropy and registration accuracy (DSC, where negative values indicate reliable uncertainty proxies):

Model Variant	AUSE \(\downarrow\)	ECE \(\downarrow\)	Mean Label Entropy	Spearman \(r\) (Entropy vs DSC)
V+D	0.0850 (0.021)	0.1285 (0.027)	0.0628 (0.002)	+0.064 (0.169)
V+LD	0.0818 (0.022)	0.1257 (0.027)	0.0705 (0.002)	+0.085 (0.165)
V+C	0.0854 (0.021)	0.1279 (0.026)	0.0595 (0.002)	+0.046 (0.164)
V+LC	0.0844 (0.021)	0.1274 (0.027)	0.0641 (0.002)	+0.033 (0.167)
I+LC	0.0861 (0.021)	0.1287 (0.027)	0.0621 (0.002)	+0.011 (0.172)
S+D	0.0897 (0.021)	0.1312 (0.026)	0.0549 (0.007)	+0.051 (0.123)
S+LD	0.0736 (0.020)	0.1159 (0.026)	0.0901 (0.004)	-0.137 (0.177)
S+C	0.0907 (0.021)	0.1324 (0.027)	0.0531 (0.006)	+0.007 (0.121)
S+LC (25L)	0.0659 (0.019)	0.1064 (0.025)	0.1079 (0.006)	-0.325 (0.218)
S+LC (10L)	0.0636 (0.018)	0.1028 (0.023)	0.1156 (0.010)	-0.419 (0.217)
Vxm5 (5-model ensemble)	0.1008 (0.015)	0.1224 (0.020)	0.0703 (0.007)	-0.410 (0.209)

Key Findings

• Substantial Oracle Selection Gap Reveals Expressive Hypotheses: In variational baselines, enriching covariance structures yields almost zero Oracle gain (0.805 vs 0.806 DSC). Under Structured SIR, S+LC (10L) expands the Oracle DSC to 0.814—a notable 0.026 margin above the proposal mean. This demonstrates that highly accurate anatomical solutions exist with substantial probability in the proposal distribution and are effectively uncovered via importance resampling.
• Uncertainty Calibration Realizes a Qualitative Breakthrough: Variational baselines and naive importance sampling exhibit counter-intuitive positive correlations between label entropy and DSC (+0.033 to +0.085), indicating broken uncertainty estimates. S+LC (10L) flips the correlation to a solid -0.419, while simultaneously reducing AUSE from 0.0850 to 0.0636 and trimming ECE by ~20%, delivering dependable clinical confidence intervals.
• Compact Rank Outperforms Over-Parameterized Covariances: Reducing the rank from 25 to 10 improves both registration accuracy (0.802 vs 0.801 mean DSC; 0.814 vs 0.812 Oracle DSC) and calibration (ECE drops from 0.1064 to 0.1028), while accelerating sampling by reducing matrix decomposition overhead.

Highlights & Insights

• Woodbury Decomposition on 1.6M-Dimensional Operators: Solving dense Gaussian inference over 1.6 million dimensions via the sum of a low-rank covariance and a sparse Cholesky precision factor transforms an intractable inversion into an \(r \times r\) eigendecomposition and sparse triangular back-substitution.
• Two-Pass Decoupled Resampling Avoids Gradient Bottlenecks: Exploring 1,200 candidates in a gradient-free first pass and backpropagating through a resampled subset allows extensive coverage of complex spatial deformation manifolds without inflating VRAM.
• Uncovering Distinct Anatomical Multi-Modality: PCA and clustering on resampled deformation fields reveal structured bi-modal distributions around the thalamus and lateral ventricle borders, where competing modes achieve equivalent DSC (~0.86) by fitting alternate tissue boundaries.

Limitations & Future Work

• Folding Rate Trade-off: As the proposal distribution explores broader displacement manifolds, the proportion of voxels with negative Jacobian determinants increases to ~0.24% in S+LC (10L). Future extensions could adopt Stationary Velocity Fields (SVF) or hyperelastic regularizers to ensure diffeomorphic guarantees.
• Unsupervised Likelihood Vulnerability: Relying strictly on NCC image similarity can misguide importance weights in low-contrast tissue or pathological lesion areas. Integrating weak segmentation supervision into the training-time likelihood could better shape proposal geometry.

Related Work & Insights

• vs. VoxelMorph (Amortized VI): VoxelMorph relies on diagonal Gaussian assumptions that fail to model inter-voxel spatial correlations and collapse to a single unimodal mean. Structured SIR retains single-pass efficiency while introducing rich spatial covariances and multi-modal posterior coverage.
• vs. PULPo: PULPo captures multi-scale uncertainty through an iterative Laplacian pyramid, incurring multi-step forward-pass latency. Structured SIR performs a single amortized prediction and solves 1,200 samples via sparse solvers in under 5 seconds.
• vs. SG-MCMC Registration: Stochastic gradient MCMC requires sequential test-time Markov chain iterations taking tens of seconds per volume. Structured SIR provides a fast amortized proposal with single-pass importance re-weighting suitable for high-throughput pipelines.

Rating

• Novelty: ⭐⭐⭐⭐⭐ Pioneering application of Sampled Importance Resampling combined with structured low-rank plus sparse Cholesky covariance in million-dimensional dense vision inference.
• Experimental Thoroughness: ⭐⭐⭐⭐⭐ Rigorous validation across 2,450 subject-to-subject brain MRI pairs, detailed AUSE/ECE calibration curves, and multi-modal mode inspections.
• Writing Quality: ⭐⭐⭐⭐⭐ Clean mathematical formulations, elegant Woodbury reductions, and thorough comparative analyses.
• Value: ⭐⭐⭐⭐⭐ Establishes a highly generalizable inference paradigm for high-dimensional ill-posed inverse problems, including medical registration, optical flow, and volumetric reconstruction.

Related Papers

• [ECCV 2026] SAM+D: Parameter-Efficient Dimensional Lifting of SAM-Family Models via Depth-Routed LoRA and Depth Shifting
• [ECCV 2026] Foundation-Guided Representation Alignment for Multimodal Medical Image Registration
• [ECCV 2026] Two-Parameter Flow Map Learning for Continuous-Time Diffeomorphic Image Registration
• [ECCV 2026] MedRepBench: Benchmarking Structured Understanding of Medical Report Images
• [ECCV 2026] Multi-Channel Uncertainty-Weighted Score Matching for Conditional Diffusion in Medical UDA