Unsupervised Deep Learning Network with Self-Attention Mechanism for Non-Rigid Registration of 3D Brain MR Images

Abstract

In non-rigid registration for medical imaging analysis, computation is complicated, and the high accuracy and robustness needed for registration are difficult to obtain. Recently, many studies have been conducted for nonrigid registration via unsupervised learning networks. This study proposes a method to improve the performance of this unsupervised learning network approach, through the use of a self-attention mechanism. In this paper, the self-attention mechanism is combined with deep learning networks to identify information of higher importance, among large amounts of data, and thereby solve specific tasks. Furthermore, the proposed method extracts both local and non-local information so that the network can create feature vectors with more information. As a result, the limitation of the existing network is addressed: alignment based solely on the entire silhouette of the brain is mitigated in favor of a network which also learns to perform registration of the parts of the brain that have internal structural characteristics. To the best of our knowledge, this is the first such utilization of the attention mechanism in this unsupervised learning network for non-rigid registration. The proposed attention network performs registration that takes into account the overall characteristics of the data, thus yielding more accurate matching results than those of the existing methods. In particular, matching is achieved with especially high accuracy in the gray matter and cortical ventricle areas, since these areas contain many of the structural features of the brain. The experiment was performed on 3D magnetic resonance images of the brains of 50 people. The measured average dice similarity coefficient after registration was 70.40%, which is an improvement of 17.48% compared to that before registration. This improvement indicates that application of the attention block can further improve the performance by an additional 8.5%, as relative to that without attention block. Ultimately, through implementation of non-rigid registration via the attention block method, the internal structure and overall shape of the brain can be addressed, without additional data input. Additionally, attention blocks have the advantage of being able to easily connect to existing networks without a significant computational overhead. Furthermore, by producing an attention map, the area of the brain around which registration was more performed can be visualized. This approach can be used for non-rigid registration with various types of medical imaging data.