Download as pdf or txt
Download as pdf or txt
You are on page 1of 65

voxelmorph.mit.

edu
+ in future FreeSurfer release

Unsupervised Learning of Image Correspondences


in Medical Imaging Analysis
Adrian V. Dalca
HMS, MIT
Small vessel disease

Stroke

Dalca et al, MICCAI, 2014


Sridharan*, Dalca* et al, MBIA 2013
Progression with age

31 years 42.5 years 54 years 65.5 years 77 years


average average average average average

Sridharan*, Dalca* et al, MBIA 2013


Scan at age 50
Scan at age 60

Dalca et al, MICCAI, 2015


Batmanghelich, Dalca et al, IPMI 2013, TMI 2016
0.5 sec

baseline (CPU) ours (CPU) ours-GPU


0.5 sec

baseline (CPU) ours (CPU) ours-GPU


Registration

field 𝜙𝜙

moving scan 𝒎𝒎 fixed scan 𝒇𝒇


Registration is fundamental in MIA
• Register scans to a template for analysis
• Register subject scans to each other for direct comparison
• Clinical data alignment
e.g. before and after surgery
• Segmentation
propagate anatomical labels
• Related to alignment in other fields
computer vision, 1D signals, computational biology
Pairwise optimization

scan 𝒎𝒎 scan 𝒇𝒇

field 𝜙𝜙

warped scan 𝑚𝑚

𝜙𝜙� 𝑚𝑚,𝑓𝑓 = argmin𝜙𝜙 ‖𝑚𝑚 ∘ 𝜙𝜙 − 𝑓𝑓‖ + 𝜆𝜆 Reg(𝜙𝜙)


optimal deformation field images match smooth field

• significant development
• slow for two images
Learning-based methods

CNN

𝑔𝑔𝜃𝜃 𝑚𝑚, 𝑓𝑓 field 𝜙𝜙

• Supervised (have example triplets {𝑚𝑚, 𝑓𝑓, 𝜙𝜙})


• Unsupervised (only have images {𝑚𝑚, 𝑓𝑓}) (voxelmorph)
Learning-based methods

CNN

𝑔𝑔𝜃𝜃 𝑚𝑚, 𝑓𝑓 field 𝜙𝜙

• Supervised (have example triplets {𝑚𝑚, 𝑓𝑓, 𝜙𝜙})


• Unsupervised (only have images {𝑚𝑚, 𝑓𝑓}) (voxelmorph)

• limited use of classical modelling


• fast for new image pair
Outline
• Model
• Variational Inference with neural networks
• Optimization interpretation

• Results (runtime and accuracy) Loss (ℒ)

• Extensions & Applications


• Training with limited data
• Training with auxiliary data
• Training with clinical data
• Learn deformable templates
Framework
deformation
Moving image
(𝒎𝒎)


Fixed image
(𝒇𝒇) network 𝑔𝑔𝜃𝜃
parameters 𝜃𝜃
Supervised:
2
ℒ = 𝜙𝜙 − 𝜙𝜙𝑔𝑔𝑔𝑔
VoxelMorph
deformation
Moving image
(𝒎𝒎)


Fixed image
(𝒇𝒇) network 𝑔𝑔𝜃𝜃
parameters 𝜃𝜃
Unsupervised:
ℒ = ‖𝑚𝑚 ∘ 𝜙𝜙 − 𝑓𝑓‖ + 𝜆𝜆 Reg(𝜙𝜙)
images match smooth field
VoxelMorph Loss
deformation
Moving image
(𝒎𝒎)


Fixed image
(𝒇𝒇) network 𝑔𝑔𝜃𝜃
parameters 𝜃𝜃

ℒ = � ‖𝑚𝑚𝑖𝑖 ∘ 𝜙𝜙𝑖𝑖𝑖𝑖 − 𝑓𝑓𝑖𝑖𝑖𝑖 ‖ + 𝜆𝜆 Reg(𝜙𝜙𝑖𝑖𝑖𝑖 )


𝑖𝑖,𝑗𝑗
VoxelMorph Loss
deformation
Moving image
(𝒎𝒎)


Fixed image
(𝒇𝒇) network 𝑔𝑔𝜃𝜃
parameters 𝜃𝜃

ℒ(𝜃𝜃; data) = � ‖𝑚𝑚𝑖𝑖 ∘ 𝑔𝑔𝜃𝜃 (𝑚𝑚𝑖𝑖. 𝑓𝑓𝑖𝑖 ) − 𝑓𝑓𝑖𝑖𝑖𝑖 ‖ + 𝜆𝜆 Reg(𝑔𝑔𝜃𝜃 (𝑚𝑚𝑖𝑖. 𝑓𝑓𝑖𝑖 ))
𝑖𝑖,𝑗𝑗
𝜙𝜙𝑖𝑖𝑖𝑖 𝜙𝜙𝑖𝑖𝑖𝑖
VoxelMorph Loss
deformation
Moving image
Moved (𝒎𝒎 ∘ 𝝓𝝓𝑧𝑧 )
(𝒎𝒎)


spatial
Fixed image transform
(𝒇𝒇) network 𝑔𝑔𝜃𝜃
parameters 𝜃𝜃 moved image 𝑚𝑚 ∘ 𝜙𝜙
should be close to 𝑓𝑓

ℒ(𝜃𝜃; data) = � ‖𝑚𝑚𝑖𝑖 ∘ 𝑔𝑔𝜃𝜃 (𝑚𝑚𝑖𝑖. 𝑓𝑓𝑖𝑖 ) − 𝑓𝑓𝑖𝑖𝑖𝑖 ‖ + 𝜆𝜆 Reg(𝑔𝑔𝜃𝜃 (𝑚𝑚𝑖𝑖. 𝑓𝑓𝑖𝑖 ))
𝑖𝑖,𝑗𝑗
𝜙𝜙𝑖𝑖𝑖𝑖 𝜙𝜙𝑖𝑖𝑖𝑖
Training
deformation
Moving image
Moved (𝒎𝒎 ∘ 𝝓𝝓𝑧𝑧 )
(𝒎𝒎)


spatial
Fixed image transform
(𝒇𝒇) network 𝑔𝑔𝜃𝜃
parameters 𝜃𝜃 moved image 𝑚𝑚 ∘ 𝜙𝜙
should be close to 𝑓𝑓
• SGD based techniques
• Each image pair contributes slightly to 𝜃𝜃
Classical optimization: slightly update 𝜙𝜙 for an image pair
Registration
deformation
Moving image
(𝒎𝒎)


Fixed image
(𝒇𝒇) network 𝑔𝑔𝜃𝜃
parameters 𝜃𝜃
Probabilistic model

𝑚𝑚 ∘ 𝜙𝜙𝑧𝑧 + 𝜖𝜖 = 𝑓𝑓
𝑧𝑧 ~ 𝒩𝒩 𝑧𝑧; 0, Λ−1
stationary velocity field smoothness via Laplacian

Goal: 𝑝𝑝(𝑧𝑧|𝑚𝑚, 𝑓𝑓) posterior probability of registration


Atlas-based registration
Data: 7000 training volumes, 250 validate, 250 test
Baseline: ANTs optimization method
Runtime for a new 3D image pair

51 sec 0.5 sec

ANTs [CPU] NiftyReg [CPU] VoxelMorph [CPU] VoxelMorph [GPU]


Anatomical volume overlap

scan m scan f warped m (VM) warped m (baseline)


pair 1

*algorithms only see images, no segmentation maps


Accuracy via volume overlap (Dice)
baseline
VoxelMorph
Outline
• Model
• Variational Inference with neural networks
• Optimization interpretation

• Results (runtime and accuracy) Loss (ℒ)

• Extensions & Applications


• Training with limited data
• Training with auxiliary data
• Training with clinical data
• Learn deformable templates
Amortized analysis: training with limited data

0.9
ANTs (Syn)
0.85 Test Set
Test Set (Instance-Specific Opt.)
0.8
Dice score
0.75

0.7

0.65

0.6
10 25 50 100
Training set size
Segmentation Maps available at training
Moving 3D Image (𝒎𝒎)

deformation
field (𝝓𝝓𝑧𝑧 )

… Moved (𝒎𝒎 ∘ 𝝓𝝓𝑧𝑧 )


Fixed 3D Image (𝒇𝒇)
spatial
transform
param: 𝜓𝜓

Moving Image
Segmentations
(𝒔𝒔𝒎𝒎 ) Move Image
spatial
transform Segmentations
(𝒔𝒔𝒎𝒎 ∘ 𝝓𝝓𝑧𝑧 )
Fixed Image
Segmentations
(𝒔𝒔𝒇𝒇 )
Dice Loss (ℒ)
Test time performance
0.9

0.85

0.8
Dice score
0.75

0.7

0.65

0.6
VoxelMorph VoxelMorph w/ seg
SynthMorph
(do we need real data?)
Hoffmann et al in submission
Inputlabels Deformed labels GMM Sampling Blurred image Finalimage

Billot MIDL 2020


https://github.com/BBillot/lab2im Billot MICCAI 2020
Do we need anatomical images to train?
Do we need anatomical images to train?

Moving Deformation 𝜙𝜙𝑣𝑣


image 𝑚𝑚 via 𝑣𝑣 integration

Modality
agnostic
CNN ℎ𝜃𝜃 (𝑚𝑚, 𝑓𝑓)
Fixed
image 𝑓𝑓

Hoffman et al, in revision


Do we need anatomical images to train?

Moving Moving
labels 𝑠𝑠𝑚𝑚 image 𝑚𝑚
Label
Image sampling Modality
sampling
model agnostic
model
𝑔𝑔𝐼𝐼 (𝑠𝑠𝑚𝑚 , 𝑠𝑠𝑓𝑓 , 𝑧𝑧) CNN ℎ𝜃𝜃 (𝑚𝑚, 𝑓𝑓)
𝑔𝑔𝑠𝑠 (𝑧𝑧)
Fixed Fixed
labels 𝑠𝑠𝑓𝑓 image 𝑓𝑓

Hoffman et al, in revision


Do we need anatomical images to train?

Warped
Moving Moving Deformation 𝜙𝜙𝑣𝑣 segmentation
labels 𝑠𝑠𝑚𝑚 image 𝑚𝑚 via 𝑣𝑣 integration 𝑠𝑠𝑚𝑚 ∘ 𝜙𝜙𝑣𝑣
Label
Image sampling Modality
sampling
model agnostic
model
𝑔𝑔𝐼𝐼 (𝑠𝑠𝑚𝑚 , 𝑠𝑠𝑓𝑓 , 𝑧𝑧) CNN ℎ𝜃𝜃 (𝑚𝑚, 𝑓𝑓)
𝑔𝑔𝑠𝑠 (𝑧𝑧)
Fixed Fixed
labels 𝑠𝑠𝑓𝑓 image 𝑓𝑓 Regularization term ℒ 𝑟𝑟𝑟𝑟𝑟𝑟
Dice term ℒ𝑑𝑑𝑑𝑑𝑑𝑑

Hoffman et al, in revision


b Skull-stripped x Cross-dataset (e.g. HCP-OASIS)
VoxelMorph - NMI SynthMorph (brains)
HyperMorph:
Amortized parameter learning
Hoopes, Hoffmann, Fischl, Guttag, Dalca, IPMI 2021
Regularization Analysis (hyperparameters)

ℒ = ‖𝑚𝑚 ∘ 𝜙𝜙 − 𝑓𝑓‖ + 𝜆𝜆 Reg(𝜙𝜙)


images match smooth field
HyperMorph
VoxelMorph Strategy

Train Multiple Models Validate Each Choose Optimal


Model Hyperparameter λ

Model Using λ1

Model Using λ2

Model Using λ3 λ

Repeat with More Precise λ


Values

Hoopes et al, IPMI 2021


HyperMorph
VoxelMorph Strategy HyperMorph Strategy

Train Multiple Models Validate Each Choose Optimal Train Single Efficiently Optimize
Model Hyperparameter λ Model Hyperparameter λ Using
Validation Data
Model Using λ1

Model Using λ2 HyperMorph


Model
Model Using λ3 λ λ

Repeat with More Precise λ {λ}


Values

Hoopes et al, IPMI 2021


HyperMorph
Learned template (𝒕𝒕)

deformation
velocity field (𝝓𝝓𝑖𝑖 ) Moved (𝒕𝒕 ∘ 𝝓𝝓𝑖𝑖 )
field
Spatial
𝑔𝑔𝜽𝜽𝑣𝑣 𝒗𝒗𝑖𝑖 Transform
Image (𝒙𝒙𝑖𝑖 )
integration
layer

Loss (ℒ)

Hyperparameter
Input 𝜆𝜆
HyperMorph
Learned template (𝒕𝒕)

deformation
velocity field (𝝓𝝓𝑖𝑖 ) Moved (𝒕𝒕 ∘ 𝝓𝝓𝑖𝑖 )
field
Spatial
𝑔𝑔𝜽𝜽𝑣𝑣 𝒗𝒗𝑖𝑖 Transform
Image (𝒙𝒙𝑖𝑖 )
integration
θv layer

Loss (ℒ)

Hypernetwork (h)
Hyperparameter
Input 𝜆𝜆
Baseline Comparison

Dice

Runtime (GPU-hours)
VoxelMorph (~10 models): 765
HyperMorph: 147
Optimal Hyperparameters vary by dataset
Optimal Hyperparameters vary by task
… even by anatomical region!
Template construction
Dalca, Rakic, Guttag, Sabuncu, NeurIPS 2019
Joshi et al, 2004
Template Construction
Learned template (𝒕𝒕)

param:
deformation
Ψatl
velocity field (𝝓𝝓𝑖𝑖 ) Moved (𝒕𝒕 ∘ 𝝓𝝓𝑖𝑖 )
field
Spatial
𝑔𝑔𝑣𝑣,𝜽𝜽𝑣𝑣 𝒙𝒙𝑖𝑖 , 𝒕𝒕 𝒗𝒗𝑖𝑖 Transform
Image (𝒙𝒙𝑖𝑖 )
integration
layer

Loss (ℒ)
Conditional template construction
Learned template (𝒕𝒕)
Attribute (𝒂𝒂𝑖𝑖 )

deformation
velocity field (𝝓𝝓𝑖𝑖 ) Moved (𝒕𝒕 ∘ 𝝓𝝓𝑖𝑖 )
field
Spatial
𝑔𝑔𝑣𝑣,𝜽𝜽𝑣𝑣 𝒙𝒙𝑖𝑖 , 𝒕𝒕 𝒗𝒗𝑖𝑖 Transform
Image (𝒙𝒙𝑖𝑖 )
integration
layer

Loss (ℒ)
15 90
29 weeks 32 35 38 41 44

dHCP atlas Dey et al, in submission


Acknowledgements
Guha Balakrishnan (MIT CSAIL DDIG)
Benjamin Billot (UCL CMIC)
Bruce Fischl (HMS/MGH LCN)
John Guttag (MIT CSAIL DDIG)
Malte Hoffmann (MGH LCN)
Andrew Hoopes (MGH LCN)
Eugenio Iglesias (MIT CSAIL, MGH HMS, UCL CMIC)
Kathleen Lewis (MIT CSAIL DDIG)
Marianne Rakic (MIT CSAIL DDIG, ETH)
Mert Sabuncu (Cornell ECE, HMS/MGH LCN)
Amy Zhao (MIT CSAIL DDIG)
vo elmorph
• Probabilistic generative model • Very fast for new image pair
for diffeomorphisms • State-of-the-art accuracy
• Variational Inference • Diffeomorphic deformations
• Unsupervised Neural Network • Uncertainty estimation

voxelmorph.mit.edu
vo elmorph
• Probabilistic generative model • Very fast for new image pair
for diffeomorphisms • State-of-the-art accuracy
• Variational Inference • Diffeomorphic deformations
• Unsupervised Neural Network • Uncertainty estimation

voxelmorph.mit.edu

• Limited training data  use VM as initialization


• Segmentation at training  better test Dice performance
• No atlas  construct atlas automatically
• Synthesis  invariant representations
• Can apply to wider domains
vo elmorph
• Probabilistic generative model • Very fast for new image pair
for diffeomorphisms • State-of-the-art accuracy
• Variational Inference • Diffeomorphic deformations
• Unsupervised Neural Network • Uncertainty estimation

voxelmorph.mit.edu

• Limited training data  use VM as initialization


• Segmentation at training  better test Dice performance
• No atlas  construct atlas automatically
• Synthesis  invariant representations
• Can apply to wider domains

You might also like