DiffuseMorph: Unsupervised Deformable Image
Registration Using Diffusion Model
Boah Kim, Inhwa Han, and Jong Chul Ye
DiffuseMorph: Unsupervised Deformable Image Registration Using Diffusion Model
Research objectives
Deformable image registration
• To estimate non-rigid voxel correspondences between moving & fixed images
• Important for medical image analysis (ex. disease diagnosis)
Multiphase abdominal CT
PET MRI
Types Advantage Limitation
Classical iterative method High topology preservation Substantial time, high computational cost
Supervised learning No parameter tuning Depending on the ground-truth fields
Unsupervised learning No parameter tuning + No labels Potential for folding problem on the fields
[1] Du, Juan, et al., International Journal of Computer Mathematics, 2006.
2
DiffuseMorph: Unsupervised Deformable Image Registration Using Diffusion Model
Research objectives
Deformable image registration
• To estimate non-rigid voxel correspondences between moving & fixed images
• Important for medical image analysis (ex. disease diagnosis)
Multiphase abdominal CT
PET MRI
• To perform image registration along the continuous trajectory
• To apply the denoising diffusion probabilistic model to image registration
Research goal
[1] Du, Juan, et al., International Journal of Computer Mathematics, 2006.
3
DiffuseMorph: Unsupervised Deformable Image Registration Using Diffusion Model
Background: Deformable image registration
• Deformation field: A vector field of displacement vectors for all coordinates in images
• Spatial transformation: Grid sampling to warp moving image into fixed image (ex. linear interpolation)
𝜙∗
= argmin
𝜙
𝑳𝒔𝒊𝒎 𝑚 𝜙 , 𝑓 + 𝑳𝒓𝒆𝒈(𝜙)
Similarity Regularization
Optimization problem in classical algorithms
• Similarity: To evaluate the shape differences between
deformed images and fixed images
• Regularization: To penalize deformation fields
ST: Spatial transformation
𝑓
Fixed image
𝑚
Moving image
ST
𝜙
Deformation field
𝑚(𝜙)
Deformed image
4
DiffuseMorph: Unsupervised Deformable Image Registration Using Diffusion Model
Background: Denoising diffusion probabilistic model (DDPM)
[1] Jonathan Ho et al., NeurIPS 2020.
• To convert a data distribution to Gaussian noise
Forward diffusion process (𝒙𝟎 → 𝒙𝑻)
Gaussian transition for each target data 𝒙0:
⇒ 𝑝 𝒙𝑡 𝒙0 = 𝒩 𝒙𝑡; 𝛼𝑡𝒙0, 1 − 𝛼𝑡 𝑰
where 𝛼𝑡 = ς𝑠=1
𝑡
(1 − 𝛽𝑠)
𝑝 𝒙𝑡 𝒙𝑡−1 = 𝒩(𝒙𝑡; 1 − 𝛽𝑡𝒙𝑡−1, 𝛽𝑡𝑰)
• To sample image from Gaussian noise
Reverse generative process (𝒙𝑻 → 𝒙𝟎)
Parameterized Gaussian process:
𝑝𝜃 𝒙𝑡−1 𝒙𝑡 = 𝒩 𝒙𝑡−1; 𝜇𝜃 𝒙𝑡, 𝑡 , 𝜎𝑡
2
𝑰
𝒙𝑁~𝒩(0, 𝑰)
𝒙𝑡−1 =
1
1−𝛽𝑡
𝒙𝑡 −
𝛽𝑡
1−𝛼𝑡
𝝐𝜃 𝒙𝑡, 𝑡 + 𝜎𝑡𝒛
Generation for 𝑡 = 𝑁, 𝑁 − 1, … , 1:
Parameterized model
Image perturbation: 𝒙𝑡 = 𝛼𝑡𝒙0 + 1 − 𝛼𝑡𝝐
5
DiffuseMorph: Unsupervised Deformable Image Registration Using Diffusion Model
Proposed method: DiffuseMorph
Training phase
min
𝐺𝜃,𝑀𝜓
𝑳𝒅𝒊𝒇𝒇𝒖𝒔𝒊𝒐𝒏 𝑐, 𝑥𝑡, 𝑡 + 𝜆𝑳𝒓𝒆𝒈𝒊𝒔𝒕(𝑚, 𝑓)
𝑳𝒅𝒊𝒇𝒇𝒖𝒔𝒊𝒐𝒏 𝑐, 𝑥𝑡, 𝑡 = 𝔼𝜖,𝑥𝑡,𝑡 𝐺𝜃 𝑐, 𝑥𝑡, 𝑡 − 𝜖 2
2
𝑳𝒓𝒆𝒈𝒊𝒔𝒕 𝑚, 𝑓 = − 𝑚 𝜙 ⊗ 𝑓 + 𝜆𝜙∑ 𝛻𝜙 2
• Diffusion network: To estimate a conditional score function
• Deformation network: To yield the registration field & provide the deformed image
Training in an end-to-end learning manner
Loss function
6
DiffuseMorph: Unsupervised Deformable Image Registration Using Diffusion Model
Proposed method: DiffuseMorph
Inference phase
 Continuous image registration  Synthetic image generation
7
DiffuseMorph: Unsupervised Deformable Image Registration Using Diffusion Model
Experiment 1
Intra-subject 2D face image registration
[1] Balakrishnan, G. et al., CVPR 2018. [2] Dalca, A.V. MICCAI 2018.
8
Methods NMSE x10−1
SSIM 𝐽𝜙 ≤ 0 (%)
Initial 0.363 (0.268) 0.823 (0.066) -
VM [1] 0.047 (0.057) 0.936 (0.024) 0.050 (0.106)
VM-diff [2] 0.034 (0.015) 0.957 (0.013) 0.014 (0.065)
Ours 0.032 (0.017) 0.964 (0.011) 0.017 (0.056)
• RaFD dataset (3 different gazed 8 expressions / Input size = 128x128)
• Training : Validation: Test = 53 : 7 : 7 (subjects)
DiffuseMorph: Unsupervised Deformable Image Registration Using Diffusion Model
Experiment 2
Intra-subject 3D cardiac MR image registration
[1] Balakrishnan, G. et al., CVPR 2018. [2] Dalca, A.V. MICCAI 2018.
9
Methods Dice 𝐽𝜙 ≤ 0 (%)
Initial 0.642 (0.188) -
VM [1] 0.787 (0.113) 0.169 (0.109)
VM-diff [2] 0.794 (0.104) 0.291 (0.188)
Ours 0.802 (0.109) 0.161 (0.082)
• ACDC dataset (4D temporal cardiac MRI data/ input size = 128×128×32)
• Training : Test = 90 : 10 (scans)
DiffuseMorph: Unsupervised Deformable Image Registration Using Diffusion Model
Experiment 3
Atlas-based 3D Brain MR Image Registration
10
Methods Dice 𝐽𝜙 ≤ 0 (%) Time (min)
Initial 0.616 (0.171) - -
SyN [1] 0.752 (0.140) 0.400 (0.100) 122, CPU
VM [2] 0.749 (0.145) 0.553 (0.075) 0.01, GPU
VM-diff [3] 0.731 (0.139) 0.631 (0.073) 0.01, GPU
SYMNet [4] 0.733 (0.148) 0.547 (0.049) 0.43, GPU
MSDIRNet [5] 0.751 (0.142) 0.804 (0.089) 2.06, GPU
CM [6] 0.750 (0.144) 0.510 (0.087) 0.01, GPU
Ours 0.756 (0.139) 0.505 (0.058) 0.01, GPU
• OASIS-3 dataset (T1-weights MRI data / input size = 160×192×224)
• Training : Validation : Test = 1027 : 93 : 129 (scans)
[1] Avants, B.B.et al., Medical image analysis, 2008 [2] Balakrishnan, G. et al., CVPR 2018. [3] Dalca, A.V. MICCAI 2018.
[4] Mok, T.C. et al., CVPR 2020 [5] Lei, Y. et al., Physics in Medicine & Biology, 2020. [6] Kim, B. et al., Medical image analysis 2021.
DiffuseMorph: Unsupervised Deformable Image Registration Using Diffusion Model
Conclusion
• DiffuseMorph learns the conditional score function of deformation
→ To generate synthetic deformed images
→ To provide high-quality image registration from the continuous deformation
• Experimental results on 2D and 3D data suggest the superiority of our method in image registration.
• DiffuseMorph can be used to generate temporal data from the moving to the fixed images.
Propose a novel diffusion-based unsupervised image registration model, called DiffuseMorph.
11
Thank you.
12

DiffuseMorph: Unsupervised Deformable Image Registration Using Diffusion Model

  • 1.
    DiffuseMorph: Unsupervised DeformableImage Registration Using Diffusion Model Boah Kim, Inhwa Han, and Jong Chul Ye
  • 2.
    DiffuseMorph: Unsupervised DeformableImage Registration Using Diffusion Model Research objectives Deformable image registration • To estimate non-rigid voxel correspondences between moving & fixed images • Important for medical image analysis (ex. disease diagnosis) Multiphase abdominal CT PET MRI Types Advantage Limitation Classical iterative method High topology preservation Substantial time, high computational cost Supervised learning No parameter tuning Depending on the ground-truth fields Unsupervised learning No parameter tuning + No labels Potential for folding problem on the fields [1] Du, Juan, et al., International Journal of Computer Mathematics, 2006. 2
  • 3.
    DiffuseMorph: Unsupervised DeformableImage Registration Using Diffusion Model Research objectives Deformable image registration • To estimate non-rigid voxel correspondences between moving & fixed images • Important for medical image analysis (ex. disease diagnosis) Multiphase abdominal CT PET MRI • To perform image registration along the continuous trajectory • To apply the denoising diffusion probabilistic model to image registration Research goal [1] Du, Juan, et al., International Journal of Computer Mathematics, 2006. 3
  • 4.
    DiffuseMorph: Unsupervised DeformableImage Registration Using Diffusion Model Background: Deformable image registration • Deformation field: A vector field of displacement vectors for all coordinates in images • Spatial transformation: Grid sampling to warp moving image into fixed image (ex. linear interpolation) 𝜙∗ = argmin 𝜙 𝑳𝒔𝒊𝒎 𝑚 𝜙 , 𝑓 + 𝑳𝒓𝒆𝒈(𝜙) Similarity Regularization Optimization problem in classical algorithms • Similarity: To evaluate the shape differences between deformed images and fixed images • Regularization: To penalize deformation fields ST: Spatial transformation 𝑓 Fixed image 𝑚 Moving image ST 𝜙 Deformation field 𝑚(𝜙) Deformed image 4
  • 5.
    DiffuseMorph: Unsupervised DeformableImage Registration Using Diffusion Model Background: Denoising diffusion probabilistic model (DDPM) [1] Jonathan Ho et al., NeurIPS 2020. • To convert a data distribution to Gaussian noise Forward diffusion process (𝒙𝟎 → 𝒙𝑻) Gaussian transition for each target data 𝒙0: ⇒ 𝑝 𝒙𝑡 𝒙0 = 𝒩 𝒙𝑡; 𝛼𝑡𝒙0, 1 − 𝛼𝑡 𝑰 where 𝛼𝑡 = ς𝑠=1 𝑡 (1 − 𝛽𝑠) 𝑝 𝒙𝑡 𝒙𝑡−1 = 𝒩(𝒙𝑡; 1 − 𝛽𝑡𝒙𝑡−1, 𝛽𝑡𝑰) • To sample image from Gaussian noise Reverse generative process (𝒙𝑻 → 𝒙𝟎) Parameterized Gaussian process: 𝑝𝜃 𝒙𝑡−1 𝒙𝑡 = 𝒩 𝒙𝑡−1; 𝜇𝜃 𝒙𝑡, 𝑡 , 𝜎𝑡 2 𝑰 𝒙𝑁~𝒩(0, 𝑰) 𝒙𝑡−1 = 1 1−𝛽𝑡 𝒙𝑡 − 𝛽𝑡 1−𝛼𝑡 𝝐𝜃 𝒙𝑡, 𝑡 + 𝜎𝑡𝒛 Generation for 𝑡 = 𝑁, 𝑁 − 1, … , 1: Parameterized model Image perturbation: 𝒙𝑡 = 𝛼𝑡𝒙0 + 1 − 𝛼𝑡𝝐 5
  • 6.
    DiffuseMorph: Unsupervised DeformableImage Registration Using Diffusion Model Proposed method: DiffuseMorph Training phase min 𝐺𝜃,𝑀𝜓 𝑳𝒅𝒊𝒇𝒇𝒖𝒔𝒊𝒐𝒏 𝑐, 𝑥𝑡, 𝑡 + 𝜆𝑳𝒓𝒆𝒈𝒊𝒔𝒕(𝑚, 𝑓) 𝑳𝒅𝒊𝒇𝒇𝒖𝒔𝒊𝒐𝒏 𝑐, 𝑥𝑡, 𝑡 = 𝔼𝜖,𝑥𝑡,𝑡 𝐺𝜃 𝑐, 𝑥𝑡, 𝑡 − 𝜖 2 2 𝑳𝒓𝒆𝒈𝒊𝒔𝒕 𝑚, 𝑓 = − 𝑚 𝜙 ⊗ 𝑓 + 𝜆𝜙∑ 𝛻𝜙 2 • Diffusion network: To estimate a conditional score function • Deformation network: To yield the registration field & provide the deformed image Training in an end-to-end learning manner Loss function 6
  • 7.
    DiffuseMorph: Unsupervised DeformableImage Registration Using Diffusion Model Proposed method: DiffuseMorph Inference phase  Continuous image registration  Synthetic image generation 7
  • 8.
    DiffuseMorph: Unsupervised DeformableImage Registration Using Diffusion Model Experiment 1 Intra-subject 2D face image registration [1] Balakrishnan, G. et al., CVPR 2018. [2] Dalca, A.V. MICCAI 2018. 8 Methods NMSE x10−1 SSIM 𝐽𝜙 ≤ 0 (%) Initial 0.363 (0.268) 0.823 (0.066) - VM [1] 0.047 (0.057) 0.936 (0.024) 0.050 (0.106) VM-diff [2] 0.034 (0.015) 0.957 (0.013) 0.014 (0.065) Ours 0.032 (0.017) 0.964 (0.011) 0.017 (0.056) • RaFD dataset (3 different gazed 8 expressions / Input size = 128x128) • Training : Validation: Test = 53 : 7 : 7 (subjects)
  • 9.
    DiffuseMorph: Unsupervised DeformableImage Registration Using Diffusion Model Experiment 2 Intra-subject 3D cardiac MR image registration [1] Balakrishnan, G. et al., CVPR 2018. [2] Dalca, A.V. MICCAI 2018. 9 Methods Dice 𝐽𝜙 ≤ 0 (%) Initial 0.642 (0.188) - VM [1] 0.787 (0.113) 0.169 (0.109) VM-diff [2] 0.794 (0.104) 0.291 (0.188) Ours 0.802 (0.109) 0.161 (0.082) • ACDC dataset (4D temporal cardiac MRI data/ input size = 128×128×32) • Training : Test = 90 : 10 (scans)
  • 10.
    DiffuseMorph: Unsupervised DeformableImage Registration Using Diffusion Model Experiment 3 Atlas-based 3D Brain MR Image Registration 10 Methods Dice 𝐽𝜙 ≤ 0 (%) Time (min) Initial 0.616 (0.171) - - SyN [1] 0.752 (0.140) 0.400 (0.100) 122, CPU VM [2] 0.749 (0.145) 0.553 (0.075) 0.01, GPU VM-diff [3] 0.731 (0.139) 0.631 (0.073) 0.01, GPU SYMNet [4] 0.733 (0.148) 0.547 (0.049) 0.43, GPU MSDIRNet [5] 0.751 (0.142) 0.804 (0.089) 2.06, GPU CM [6] 0.750 (0.144) 0.510 (0.087) 0.01, GPU Ours 0.756 (0.139) 0.505 (0.058) 0.01, GPU • OASIS-3 dataset (T1-weights MRI data / input size = 160×192×224) • Training : Validation : Test = 1027 : 93 : 129 (scans) [1] Avants, B.B.et al., Medical image analysis, 2008 [2] Balakrishnan, G. et al., CVPR 2018. [3] Dalca, A.V. MICCAI 2018. [4] Mok, T.C. et al., CVPR 2020 [5] Lei, Y. et al., Physics in Medicine & Biology, 2020. [6] Kim, B. et al., Medical image analysis 2021.
  • 11.
    DiffuseMorph: Unsupervised DeformableImage Registration Using Diffusion Model Conclusion • DiffuseMorph learns the conditional score function of deformation → To generate synthetic deformed images → To provide high-quality image registration from the continuous deformation • Experimental results on 2D and 3D data suggest the superiority of our method in image registration. • DiffuseMorph can be used to generate temporal data from the moving to the fixed images. Propose a novel diffusion-based unsupervised image registration model, called DiffuseMorph. 11
  • 12.