Adaptive Wavelet Decomposition with Cross-Band Attention for Unsupervised 3D Medical Image Registration

Abstract

This report introduces a new unsupervised deformable 3D registration network called AWaRe Net developed for the purpose of overcoming the drawbacks of existing techniques with regard to multi scale fusion and losing key information. AWaRe Net replaces fixed Haar wavelets with the ability to create adaptive wavelet filters through a technique called Adaptive Learnable Wavelet Decomposition (ALWD) that allows each filter to have parameters optimized using backpropagation while simultaneously allowing for attention to dynamically adjust weights across sub bands using CrossBand Attention Fusion (CBAF). Additionally, a hybrid lightweight bottleneck architecture that combines ConvNexT and axial attention is used to create long range dependencies between data points in a Linearithmic (i.e, O(N log N)) fashion, whilst an explicit frequency loss provides additional alignment of wavelets and other properties found within traditional 2D methods through conventional spatial similarity. AWaRe Net has been evaluated against two publicly available, currently utilized datasets: IXI (image registration from Atlas to Patient) and OASIS (inter-patient), and demonstrates mean values of 79.1% and 83.9% Dice scores, respectively which significantly outperform comparable algorithms (Wave Morph, Trans Morph, and Voxel Morph) (p < 0.001). This allowed for achieving a lower folding ratio of approx. 0.168% relative to previously published algorithms on the same dataset (OASIS), with an average inference time of 68 ms per image pair, and only 0.82 million parameters. Each proposed component was confirmed through an ablation study. Additionally, generalization tests on synthetic multi organ datasets have demonstrated enhanced transferability of registration results to abdominal, cardiac and lung images. This study demonstrates that the use of adaptive wavelet decomposition with cross band attention substantially increases accuracy within the registration process, the smoothness of deformation across the registered image, as well as the computational efficiency of the entire process as it relates to clinical deployment and longitudinal disease monitoring in real time.