Uncertainty-Aware Test-Time Adaptation for Inverse Consistent Diffeomorphic Lung Image Registration

TL;DR

This work proposes an uncertainty-aware test-time adaptation framework for inverse consistent diffeomorphic lung registration that uses Monte Carlo (MC) dropout to estimate spatial uncertainty that is used to improve model performance.

Abstract

Diffeomorphic deformable image registration ensures smooth invertible transformations across inspiratory and expiratory chest CT scans. Yet, in practice, deep learning-based diffeomorphic methods struggle to capture large deformations between inspiratory and expiratory volumes, and therefore lack inverse consistency. Existing methods also fail to account for model uncertainty, which can be useful for improving performance. We propose an uncertainty-aware test-time adaptation framework for inverse consistent diffeomorphic lung registration. Our method uses Monte Carlo (MC) dropout to estimate spatial uncertainty that is used to improve model performance. We train and evaluate our method for inspiratory-to-expiratory CT registration on a large cohort of 675 subjects from the COPDGene study, achieving a higher Dice similarity coefficient (DSC) between the lung boundaries (0.966) compared to both VoxelMorph (0.953) and TransMorph (0.953). Our method demonstrates consistent improvements in the inverse registration direction as well with an overall DSC of 0.966, higher than VoxelMorph (0.958) and TransMorph (0.956). Paired t-tests indicate statistically significant improvements.

Figures (4)

• Figure 1: Inverse consistent test-time adaptation framework. During training, the network $\mathcal{G}_{\theta}$ learns to predict SVF $\boldsymbol{v}$, from a fixed image $\boldsymbol{I}_\mathrm{F} = \boldsymbol{I}_\mathrm{FRC}$ and a moving image $\boldsymbol{I}_\mathrm{M} = \boldsymbol{I}_\mathrm{TLC}$. The SVF is then integrated using the scaling and squaring (SS) method to get the final displacement field $\mathbf{\Phi}_1$arsigny2006log. During inference, two different pathways (forward and inverse) are defined for uncertainty estimation and adaptation, details of which are described below. The LDDMM framework allows for adaptation in inverse direction by simply negating the SVF ($-\boldsymbol{v}$) and then integrating it through SS.
• Figure 2: Qualitative visualization of two large deformation forward registration cases with $\delta V = 1.1$ L and $\delta V = 1.6$ L. The fixed image outlines, shown in yellow, are overlaid on the deformed images for references. Change in volume was defined in liters (L).
• Figure 3: Visualization of inverse image registration with large $\delta V = 1.1$ L and $\delta V = 1.6$ L. The inverse transformation $\mathbf{\Phi}_{1}^{-1}$ was obtained by the integration of negative SVF $-\boldsymbol{v}$.
• Figure 4: Change in lung mask overlap metrics [(A) DSC and (B) ASSD] for TLC to FRC (forward) with increasing $\delta V$s. Regression lines for each method and their slopes indicate improved model performance with increasing test-time adaptation steps $\mathcal{T}$.