Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Toward Universal Medical Image Registration via Sharpness-Aware Meta-Continual Learning

Bomin Wang, Xinzhe Luo, Xiahai Zhuang

TL;DR

This paper presents the first attempt to achieve the goal of universal 3D medical image registration in sequential learning scenarios by proposing a continual learning method that utilizes meta-learning with experience replay to mitigating the problem of catastrophic forgetting.

Abstract

Current deep learning approaches in medical image registration usually face the challenges of distribution shift and data collection, hindering real-world deployment. In contrast, universal medical image registration aims to perform registration on a wide range of clinically relevant tasks simultaneously, thus having tremendous potential for clinical applications. In this paper, we present the first attempt to achieve the goal of universal 3D medical image registration in sequential learning scenarios by proposing a continual learning method. Specifically, we utilize meta-learning with experience replay to mitigating the problem of catastrophic forgetting. To promote the generalizability of meta-continual learning, we further propose sharpness-aware meta-continual learning (SAMCL). We validate the effectiveness of our method on four datasets in a continual learning setup, including brain MR, abdomen CT, lung CT, and abdomen MR-CT image pairs. Results have shown the potential of SAMCL in realizing universal image registration, which performs better than or on par with vanilla sequential or centralized multi-task training strategies.The source code will be available from https://github.com/xzluo97/Continual-Reg.

Toward Universal Medical Image Registration via Sharpness-Aware Meta-Continual Learning

TL;DR

This paper presents the first attempt to achieve the goal of universal 3D medical image registration in sequential learning scenarios by proposing a continual learning method that utilizes meta-learning with experience replay to mitigating the problem of catastrophic forgetting.

Abstract

Current deep learning approaches in medical image registration usually face the challenges of distribution shift and data collection, hindering real-world deployment. In contrast, universal medical image registration aims to perform registration on a wide range of clinically relevant tasks simultaneously, thus having tremendous potential for clinical applications. In this paper, we present the first attempt to achieve the goal of universal 3D medical image registration in sequential learning scenarios by proposing a continual learning method. Specifically, we utilize meta-learning with experience replay to mitigating the problem of catastrophic forgetting. To promote the generalizability of meta-continual learning, we further propose sharpness-aware meta-continual learning (SAMCL). We validate the effectiveness of our method on four datasets in a continual learning setup, including brain MR, abdomen CT, lung CT, and abdomen MR-CT image pairs. Results have shown the potential of SAMCL in realizing universal image registration, which performs better than or on par with vanilla sequential or centralized multi-task training strategies.The source code will be available from https://github.com/xzluo97/Continual-Reg.

Paper Structure

This paper contains 13 sections, 6 equations, 5 figures, 1 table, 1 algorithm.

Table of Contents

  1. Introduction
  2. Methodology
  3. Sharpness-Aware Meta-Continual Learning (SAMCL)
  4. Toward Universal Medical Image Registration.
  5. Experiments and Results
  6. Materials and Compared Methods.
  7. Evaluation Metrics.
  8. Implementation Details.
  9. Comparison Study.
  10. Ablation and Generalization Study.
  11. Conclusion and Discussion
  12. Acknowledgments.
  13. Disclosure of Interests.

Figures (5)

  • Figure 1: Continual image registration. One can see that the image pairs differ significantly in appearance and anatomy across tasks, i.e., serious distribution shift, posing a major challenge to modern learning-based registration methods.
  • Figure 2: Registration results on an exemplar OASIS and Abdomen CT image pair using different training strategies. The top row visualizes the difference map between the fixed image and the registered moving image. Note that CT intensities were clipped to the range $[-200, 300]$.
  • Figure 3: Ablation and generalization study of the proposed SAMCL on (a) using various memory buffer sizes for experience replay, and (b) using sharpness-aware minimization (SAM) for in- and out-of-domain generalization tasks.
  • Figure 4: Violin plot to show the registration performance of different methods on (a)OASIS, (b)Abdomen CT-CT, (c)NLST and (d) Abdomen MR-CT.
  • Figure 5: Registration results on an exemplar NLST image pair using different training strategies. The top row visualizes the difference map between the fixed image and the registered moving image.