(2021)HyperMorph: Amortized Hyperparameter Learning for Image Registration

0. Abstract

본 논문에서는 hyperparameters를 설정할 필요가 없는 학습 기반의 deformable image registration 모델 HyperMorph를 제안하였다.

Hyperparameters는 registration 결과의 quality에 큰 영향을 미친다. 하지만 hyperparameter tuning은 시간 비용이 매우 많이 드는 과정이다.

이러한 문제를 해결하기 위해 본 논문은 deformation fields에서 hyperparameters의 영향을 학습하는 새로운 전략을 제시하였다.

제안 된 모델은 hyperparameters를 input으로 받고 registration network를 조정하는 hypernetwork를 학습함으로 hyperparameters를 찾는 수고를 덜어주고 flexibility를 높였다.

1. Introduction

학습 기반의 registration 방법은 deformation field를 빠르게 출력하는 function을 학습하기 위해 image pairs가 필요하다. 이 때 최적의 hyperparameters는 modality, anatomy에 따라 다를 수 있으며 조금의 변화에도 성능에 큰 영향을 줄 수 있다.

본 논문에서 제안한 HyperMorph는 hyperparameters에 따른 registration networks의 landscape를 근사화 한다.

HyperMorph는 initialization에 robust하여 다른 initialization으로 재학습 할 수고도 덜어주고 local minima에 빠질 위험도 적다. 또한 test 시, 최적의 hyperparameters를 빠르게 찾을 수 있다.

2. Related Work

최근에 학습 기반의 registration 방법은 image pair를 이용해 빠르게 deformation field를 출력하는 function을 학습한다. 방식은 두 가지가 있는데 deformation filed를 ground truth로 이용해 supervised 방식으로 학습하는 방법과 ground truth 없이 similarity를 이용해 loss function 만으로 학습을 하는 unsupervised 방식이 있다. 일반적으로 두 방식 모두 최소 하나 이상의 hyperparameters에 의존한다.

Hypernetworks는 gradient 기반으로 전체 model에 대해 미분이 가능해 hyperparameters를 쉽게 찾을 수 있다는 장점이 있다. HyperMorph는 hypernetworks의 연장선이다.

3. Methods

1) HyperMorph

본 논문의 핵심 idea는 hyperparameters가 registration에 미치는 영향을 학습하는 것이다.

$\mathcal{L}_h(\theta_h;\mathcal{D})=\mathbb{E}_{\Lambda\sim p(\Lambda)}\left[\mathcal{L}(\theta_h;\mathcal{D},\Lambda)\right]$

위 식에서 최적의 optimal hypernetwork parameters는 $\theta_h$, image dataset은 $\mathcal{D}$, hyperparameters에 대한 prior probability는 $p(\Lambda)$, $\mathcal{L}$은 registration loss이다.

$p(\Lambda)$는 미리 정의된 범위에서 uniform할 수 있으며, 또는 prior expectation을 기반으로 적용될 수 있다.

• Unsupervised Model Instantiation

  $\mathcal{L}_h(\theta_h;\mathcal{D})=\mathbb{E}_\Lambda\left[\sum\limits_{m,f\in\mathcal{D}}\big((1-\lambda)\mathcal{L}_{sim}(f,m\circ\phi;\lambda_{sim})+\lambda\mathcal{L}_{reg}(\phi;\lambda_{reg})\big)\right]$

  Unsupervised 방식의 registration에서 loss function은 위 식과 같다.

  $m\circ\phi$는 deformation filed $\phi$에 의해 warped된 moved image를 나타낸다. Loss term $\mathcal{L}_{sim}$은 moved image와 fixed image간의 similarity를 측정하며 hyperparameter $\lambda_{sim}$을 포함한다. $\lambda$는 loss term간의 적절한 importance를 맞춰주는 역할을 한다. 즉, $\Lambda$는 다음과 같이 구성된다.

  $\Lambda=\{\lambda,\lambda_{sim},\lambda_{reg}\}$

  같은 modality의 images를 registration 할 때, similarity metrics로는 MSE(Mean-Squared Error)와 local NCC(Normalized Cross-Correlation)을 사용한다. 반대로 다른 modality의 images를 registration 할 때는 NMI(Normalized Mutual Information)을 사용한다.

  Regularization을 위해 deformation filed $\phi$를 SVF(Stationary velocity filed)로 parameterize하고 diffeomorphism을 얻기 위해 network와 통합한다. Regularization 식은 다음과 같다.

  $\mathcal{L}_{reg}(\phi)=\frac{1}{2}||\nabla v||^2$
  

• Semi-supervised Model Instantiation

  $\mathcal{L}_h(\theta_h;\mathcal{D})=$

  $\mathbb{E}_\Lambda\left[\sum\limits_{m,f\in\mathcal{D}}(1-\lambda)(1-\gamma)\mathcal{L}_{sim}(f,m\circ\phi;\lambda_{sim})+\lambda\mathcal{L}_{reg}(\phi;\lambda_{reg})+(1-\lambda)\gamma\mathcal{L_{seg}(s_f,s_m\circ\phi)}\right]$

  본 논문에서는 HyperMorph의 loss를 수정하여 semi-supervised setting도 적용하였다.

  $\mathcal{L}_{seg}$는 segmentation similarity metric으로 보통 Dice coefficient를 사용한다. 다른 loss term과의 importance를 맞추기 위해 hyperparameter $\gamma$를 사용하였다.
  

2) Hyperparameter Tuning

Image pair $\{m,f\}$가 주어지면 HyperMorph는 hyperparameters로 부터 deformation field를 생성한다. 만약 추가적인 정보가 없다면 최적의 hyperparameters를 iterative fashion으로 빠르게 얻을 수 있다.

그러나 만약 validation subjects의 landmark 또는 segmentation maps이 주어진다면 빠른 automatic tuning이 가능하다.

• Iterative

  Hyperparameters를 변경하는데 slider가 사용될 수 있으며 실시간으로 hyperparameter를 바꿔가며 registration 결과를 확인하며 조정할 수 있다.

  특정 region에 관심이 있는 경우에는 $\lambda$를 조절해 anatomical structure를 변경할 수 있다.

• Automatic

  $\Lambda^*=\arg\max\limits_{\Lambda}\mathcal{L}(\Lambda;\theta_h,\mathcal{D},\mathcal{V})=\arg\max\limits_{\Lambda}\sum\limits_{(m,f)\in\mathcal{D^2} \atop (s_m,s_f)\in\mathcal{V}^2}\mathcal{L}_{val}(s_f,s_m\circ\phi)$

  만약 validation 시 segmentation maps $\{s_m,s_f\}$가 주어진다면, HyperMorph는 hyperparameter를 위 식과 같이 최적화 할 수 있다.
  

3) Implementation

Hypernetwork는 64 units의 4개의 fully connected layer와 ReLU activation function으로 구성되어 있으며 마지막 layer는 Tanh activation function이다.

HyperMorph는 U-Net을 registration architecture로 사용한 VoxelMorph library를 기반으로 개발되었다.

Optimizer는 Adam을 사용하였다.

4. Experiments

실험을 통해 HyperMorph가 개별적으로 hyperparameters를 다르게 설정해 학습한 model과 동등한 landscape를 그리고 있음을 검증했다. 그리고 다른 subpopulation, registration types, anatomical structures의 경우에 바르게 hyperparameters를 최적화 할 수 있음을 확인하였다.

• Dataset

  2개의 큰 3D brain MR image dataset을 사용하였다. 첫 번째로, 30495장의 T1_weighted scans가 모인 ABIDE, ADHD200, ADNI, GSP, MCIC, PPMI, OASIS, UK Biobank dataset이다.

  두 번째, 1558 T1_weighted, T2-weighted, multi-flip-angle, multi-inversion-time images가 모인 ADNI, HCP dataset을 사용하였다.

  모든 MRI scans는 $256\times256\times256$ volumes, $1-mm$ isotropic grid space에 conformed 되었으며 FreeSurfer를 이용해 bias-correcting, skull-stripping을 적용했다. 추가적으로 evaluation을 segmentation maps도 얻었다.

  전체 데이터는 affinely normalization한 후, $160\times192\times224$ volumes로 일정하게 crop되었다.
  

• Evaluation

  Evaluation을 위해 anatomical label maps에 대한 Dice metric을 사용하였다.
  

• Baseline Models

  HyperMorph는 학습 기반의 어떤 registration architecture에도 적용될 수 있다. 본 논문에서는 standard model로 VoxelMorph를 선정하였다.
  

1) Experiment 1: HyperMorph Efficiency and Capacity

단일 model HyperMorph가 baseline model의 landscape를 따라가는 지를 평가하는 것을 목표로 한다.

HyperMorph는 baseline model의 grid search 방식보다 5배 이상 적은 시간이 소요된다. 만약 2개의 hyperparameters가 있다면 10배 이상 적은 시간이 소요된다.

위 두 그림은 HyperMorph가 computational cost가 적음에도 optimal hyperparameters를 baseline model의 dense grid search 결과와 동일하게 잘 찾음을 보여준다.

두 모델의 optimal hyperparameter $\lambda^*$는 $0.04\pm0.02$, Dice difference는 $0.16\pm0.03$ 차이 밖에 나지 않는다. 동일하게 multi-hyperparameters 인 경우에도 Dice difference가 $0.32\pm0.02$로 매우 작은 차이가 난다.

2) Experiment 2: Robustness to Initialization

Network initialization의 robustness를 평가하기 위해 4개의 HyperMorph model과 5개의 서로 다른 hyperparameter $\lambda$로 학습한 baseline model을 비교하였다.

위 그림을 보면 HyperMorph가 baseline model에 비해 initialization에 standard deviation이 작아 robustness가 우수함을 확인할 수 있다.

3) Experiment 3: Hyperparameter-Tuning Utility

HyperMorph의 hyperparameter visual optimization tool은 http://voxelmorph.mit.edu에 제공된다.

Validation 시 anatomical annotations가 제공되는 경우 hyperparameters의 automatic optimization이 가능하다. 위 두 그림을 보면 anatomical region의 종류에 따라 최적의 hyperparameter $\lambda$가 다름을 확인할 수 있다.

HyperMorph는 하나의 model로 관심이 있는 anatomical region에 맞춰 hyperparameter를 수정함과 동시에 deformation field를 바로 얻을 수 있다는 장점이 있다.

5. Conclusion

본 논문에서는 HyperMorph를 제안함으로, hyperparameters의 model 성능에 대한 영향을 정량화 하기 위해 반복해서 model을 학습하는 수고를 덜어주었다.

기존의 model과는 달리, HyperMorph는 입력되는 hyperparameters에 해당하는 최적의 deformation field를 예측하여 기존 registration model이 그리는 landscape를 비슷하게 잘 따라감을 확인하였다.

또한 다른 subpopulation, registration task, anatomical region에 따라 최적의 hyperparameters를 빠르게 찾을 수 있는 장점이 있으며 각각의 image pair 마다 최적의 hyperparameters를 적용할 수도 있다.

HyperMorph는 최적의 hyperparameters를 찾기 위해 재 학습을 반복하는 수고를 덜어줌으로 효과적으로 registration model을 최적화 할 수 있게 해주었다.