Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

AdamFlow: Adam-based Wasserstein Gradient Flows for Surface Registration in Medical Imaging

Qiang Ma, Qingjie Meng, Xin Hu, Yicheng Wu, Wenjia Bai

Abstract

Surface registration plays an important role for anatomical shape analysis in medical imaging. Existing surface registration methods often face a trade-off between efficiency and robustness. Local point matching methods are computationally efficient, but vulnerable to noise and initialisation. Methods designed for global point set alignment tend to incur a high computational cost. To address the challenge, here we present a fast surface registration method, which formulates surface meshes as probability measures and surface registration as a distributional optimisation problem. The discrepancy between two meshes is measured using an efficient sliced Wasserstein distance with log-linear computational complexity. We propose a novel optimisation method, AdamFlow, which generalises the well-known Adam optimisation method from the Euclidean space to the probability space for minimising the sliced Wasserstein distance. We theoretically analyse the asymptotic convergence of AdamFlow and empirically demonstrate its superior performance in both affine and non-rigid surface registration across various anatomical structures.

AdamFlow: Adam-based Wasserstein Gradient Flows for Surface Registration in Medical Imaging

Abstract

Surface registration plays an important role for anatomical shape analysis in medical imaging. Existing surface registration methods often face a trade-off between efficiency and robustness. Local point matching methods are computationally efficient, but vulnerable to noise and initialisation. Methods designed for global point set alignment tend to incur a high computational cost. To address the challenge, here we present a fast surface registration method, which formulates surface meshes as probability measures and surface registration as a distributional optimisation problem. The discrepancy between two meshes is measured using an efficient sliced Wasserstein distance with log-linear computational complexity. We propose a novel optimisation method, AdamFlow, which generalises the well-known Adam optimisation method from the Euclidean space to the probability space for minimising the sliced Wasserstein distance. We theoretically analyse the asymptotic convergence of AdamFlow and empirically demonstrate its superior performance in both affine and non-rigid surface registration across various anatomical structures.

Paper Structure

This paper contains 43 sections, 4 theorems, 59 equations, 9 figures, 7 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. Distributional optimisation.
  4. Wasserstein distance.
  5. Sliced Wasserstein distance.
  6. Wasserstein gradient.
  7. Wasserstein gradient flow.
  8. Surface registration
  9. Surface mesh as a probability measure
  10. Surface registration in the Wasserstein space
  11. Affine surface registration
  12. Non-rigid surface registration
  13. AdamFlow for distributional optimisation
  14. Formulation of AdamFlow
  15. Discrete Adam.
  16. ...and 28 more sections

Key Result

Proposition 1

If $(x(t),m(t),v(t))$ is a solution of Adam ODE (eq:adam_ode) with an objective function $f$, then the trajectory of the Dirac measure $\delta_{(x(t),m(t),v(t))}$ is a solution of AdamFlow (eq:adamflow) with the objective functional $F_f[\mu]=\int_{\mathbb{R}^d}f(x)\mathrm{d}\mu(x)$ in the distribut

Figures (9)

  • Figure 1: Illustration of affine surface registration of the left ventricle of the heart and non-rigid surface registration of the liver. The surface registration is formulated as a distributional optimisation problem between a source surface ($t=0$) and a target surface. The registration is performed by integrating a Wasserstein gradient flow for $t\in[0,T]$ to optimise an objective functional.
  • Figure 2: Comparative results and convergence curves for affine and non-rigid surface registration. The registration errors are measured by ASSD (mm) and HD90 (mm). Registration errors decrease with increasing numbers of optimisation iterations.
  • Figure 3: Qualitative comparisons for affine surface registration of the liver, pancreas, and left ventricle from a source mesh to a target mesh. AdamFlow shows consistently better affine registration accuracy than the traditional ICP algorithm.
  • Figure 4: Qualitative comparisons for non-rigid surface registration of the liver, pancreas, and left ventricle from a source mesh to a target mesh. AdamFlow is used to minimises the SWD for global surface alignment (coarse), and then optimises the Chamfer distance for local mesh refinement (fine), enabling coarse-to-fine surface registration.
  • Figure 5: Runtime (millisecond) for computing the discrepancy metrics across different numbers of points.
  • ...and 4 more figures

Theorems & Definitions (8)

  • Proposition 1
  • Proposition 2
  • Theorem 1
  • Proposition 3
  • proof
  • proof
  • proof
  • proof