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

2026-04-02

AI Summary

This work addresses the challenge of simultaneously achieving efficiency and robustness in medical image surface registration by modeling surface meshes as probability measures and reformulating registration as a distributional optimization problem. To measure discrepancies between distributions, the method employs the computationally efficient sliced Wasserstein distance. The core contribution is AdamFlow, a novel algorithm that extends the Adam optimizer to the space of probability measures, enabling efficient minimization of the sliced Wasserstein distance via gradient flows. AdamFlow offers both theoretical convergence guarantees and practical scalability. Experimental results demonstrate that the proposed approach significantly outperforms existing methods across diverse anatomical structures, achieving high accuracy, strong robustness, and rapid convergence in both affine and non-rigid registration settings.

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.

Problem

• surface registration
• medical imaging
• efficiency-robustness trade-off
• noise sensitivity
• computational cost

Innovation

• AdamFlow
• sliced Wasserstein distance
• surface registration
• probability measures
• distributional optimisation